package extproc

import (
	"testing"

	"github.com/openai/openai-go"
	"github.com/stretchr/testify/assert"

	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/cache"
	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/config"
)

// =====================================================================
// Mock
// =====================================================================

type mockStreamingCache struct {
	addEntryCalled     bool
	addPendingCalled   bool
	findSimilarCalled  bool
	updateCalled       bool
	updateRequestID    string
	updateResponseBody []byte
	addEntryErr        error
	updateErr          error
	lastTTLSeconds     int
}

func (m *mockStreamingCache) IsEnabled() bool { return true }

func (m *mockStreamingCache) CheckConnection() error { return nil }

func (m *mockStreamingCache) AddPendingRequest(
	_ string,
	_ string,
	_ string,
	_ []byte,
	_ int,
) error {
	m.addPendingCalled = true
	return nil
}

func (m *mockStreamingCache) UpdateWithResponse(requestID string, responseBody []byte, ttlSeconds int) error {
	m.updateCalled = true
	m.updateRequestID = requestID
	m.updateResponseBody = append([]byte(nil), responseBody...)
	m.lastTTLSeconds = ttlSeconds
	return m.updateErr
}

func (m *mockStreamingCache) AddEntry(
	_ string,
	_ string,
	_ string,
	_ []byte,
	_ []byte,
	ttlSeconds int,
) error {
	m.addEntryCalled = true
	m.lastTTLSeconds = ttlSeconds
	return m.addEntryErr
}

func (m *mockStreamingCache) FindSimilar(_ string, _ string) ([]byte, bool, error) {
	return nil, false, nil
}

func (m *mockStreamingCache) FindSimilarWithThreshold(
	_ string,
	_ string,
	_ float32,
) ([]byte, bool, error) {
	m.findSimilarCalled = true
	return nil, false, nil
}

func (m *mockStreamingCache) LastSimilarity() float32 { return 0 }

func (m *mockStreamingCache) Close() error { return nil }

func (m *mockStreamingCache) GetStats() cache.CacheStats { return cache.CacheStats{} }

// =====================================================================
// NON-STREAMING: updateResponseCache
// =====================================================================

func TestUpdateResponseCache_SkipsWhenDecisionCacheDisabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "no-cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		VSRSelectedDecisionName: "no-cache-decision",
	}

	router.updateResponseCache(ctx, []byte(`{"ok":true}`))
	assert.False(t, mockCache.updateCalled, "should not store response when decision has no semantic-cache plugin")
}

func TestUpdateResponseCache_SkipsWhenDecisionCacheExplicitlyDisabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "disabled-cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
					Plugins: []config.DecisionPlugin{
						{
							Type:          "semantic-cache",
							Configuration: config.MustStructuredPayload(map[string]interface{}{"enabled": false}),
						},
					},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		VSRSelectedDecisionName: "disabled-cache-decision",
	}

	router.updateResponseCache(ctx, []byte(`{"ok":true}`))
	assert.False(t, mockCache.updateCalled, "should not store response when decision has semantic-cache disabled")
}

func TestUpdateResponseCache_StoresWhenDecisionCacheEnabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
					Plugins: []config.DecisionPlugin{
						{
							Type:          "semantic-cache",
							Configuration: config.MustStructuredPayload(map[string]interface{}{"enabled": true}),
						},
					},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		VSRSelectedDecisionName: "cache-decision",
	}

	router.updateResponseCache(ctx, []byte(`{"ok":true}`))
	assert.True(t, mockCache.updateCalled, "should store response when decision has semantic-cache enabled")
}

func TestUpdateResponseCache_StoresWhenNoDecisionSelectedAndNoDecisionsConfigured(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		VSRSelectedDecisionName: "",
	}

	router.updateResponseCache(ctx, []byte(`{"ok":true}`))
	assert.True(t, mockCache.updateCalled, "should store response when no decision is selected (global cache applies)")
}

func TestUpdateResponseCache_SkipsWhenNoDecisionSelectedButDecisionsConfigured(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "default-route",
					ModelRefs: []config.ModelRef{{Model: "test"}},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		VSRSelectedDecisionName: "",
	}

	router.updateResponseCache(ctx, []byte(`{"ok":true}`))
	assert.False(t, mockCache.updateCalled, "should not store response when decisions exist but no decision matched")
}

// =====================================================================
// STREAMING: cacheReconstructedStreamingResponse
// =====================================================================

// Regression for #2030: streaming finalize must update the pending entry in
// place (by request_id), not mint a new key via AddEntry that orphans it.
func TestCacheReconstructedStreamingResponseUpdatesInPlace(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:    "req-1",
		RequestModel: "test-model",
		RequestQuery: "hello",
	}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.True(t, mockCache.updateCalled, "streaming finalize must update the pending entry in place")
	assert.False(t, mockCache.addEntryCalled, "streaming finalize must not mint a new cache key via AddEntry")
	assert.Equal(t, "req-1", mockCache.updateRequestID)
	assert.JSONEq(t, `{"ok":true}`, string(mockCache.updateResponseBody))
}

// Regression for #2030: ctx.RequestModel is overwritten from the client alias
// to the resolved model by finalize time, so finalize must key off request_id.
func TestCacheReconstructedStreamingResponseUpdatesAfterModelResolved(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:    "req-1",
		RequestModel: "doubao-1-5-lite-32k-250115",
		RequestQuery: "who discovered gravity?",
	}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.True(t, mockCache.updateCalled)
	assert.False(t, mockCache.addEntryCalled)
	assert.Equal(t, "req-1", mockCache.updateRequestID)
}

func TestCacheReconstructedStreamingResponseUpdatesWithoutQueryMetadata(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{RequestID: "req-1"}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled)
	assert.True(t, mockCache.updateCalled)
}

func TestCacheReconstructedStreamingResponseSkipsWithoutRequestID(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{RequestModel: "test-model", RequestQuery: "hello"}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled)
	assert.False(t, mockCache.updateCalled)
}

// =====================================================================
// STREAMING: cacheReconstructedStreamingResponse — per-decision
// =====================================================================

func TestCacheReconstructedStreamingResponse_SkipsWhenDecisionCacheDisabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "no-cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		RequestModel:            "test-model",
		RequestQuery:            "hello",
		VSRSelectedDecisionName: "no-cache-decision",
	}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled, "should not call AddEntry when decision has no semantic-cache plugin")
	assert.False(t, mockCache.updateCalled, "should not call UpdateWithResponse when decision has no semantic-cache plugin")
}

func TestCacheReconstructedStreamingResponse_SkipsWhenDecisionCacheExplicitlyDisabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "disabled-cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
					Plugins: []config.DecisionPlugin{
						{
							Type:          "semantic-cache",
							Configuration: config.MustStructuredPayload(map[string]interface{}{"enabled": false}),
						},
					},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		RequestModel:            "test-model",
		RequestQuery:            "hello",
		VSRSelectedDecisionName: "disabled-cache-decision",
	}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled, "should not call AddEntry when decision has semantic-cache disabled")
	assert.False(t, mockCache.updateCalled, "should not call UpdateWithResponse when decision has semantic-cache disabled")
}

func TestCacheReconstructedStreamingResponse_StoresWhenDecisionCacheEnabled(t *testing.T) {
	mockCache := &mockStreamingCache{}
	cfg := &config.RouterConfig{
		SemanticCache: config.SemanticCache{Enabled: true},
		IntelligentRouting: config.IntelligentRouting{
			Decisions: []config.Decision{
				{
					Name:      "cache-decision",
					ModelRefs: []config.ModelRef{{Model: "test"}},
					Plugins: []config.DecisionPlugin{
						{
							Type:          "semantic-cache",
							Configuration: config.MustStructuredPayload(map[string]interface{}{"enabled": true}),
						},
					},
				},
			},
		},
	}
	router := &OpenAIRouter{Cache: mockCache, Config: cfg}
	ctx := &RequestContext{
		RequestID:               "req-1",
		RequestModel:            "test-model",
		RequestQuery:            "hello",
		VSRSelectedDecisionName: "cache-decision",
	}

	err := router.cacheReconstructedStreamingResponse(ctx, []byte(`{"ok":true}`))
	assert.NoError(t, err)
	assert.True(t, mockCache.updateCalled, "should finalize the pending entry when decision has semantic-cache enabled")
	assert.False(t, mockCache.addEntryCalled, "should not mint a new cache key via AddEntry")
}

// =====================================================================
// STREAMING: cacheStreamingResponse — personalized context
// =====================================================================

func TestCacheStreamingSkippedWhenRAGContextPresent(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:           "req-rag",
		RequestModel:        "test-model",
		RequestQuery:        "what is our refund policy",
		StreamingComplete:   true,
		StreamingContent:    "Based on your documents...",
		StreamingMetadata:   map[string]interface{}{"id": "chatcmpl-1", "model": "test-model", "created": int64(1234567890)},
		RAGRetrievedContext: "Internal policy doc: refunds within 30 days",
	}

	err := router.cacheStreamingResponse(ctx)
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled, "cache write must be skipped when RAG context is present")
	assert.False(t, mockCache.updateCalled, "cache update must be skipped when RAG context is present")
}

func TestCacheStreamingSkippedWhenMemoryContextPresent(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:         "req-mem",
		RequestModel:      "test-model",
		RequestQuery:      "remind me what we discussed",
		StreamingComplete: true,
		StreamingContent:  "Last time you mentioned...",
		StreamingMetadata: map[string]interface{}{"id": "chatcmpl-2", "model": "test-model", "created": int64(1234567890)},
		MemoryContext:     "User previously discussed project deadlines",
	}

	err := router.cacheStreamingResponse(ctx)
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled, "cache write must be skipped when memory context is present")
	assert.False(t, mockCache.updateCalled, "cache update must be skipped when memory context is present")
}

func TestCacheStreamingSkippedWhenPIIDetected(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:         "req-pii",
		RequestModel:      "test-model",
		RequestQuery:      "parse this email",
		StreamingComplete: true,
		StreamingContent:  "The email contains...",
		StreamingMetadata: map[string]interface{}{"id": "chatcmpl-3", "model": "test-model", "created": int64(1234567890)},
		PIIDetected:       true,
	}

	err := router.cacheStreamingResponse(ctx)
	assert.NoError(t, err)
	assert.False(t, mockCache.addEntryCalled, "cache write must be skipped when PII is detected")
	assert.False(t, mockCache.updateCalled, "cache update must be skipped when PII is detected")
}

func TestCacheStreamingAllowedForGenericRequest(t *testing.T) {
	mockCache := &mockStreamingCache{}
	router := &OpenAIRouter{Cache: mockCache}
	ctx := &RequestContext{
		RequestID:         "req-generic",
		RequestModel:      "test-model",
		RequestQuery:      "explain quantum computing",
		StreamingComplete: true,
		StreamingContent:  "Quantum computing uses qubits...",
		StreamingMetadata: map[string]interface{}{"id": "chatcmpl-4", "model": "test-model", "created": int64(1234567890)},
	}

	err := router.cacheStreamingResponse(ctx)
	assert.NoError(t, err)
	assert.True(t, mockCache.updateCalled,
		"cache write must proceed for generic requests without personalized context")
}

func TestBuildReconstructedStreamingResponseIncludesToolCallsForReplay(t *testing.T) {
	ctx := &RequestContext{
		StreamingMetadata: map[string]interface{}{
			"id":            "chatcmpl-123",
			"model":         "test-model",
			"created":       int64(1234567890),
			"finish_reason": "tool_calls",
		},
		StreamingToolCalls: map[int]*StreamingToolCallState{
			0: {
				ID:        "call_weather",
				Name:      "get_weather",
				Arguments: "{\"location\":\"San Francisco\"}",
			},
		},
	}

	body, err := buildReconstructedStreamingResponse(ctx, openai.CompletionUsage{
		PromptTokens:     10,
		CompletionTokens: 2,
		TotalTokens:      12,
	}, true)
	assert.NoError(t, err)
	assert.JSONEq(t, `{
		"id": "chatcmpl-123",
		"object": "chat.completion",
		"created": 1234567890,
		"model": "test-model",
		"choices": [
			{
				"index": 0,
				"message": {
					"role": "assistant",
					"content": null,
					"tool_calls": [
						{
							"id": "call_weather",
							"type": "function",
							"function": {
								"name": "get_weather",
								"arguments": "{\"location\":\"San Francisco\"}"
							}
						}
					]
				},
				"finish_reason": "tool_calls"
			}
		],
		"usage": {
			"prompt_tokens": 10,
			"completion_tokens": 2,
			"total_tokens": 12
		}
	}`, string(body))
}

func TestBuildReconstructedStreamingResponseSkipsToolCallsOnlyForCache(t *testing.T) {
	ctx := &RequestContext{
		StreamingMetadata: map[string]interface{}{
			"id":      "chatcmpl-123",
			"model":   "test-model",
			"created": int64(1234567890),
		},
		StreamingToolCalls: map[int]*StreamingToolCallState{
			0: {
				ID:        "call_weather",
				Name:      "get_weather",
				Arguments: "{\"location\":\"San Francisco\"}",
			},
		},
	}

	_, err := buildReconstructedStreamingResponse(ctx, openai.CompletionUsage{}, false)
	assert.ErrorIs(t, err, errSkipStreamingCache)
}
