package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
	"time"

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

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

type requestResolver struct {
	Provider types.LLMProvider
	Model    string
	Body     []byte
}

func (handler *LLMHandler) resolveRequest(ctx *gin.Context) (*requestResolver, bool) {
	rawBody, err := io.ReadAll(ctx.Request.Body)
	if err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "failed to read request body"})
		return nil, false
	}

	var request struct {
		Model string `json:"model"`
	}
	if err := json.Unmarshal(rawBody, &request); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
		return nil, false
	}

	authorized, ok := llmProxyAuthorizationFromContext(ctx.Request.Context())
	if !ok {
		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing llm proxy authorization"})
		return nil, false
	}
	provider, normalizedModel := handler.resolveAuthorizedProviderModel(authorized, request.Model)
	if provider == nil {
		if handler.logger != nil {
			handler.logger.Warn("llm request rejected: unsupported model",
				"requested_model", strings.TrimSpace(request.Model),
				"workspace_id", strings.TrimSpace(authorized.WorkspaceID),
				"organization_id", strings.TrimSpace(authorized.OrganizationID),
				"allowed_provider_ids", authorized.AllowedProviderIDs,
				"allowed_models", authorized.AllowedModels,
			)
		}
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "unsupported model: " + request.Model})
		return nil, false
	}

	rewrittenBody, err := rewriteModelField(rawBody, normalizedModel)
	if err != nil {
		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "failed to rewrite request model"})
		return nil, false
	}

	return &requestResolver{Provider: provider, Model: normalizedModel, Body: rewrittenBody}, true
}

func (handler *LLMHandler) acquireConcurrency(ctx *gin.Context, resolved *requestResolver, operation string) (time.Duration, func(), bool) {
	if resolved == nil || resolved.Provider == nil {
		return 0, func() {}, true
	}
	limiter := concurrency.DefaultLimiter()
	key := strings.TrimSpace(resolved.Provider.Name()) + "/" + strings.TrimSpace(resolved.Model)
	waitDuration, release, err := limiter.Acquire(ctx.Request.Context(), key)
	if err != nil {
		if handler.logger != nil {
			handler.logger.Warn("llm concurrency acquire failed", "provider", resolved.Provider.Name(), "model", resolved.Model, "operation", operation, "error", err)
		}
		ctx.JSON(http.StatusTooManyRequests, gin.H{"error": "llm concurrency limit exceeded"})
		return 0, nil, false
	}
	if handler.logger != nil && waitDuration > 0 {
		handler.logger.Debug("llm concurrency wait", "provider", resolved.Provider.Name(), "model", resolved.Model, "operation", operation, "wait_ms", waitDuration.Milliseconds())
	}
	return waitDuration, release, true
}

func llmProxyAuthorizationFromContext(ctx context.Context) (llmProxyAuthorization, bool) {
	value := ctx.Value(ctxKeyLLMProxyAuthorization)
	if value == nil {
		return llmProxyAuthorization{}, false
	}
	authorization, ok := value.(llmProxyAuthorization)
	return authorization, ok
}

func (handler *LLMHandler) resolveAuthorizedProviderModel(authorized llmProxyAuthorization, requestedModel string) (types.LLMProvider, string) {
	if handler.aprDispatcher == nil {
		return nil, ""
	}
	if strings.TrimSpace(authorized.WorkspaceID) != "" {
		normalizedModel, ok := normalizeAuthorizedModel(
			requestedModel,
			authorized.AllowedModels,
			"",
		)
		if !ok {
			return nil, ""
		}
		provider, resolvedModel := handler.aprDispatcher.LookupAmongProviders("", normalizedModel, authorized.AllowedProviderIDs)
		if provider != nil {
			return provider, resolvedModel
		}
		return handler.resolveDirectAuthorizedProviderModel(normalizedModel, authorized.AllowedProviderIDs)
	}
	trimmedProviderID := strings.TrimSpace(authorized.OrganizationProviderID)
	if trimmedProviderID == "" {
		return nil, ""
	}
	providerModels := providerModelIDs(handler.aprDispatcher.ProviderModels()[trimmedProviderID])
	normalizedModel, ok := normalizeAuthorizedModel(requestedModel, providerModels, trimmedProviderID)
	if !ok {
		return nil, ""
	}
	return handler.aprDispatcher.Lookup("", trimmedProviderID+"/"+normalizedModel)
}

func (handler *LLMHandler) resolveDirectAuthorizedProviderModel(requestedModel string, allowedProviderIDs []string) (types.LLMProvider, string) {
	if handler.aprDispatcher == nil {
		return nil, ""
	}
	trimmedRequested := strings.ToLower(strings.TrimSpace(requestedModel))
	if trimmedRequested == "" {
		return nil, ""
	}
	providerModels := handler.aprDispatcher.ProviderModels()
	for _, providerID := range allowedProviderIDs {
		trimmedProviderID := strings.TrimSpace(providerID)
		if trimmedProviderID == "" {
			continue
		}
		for _, model := range providerModels[trimmedProviderID] {
			modelID := strings.ToLower(strings.TrimSpace(model.ID))
			if modelID == "" {
				continue
			}
			if modelID != trimmedRequested {
				continue
			}
			return handler.aprDispatcher.Lookup("", trimmedProviderID+"/"+modelID)
		}
	}
	return nil, ""
}

func providerModelIDs(models []types.Model) []string {
	result := make([]string, 0, len(models))
	for _, model := range models {
		if trimmed := strings.TrimSpace(model.ID); trimmed != "" {
			result = append(result, trimmed)
		}
	}
	return result
}

func normalizeAuthorizedModel(requestedModel string, allowedModels []string, providerID string) (string, bool) {
	trimmedModel := strings.TrimSpace(requestedModel)
	if trimmedModel == "" {
		return "", false
	}
	loweredAllowed := make([]string, 0, len(allowedModels))
	for _, model := range allowedModels {
		trimmedAllowed := strings.TrimSpace(model)
		if trimmedAllowed == "" {
			continue
		}
		loweredAllowed = append(loweredAllowed, strings.ToLower(trimmedAllowed))
	}
	loweredRequested := strings.ToLower(trimmedModel)
	if slices.Contains(loweredAllowed, loweredRequested) {
		return loweredRequested, true
	}
	parts := strings.SplitN(loweredRequested, "/", 2)
	if len(parts) == 2 {
		if strings.EqualFold(parts[0], providerID) && slices.Contains(loweredAllowed, parts[1]) {
			return parts[1], true
		}
		if slices.Contains(loweredAllowed, parts[1]) {
			return parts[1], true
		}
	}
	return "", false
}

func rewriteModelField(rawBody []byte, model string) ([]byte, error) {
	var payload map[string]any
	if err := json.Unmarshal(rawBody, &payload); err != nil {
		return nil, err
	}
	payload["model"] = model
	return json.Marshal(payload)
}

func isStreamingResponse(contentType string) bool {
	return strings.Contains(contentType, "text/event-stream") || strings.Contains(contentType, "text/plain")
}

func (handler *LLMHandler) respondUpstreamError(ctx *gin.Context, observation *llmUsageObservation, providerName string, model string, operation string, err error) {
	if handler.logger != nil {
		handler.logger.Error(operation+" failed",
			"provider", providerName,
			"model", model,
			"error", err,
		)
	}
	ctx.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("upstream %s failed", operation)})
	handler.recordObservation(ctx, observation, llmUsageTotals{}, classifyUpstreamTransportError(err))
}