package llm

import (
	"encoding/json"
	"net/http"

	"github.com/tidbcloud/AgentTeam/internal/apr/types"

	"github.com/gin-gonic/gin"
)

// Embeddings proxies embeddings API requests through the AgentTeam APR-backed proxy.
func (handler *LLMHandler) Embeddings(ctx *gin.Context) {
	resolved, ok := handler.resolveRequest(ctx)
	if !ok {
		return
	}
	observation, _ := handler.newUsageObservation(ctx, resolved)
	if !handler.enforceQuota(ctx, observation) {
		return
	}
	queueWait, release, ok := handler.acquireConcurrency(ctx, resolved, "embeddings")
	if !ok {
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "rate_limit")
		return
	}
	if observation != nil {
		observation.queueWait = queueWait
	}
	defer release()

	embeddingsProvider, ok := resolved.Provider.(types.EmbeddingCompatible)
	if !ok {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "provider does not support embeddings"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "bad_request")
		return
	}

	var request types.EmbeddingRequest
	if err := json.Unmarshal(resolved.Body, &request); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "invalid embedding request"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "bad_request")
		return
	}
	resp, err := embeddingsProvider.CreateEmbedding(ctx.Request.Context(), &request)
	if err != nil {
		handler.respondUpstreamError(ctx, observation, resolved.Provider.Name(), resolved.Model, "embeddings", err)
		return
	}
	ctx.JSON(http.StatusOK, resp)
	if observation != nil {
		cachedTokens := 0
		if resp.Usage.PromptTokensDetails != nil && resp.Usage.PromptTokensDetails.CachedTokens != nil {
			cachedTokens = *resp.Usage.PromptTokensDetails.CachedTokens
		}
		handler.recordObservation(ctx, observation, llmUsageTotals{
			InputTokens:       uint64(maxInt(resp.Usage.PromptTokens, 0)),
			CachedInputTokens: uint64(maxInt(cachedTokens, 0)),
			OutputTokens:      uint64(maxInt(resp.Usage.CompletionTokens, 0)),
		}, "")
	}
}