package extproc

import (
	"encoding/json"
	"errors"
	"sort"

	"github.com/openai/openai-go"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"

	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability/logging"
	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability/metrics"
	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability/tracing"
)

// =====================================================================
// NON-STREAMING
// =====================================================================

func (r *OpenAIRouter) updateResponseCache(ctx *RequestContext, responseBody []byte) {
	if ctx.RequestID == "" || responseBody == nil {
		return
	}
	decisionName := ctx.VSRSelectedDecisionName
	if !r.semanticCacheEnabledForScope(decisionName) {
		return
	}

	// Status gate: never cache a non-2xx upstream body. Caching an error
	// (4xx/5xx/3xx) would let a later cache hit replay it as an HTTP 200,
	// freezing a transient upstream failure for the whole TTL.
	if skip, reason := shouldSkipCacheWriteForStatus(ctx); skip {
		metrics.RecordCacheWriteSkipped(reason)
		logging.Infof("Skipping cache write for request ID %s: upstream status %d is not 2xx (reason=%s)", ctx.RequestID, ctx.UpstreamStatusCode, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return
	}

	// Retention drop gate: DSL-author-explicit signal takes precedence over
	// the implicit personalized-context heuristic. Order is locked in
	// docs/design (issue #1747 §2.8): scope -> retention -> personalized -> TTL -> write.
	if skip, reason := ShouldSkipCacheWrite(ctx); skip {
		metrics.RecordCacheWriteSkipped(reason)
		logRetentionSkip(ctx, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return
	}

	if ok, reason := ctx.HasPersonalizedContext(); ok {
		metrics.RecordCacheWriteSkipped(reason)
		logging.Infof("Skipping cache write for request ID %s: response has personalized context (reason=%s)", ctx.RequestID, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return
	}

	ttlSeconds := -1
	if r != nil && r.Config != nil {
		ttlSeconds = r.Config.GetCacheTTLSecondsForDecision(decisionName)
	}
	// Retention ttl_turns (when set) overrides the decision/global default,
	// scoping this entry to a turn-bounded lifetime (§2.8 order: ... -> TTL -> write).
	ttlSeconds = applyRetentionTTLOverride(ttlSeconds, ctx)
	if err := r.Cache.UpdateWithResponse(ctx.RequestID, responseBody, ttlSeconds); err != nil {
		logging.Errorf("Error updating cache: %v", err)
		return
	}
	logging.Infof("Cache updated for request ID: %s", ctx.RequestID)
}

// =====================================================================
// STREAMING
// =====================================================================

// cacheStreamingResponse reconstructs a ChatCompletion from accumulated chunks and caches it.
func (r *OpenAIRouter) cacheStreamingResponse(ctx *RequestContext) error {
	if err := validateStreamingCachePreconditions(ctx); err != nil {
		return nil
	}

	decisionName := ctx.VSRSelectedDecisionName
	if !r.semanticCacheEnabledForScope(decisionName) {
		return nil
	}

	// Status gate: never cache a non-2xx upstream body (see updateResponseCache).
	// Streaming errors are usually filtered by the preconditions above, but an
	// error that still carries content must not be frozen into the cache either.
	if skip, reason := shouldSkipCacheWriteForStatus(ctx); skip {
		metrics.RecordCacheWriteSkipped(reason)
		logging.Infof("Skipping streaming cache write for request ID %s: upstream status %d is not 2xx (reason=%s)", ctx.RequestID, ctx.UpstreamStatusCode, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return nil
	}

	// Retention drop gate runs before any reconstruction work so that we never
	// even build the cache entry for a decision that asked us to skip the
	// write. Same precedence rule as the non-streaming path (§2.8).
	if skip, reason := ShouldSkipCacheWrite(ctx); skip {
		metrics.RecordCacheWriteSkipped(reason)
		logRetentionSkip(ctx, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return nil
	}

	if ok, reason := ctx.HasPersonalizedContext(); ok {
		metrics.RecordCacheWriteSkipped(reason)
		logging.Infof("Skipping cache write for streaming request ID %s: response has personalized context (reason=%s)", ctx.RequestID, reason)
		if span := trace.SpanFromContext(ctx.TraceContext); span.IsRecording() {
			span.SetAttributes(attribute.String(tracing.AttrCacheWriteSkippedReason, reason))
		}
		return nil
	}

	usage := extractStreamingUsage(ctx)
	reconstructedJSON, err := buildReconstructedStreamingResponse(ctx, usage, false)
	if err != nil {
		if errors.Is(err, errSkipStreamingCache) {
			return nil
		}
		return err
	}

	return r.cacheReconstructedStreamingResponse(ctx, reconstructedJSON)
}

func validateStreamingCachePreconditions(ctx *RequestContext) error {
	switch {
	case !ctx.StreamingComplete:
		logging.Warnf("Stream not completed (no [DONE] marker), skipping cache")
	case ctx.StreamingAborted:
		logging.Warnf("Stream was aborted, skipping cache")
	case ctx.StreamingContent == "":
		logging.Warnf("Streaming response has no content, skipping cache")
	case ctx.StreamingMetadata["id"] == nil || ctx.StreamingMetadata["model"] == nil:
		logging.Warnf("Streaming response missing required metadata, skipping cache")
	default:
		return nil
	}
	return errSkipStreamingCache
}

var errSkipStreamingCache = &streamingCacheSkipError{}

type streamingCacheSkipError struct{}

func (e *streamingCacheSkipError) Error() string { return "skip streaming cache" }

func buildReconstructedStreamingResponse(
	ctx *RequestContext,
	usage openai.CompletionUsage,
	includeToolCalls bool,
) ([]byte, error) {
	if ctx == nil {
		return nil, errSkipStreamingCache
	}

	id, idOK := ctx.StreamingMetadata["id"].(string)
	model, modelOK := ctx.StreamingMetadata["model"].(string)
	created, createdOK := ctx.StreamingMetadata["created"].(int64)
	if !idOK || !modelOK || !createdOK {
		logging.Warnf("Streaming response missing metadata required for reconstruction, skipping")
		return nil, errSkipStreamingCache
	}

	finishReason := "stop"
	if finishReasonValue, ok := ctx.StreamingMetadata["finish_reason"].(string); ok && finishReasonValue != "" {
		finishReason = finishReasonValue
	}

	message := map[string]interface{}{
		"role": "assistant",
	}
	if ctx.StreamingContent != "" {
		message["content"] = ctx.StreamingContent
	} else {
		message["content"] = nil
	}
	// Preserve reasoning so a cache hit returns the same reasoning_content the
	// live stream delivered (reasoning models stream it under delta.reasoning_content).
	if ctx.StreamingReasoning != "" {
		message["reasoning_content"] = ctx.StreamingReasoning
	}

	toolCalls := buildStreamingResponseToolCalls(ctx)
	if includeToolCalls && len(toolCalls) > 0 {
		message["tool_calls"] = toolCalls
		if finishReason == "stop" {
			finishReason = "tool_calls"
		}
	}

	if ctx.StreamingContent == "" && (!includeToolCalls || len(toolCalls) == 0) {
		logging.Warnf("Reconstructed response has no valid assistant payload, skipping cache")
		return nil, errSkipStreamingCache
	}

	reconstructed := map[string]interface{}{
		"id":      id,
		"object":  "chat.completion",
		"created": created,
		"model":   model,
		"choices": []map[string]interface{}{
			{
				"index":         0,
				"message":       message,
				"finish_reason": finishReason,
			},
		},
		"usage": map[string]interface{}{
			"prompt_tokens":     usage.PromptTokens,
			"completion_tokens": usage.CompletionTokens,
			"total_tokens":      usage.TotalTokens,
		},
	}

	reconstructedJSON, err := json.Marshal(reconstructed)
	if err != nil {
		logging.Errorf("Failed to marshal reconstructed response: %v", err)
		return nil, err
	}
	return reconstructedJSON, nil
}

func buildStreamingResponseToolCalls(ctx *RequestContext) []map[string]interface{} {
	if ctx == nil || len(ctx.StreamingToolCalls) == 0 {
		return nil
	}

	indexes := make([]int, 0, len(ctx.StreamingToolCalls))
	for index := range ctx.StreamingToolCalls {
		indexes = append(indexes, index)
	}
	sort.Ints(indexes)

	toolCalls := make([]map[string]interface{}, 0, len(indexes))
	for _, index := range indexes {
		call := ctx.StreamingToolCalls[index]
		if call == nil || (call.ID == "" && call.Name == "" && call.Arguments == "") {
			continue
		}
		toolCalls = append(toolCalls, map[string]interface{}{
			"id":   call.ID,
			"type": "function",
			"function": map[string]interface{}{
				"name":      call.Name,
				"arguments": call.Arguments,
			},
		})
	}
	return toolCalls
}

func (r *OpenAIRouter) cacheReconstructedStreamingResponse(
	ctx *RequestContext,
	reconstructedJSON []byte,
) error {
	decisionName := ctx.VSRSelectedDecisionName
	if !r.semanticCacheEnabledForScope(decisionName) {
		return nil
	}

	ttlSeconds := -1
	if r != nil && r.Config != nil {
		ttlSeconds = r.Config.GetCacheTTLSecondsForDecision(decisionName)
	}
	// Retention ttl_turns (when set) overrides the decision/global default for
	// the reconstructed streaming entry, mirroring the non-streaming path.
	ttlSeconds = applyRetentionTTLOverride(ttlSeconds, ctx)

	if ctx.RequestID == "" {
		logging.Warnf("No request ID available, cannot cache streaming response")
		return nil
	}

	// Must finalize in place (by request_id) like the non-streaming path.
	// AddEntry would mint a new key and orphan the empty pending entry,
	// which then shadows KNN lookups and causes cache misses (issue #2030).
	return r.updateStreamingCacheEntry(ctx.RequestID, reconstructedJSON, ttlSeconds)
}

func (r *OpenAIRouter) updateStreamingCacheEntry(
	requestID string,
	reconstructedJSON []byte,
	ttlSeconds int,
) error {
	if err := r.Cache.UpdateWithResponse(requestID, reconstructedJSON, ttlSeconds); err != nil {
		logging.Errorf("Cache update failed for streaming %s: %v", requestID, err)
		return err
	}
	return nil
}
