package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/store/organization"

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

type streamingResponseLineTransformer func(line []byte) []byte

type llmUsageTotals struct {
	InputTokens       uint64
	CachedInputTokens uint64
	OutputTokens      uint64
}

type llmUsageObservation struct {
	authorization organization.OrganizationProviderProxyAccess
	providerID    string
	model         string
	startedAt     time.Time
	queueWait     time.Duration
}

func (handler *LLMHandler) newUsageObservation(ctx *gin.Context, resolved *requestResolver) (*llmUsageObservation, bool) {
	if handler == nil || resolved == nil {
		return nil, false
	}
	authorization, ok := llmProxyAuthorizationFromContext(ctx.Request.Context())
	if !ok {
		return nil, false
	}
	return &llmUsageObservation{
		authorization: authorization.OrganizationProviderProxyAccess,
		providerID:    strings.TrimSpace(resolved.Provider.Name()),
		model:         strings.TrimSpace(resolved.Model),
		startedAt:     time.Now().UTC(),
	}, true
}

func (handler *LLMHandler) enforceQuota(ctx *gin.Context, observation *llmUsageObservation) bool {
	if handler == nil || handler.usageTracker == nil || observation == nil {
		return true
	}
	decision, err := handler.usageTracker.CheckOrganizationLLMQuota(ctx.Request.Context(), observation.authorization, observation.model)
	if err != nil {
		ctx.JSON(http.StatusServiceUnavailable, gin.H{"error": "llm quota check is unavailable"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "internal")
		return false
	}
	if decision.Allowed {
		return true
	}
	ctx.JSON(http.StatusTooManyRequests, gin.H{
		"error":           "llm daily quota exceeded",
		"dailyQuotaUsd":   decision.DailyQuotaUSD,
		"currentSpentUsd": decision.CurrentSpentUSD,
	})
	handler.recordObservation(ctx, observation, llmUsageTotals{}, "rate_limit")
	return false
}

func (handler *LLMHandler) recordObservation(ctx *gin.Context, observation *llmUsageObservation, usage llmUsageTotals, errorType string) {
	if handler == nil || handler.usageTracker == nil || observation == nil {
		return
	}
	recordCtx := context.WithoutCancel(ctx.Request.Context())
	_ = handler.usageTracker.RecordOrganizationLLMUsage(recordCtx, organization.OrganizationLLMUsageRecord{
		Authorization:     observation.authorization,
		ProviderID:        observation.providerID,
		Model:             observation.model,
		RequestCount:      1,
		InputTokens:       usage.InputTokens,
		CachedInputTokens: usage.CachedInputTokens,
		OutputTokens:      usage.OutputTokens,
		Duration:          time.Since(observation.startedAt),
		QueueWait:         observation.queueWait,
		ErrorType:         errorType,
		OccurredAt:        time.Now().UTC(),
	})
}

func (handler *LLMHandler) copyResponseWithObservation(ctx *gin.Context, resp *http.Response, observation *llmUsageObservation) {
	if resp == nil {
		ctx.JSON(http.StatusBadGateway, gin.H{"error": "upstream response is unavailable"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "provider")
		return
	}
	for key, values := range resp.Header {
		ctx.Header(key, strings.Join(values, ", "))
	}

	if isStreamingResponse(resp.Header.Get("Content-Type")) {
		ctx.Status(resp.StatusCode)
		usage, err := copyStreamingResponse(ctx.Writer, resp.Body, nil)
		if err != nil {
			handler.recordObservation(ctx, observation, usage, "network")
			return
		}
		handler.recordObservation(ctx, observation, usage, classifyHTTPErrorType(resp.StatusCode, nil))
		return
	}

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		ctx.JSON(http.StatusBadGateway, gin.H{"error": "failed to read upstream response"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "network")
		return
	}
	ctx.Status(resp.StatusCode)
	_, _ = ctx.Writer.Write(raw)
	handler.recordObservation(ctx, observation, extractUsageFromJSON(raw), classifyHTTPErrorType(resp.StatusCode, raw))
}

func (handler *LLMHandler) copyOpenAIChatCompletionResponseWithObservation(ctx *gin.Context, resp *http.Response, observation *llmUsageObservation) {
	if resp == nil {
		ctx.JSON(http.StatusBadGateway, gin.H{"error": "upstream response is unavailable"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "provider")
		return
	}
	for key, values := range resp.Header {
		ctx.Header(key, strings.Join(values, ", "))
	}

	if isStreamingResponse(resp.Header.Get("Content-Type")) {
		ctx.Status(resp.StatusCode)
		normalizer := newOpenAIChatCompletionSSENormalizer()
		usage, err := copyStreamingResponse(ctx.Writer, resp.Body, normalizer.TransformLine)
		if err != nil {
			handler.recordObservation(ctx, observation, usage, "network")
			return
		}
		if handler != nil && handler.logger != nil && normalizer.didNormalize() {
			handler.logger.Warn("llm chat completion stream normalized tool call ids",
				"synthetic_ids", normalizer.syntheticIDs,
				"canonicalized_ids", normalizer.canonicalizedIDs,
			)
		}
		handler.recordObservation(ctx, observation, usage, classifyHTTPErrorType(resp.StatusCode, nil))
		return
	}

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		ctx.JSON(http.StatusBadGateway, gin.H{"error": "failed to read upstream response"})
		handler.recordObservation(ctx, observation, llmUsageTotals{}, "network")
		return
	}
	ctx.Status(resp.StatusCode)
	_, _ = ctx.Writer.Write(raw)
	handler.recordObservation(ctx, observation, extractUsageFromJSON(raw), classifyHTTPErrorType(resp.StatusCode, raw))
}

func copyStreamingResponse(writer io.Writer, reader io.Reader, transform streamingResponseLineTransformer) (llmUsageTotals, error) {
	usage := llmUsageTotals{}
	if writer == nil || reader == nil {
		return usage, nil
	}
	buffered := bufio.NewReader(reader)
	for {
		line, err := buffered.ReadBytes('\n')
		if len(line) > 0 {
			if transform != nil {
				line = transform(line)
			}
			if _, writeErr := writer.Write(line); writeErr != nil {
				return usage, writeErr
			}
			if flusher, ok := writer.(http.Flusher); ok {
				flusher.Flush()
			}
			mergeUsageTotals(&usage, extractUsageFromSSELine(line))
		}
		if errors.Is(err, io.EOF) {
			return usage, nil
		}
		if err != nil {
			return usage, err
		}
	}
}

type openAIChatToolCallKey struct {
	choiceIndex int
	toolIndex   int
}

type openAIChatCompletionSSENormalizer struct {
	toolCallIDs      map[openAIChatToolCallKey]string
	syntheticIDs     int
	canonicalizedIDs int
}

func newOpenAIChatCompletionSSENormalizer() *openAIChatCompletionSSENormalizer {
	return &openAIChatCompletionSSENormalizer{
		toolCallIDs: make(map[openAIChatToolCallKey]string),
	}
}

func (normalizer *openAIChatCompletionSSENormalizer) didNormalize() bool {
	if normalizer == nil {
		return false
	}
	return normalizer.syntheticIDs > 0 || normalizer.canonicalizedIDs > 0
}

func (normalizer *openAIChatCompletionSSENormalizer) TransformLine(line []byte) []byte {
	if normalizer == nil || len(line) == 0 {
		return line
	}
	trimmedLine := bytes.TrimRight(line, "\r\n")
	newline := line[len(trimmedLine):]
	trimmed := bytes.TrimSpace(trimmedLine)
	if !bytes.HasPrefix(trimmed, []byte("data:")) {
		return line
	}
	data := bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:")))
	if len(data) == 0 || bytes.Equal(data, []byte("[DONE]")) {
		return line
	}

	var payload map[string]any
	if err := json.Unmarshal(data, &payload); err != nil {
		return line
	}
	if !normalizer.normalizePayload(payload) {
		return line
	}
	encoded, err := json.Marshal(payload)
	if err != nil {
		return line
	}
	rewritten := append([]byte("data: "), encoded...)
	rewritten = append(rewritten, newline...)
	return rewritten
}

func (normalizer *openAIChatCompletionSSENormalizer) normalizePayload(payload map[string]any) bool {
	if normalizer == nil || payload == nil {
		return false
	}
	choices, ok := payload["choices"].([]any)
	if !ok {
		return false
	}
	changed := false
	for choiceOffset, rawChoice := range choices {
		choice, ok := rawChoice.(map[string]any)
		if !ok {
			continue
		}
		choiceIndex := choiceOffset
		if _, exists := choice["index"]; exists {
			choiceIndex = intValue(choice["index"])
		}
		delta, ok := choice["delta"].(map[string]any)
		if !ok {
			continue
		}
		toolCalls, ok := delta["tool_calls"].([]any)
		if !ok {
			continue
		}
		for toolOffset, rawToolCall := range toolCalls {
			toolCall, ok := rawToolCall.(map[string]any)
			if !ok {
				continue
			}
			toolIndex := toolOffset
			if _, exists := toolCall["index"]; exists {
				toolIndex = intValue(toolCall["index"])
			}
			key := openAIChatToolCallKey{choiceIndex: choiceIndex, toolIndex: toolIndex}
			currentID, _ := toolCall["id"].(string)
			currentID = strings.TrimSpace(currentID)
			canonicalID := normalizer.toolCallIDs[key]
			if canonicalID == "" {
				if currentID == "" {
					canonicalID = randomID("call")
					toolCall["id"] = canonicalID
					normalizer.syntheticIDs++
					changed = true
				} else {
					canonicalID = currentID
				}
				normalizer.toolCallIDs[key] = canonicalID
				continue
			}
			if currentID != canonicalID {
				toolCall["id"] = canonicalID
				normalizer.canonicalizedIDs++
				changed = true
			}
		}
	}
	return changed
}

func extractUsageFromSSELine(line []byte) llmUsageTotals {
	trimmed := bytes.TrimSpace(line)
	if !bytes.HasPrefix(trimmed, []byte("data:")) {
		return llmUsageTotals{}
	}
	data := bytes.TrimSpace(bytes.TrimPrefix(trimmed, []byte("data:")))
	if len(data) == 0 || bytes.Equal(data, []byte("[DONE]")) {
		return llmUsageTotals{}
	}
	return extractUsageFromJSON(data)
}

func extractUsageFromJSON(raw []byte) llmUsageTotals {
	var payload map[string]any
	if err := json.Unmarshal(raw, &payload); err != nil {
		return llmUsageTotals{}
	}
	return extractUsageFromPayload(payload)
}

func extractUsageFromPayload(payload map[string]any) llmUsageTotals {
	if payload == nil {
		return llmUsageTotals{}
	}
	if usage, ok := payload["usage"].(map[string]any); ok {
		return extractUsageFromMap(usage)
	}
	if response, ok := payload["response"].(map[string]any); ok {
		if usage, ok := response["usage"].(map[string]any); ok {
			return extractUsageFromMap(usage)
		}
	}
	if message, ok := payload["message"].(map[string]any); ok {
		if usage, ok := message["usage"].(map[string]any); ok {
			return extractUsageFromMap(usage)
		}
	}
	return llmUsageTotals{}
}

func extractUsageFromMap(usage map[string]any) llmUsageTotals {
	inputTokens := intValue(usage["input_tokens"])
	if inputTokens == 0 {
		inputTokens = intValue(usage["prompt_tokens"])
	}
	outputTokens := intValue(usage["output_tokens"])
	if outputTokens == 0 {
		outputTokens = intValue(usage["completion_tokens"])
	}
	cachedTokens := intValue(usage["cache_read_input_tokens"])
	if cachedTokens == 0 {
		if details, ok := usage["input_tokens_details"].(map[string]any); ok {
			cachedTokens = intValue(details["cached_tokens"])
		}
	}
	if cachedTokens == 0 {
		if details, ok := usage["prompt_tokens_details"].(map[string]any); ok {
			cachedTokens = intValue(details["cached_tokens"])
		}
	}
	return llmUsageTotals{
		InputTokens:       uint64(maxInt(inputTokens, 0)),
		CachedInputTokens: uint64(maxInt(cachedTokens, 0)),
		OutputTokens:      uint64(maxInt(outputTokens, 0)),
	}
}

func mergeUsageTotals(target *llmUsageTotals, incoming llmUsageTotals) {
	if target == nil {
		return
	}
	if incoming.InputTokens > 0 {
		target.InputTokens = incoming.InputTokens
	}
	if incoming.CachedInputTokens > 0 {
		target.CachedInputTokens = incoming.CachedInputTokens
	}
	if incoming.OutputTokens > 0 {
		target.OutputTokens = incoming.OutputTokens
	}
}

func classifyHTTPErrorType(statusCode int, body []byte) string {
	if statusCode < http.StatusBadRequest {
		return ""
	}
	bodyText := strings.ToLower(string(body))
	switch statusCode {
	case http.StatusUnauthorized, http.StatusForbidden:
		return "auth"
	case http.StatusTooManyRequests:
		return "rate_limit"
	case http.StatusRequestTimeout, http.StatusGatewayTimeout:
		return "timeout"
	case http.StatusNotFound:
		if strings.Contains(bodyText, "model") {
			return "model_not_found"
		}
		return "bad_request"
	case http.StatusBadRequest:
		if strings.Contains(bodyText, "unsupported model") || strings.Contains(bodyText, "model") && strings.Contains(bodyText, "not found") {
			return "model_not_found"
		}
		return "bad_request"
	default:
		if statusCode >= http.StatusInternalServerError {
			return "provider"
		}
		if strings.Contains(bodyText, "model") && strings.Contains(bodyText, "not found") {
			return "model_not_found"
		}
		return "bad_request"
	}
}

func classifyUpstreamTransportError(err error) string {
	if err == nil {
		return ""
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return "timeout"
	}
	var netErr net.Error
	if errors.As(err, &netErr) {
		if netErr.Timeout() {
			return "timeout"
		}
		return "network"
	}
	if strings.Contains(strings.ToLower(err.Error()), "timeout") {
		return "timeout"
	}
	return "network"
}

func maxInt(value int, minimum int) int {
	if value < minimum {
		return minimum
	}
	return value
}