package openai

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// TestChatCompletionsPassthroughProviderPoolGenerationPolicy verifies that
// provider-pool omitted-mode passthrough requests apply the catalog entry's
// generation policy (default_max_tokens, min_max_tokens, default_thinking_token_budget)
// before dispatching to the provider. The fixture omits a top-level model echo,
// so the response body remains byte-identical.
func TestChatCompletionsPassthroughProviderPoolGenerationPolicy(t *testing.T) {
	providerBody := `{"id":"cmpl-provider-1","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}`
	var gotProviderReq []byte
	provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotProviderReq, _ = io.ReadAll(r.Body)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(providerBody))
	}))
	defer provider.Close()

	fake := &providerFakeRunService{
		tunnelProviderURL:  provider.URL,
		tunnelServedTarget: "served-model",
	}
	// Enable strict output to trigger think=true from policy
	srv := NewServer(config.EdgeOpenAIConf{StrictOutput: true}, fake, nil)
	srv.SetModelCatalog([]config.ModelCatalogEntry{
		{
			ID:                         "pool-model",
			Providers:                  map[string]string{"prov-1": "served-model"},
			DefaultMaxTokens:           100,
			MinMaxTokens:               50,
			DefaultThinkingTokenBudget: 30,
		},
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"pool-model",
		"messages":[{"role":"user","content":"hello"}],
		"max_completion_tokens":20
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if got := w.Body.String(); got != providerBody {
		t.Fatalf("body not byte-identical to provider:\n got: %q\nwant: %q", got, providerBody)
	}

	var parsedReq map[string]interface{}
	if err := json.Unmarshal(gotProviderReq, &parsedReq); err != nil {
		t.Fatalf("failed to parse provider request: %v", err)
	}

	if model, ok := parsedReq["model"].(string); !ok || model != "served-model" {
		t.Errorf("expected model served-model, got %v", parsedReq["model"])
	}

	// MinMaxTokens (50) should reach max_tokens (overriding max_completion_tokens 20)
	if maxTokens, ok := parsedReq["max_tokens"].(float64); !ok || maxTokens != 50 {
		t.Errorf("expected max_tokens 50, got %v", parsedReq["max_tokens"])
	}
	if _, ok := parsedReq["max_completion_tokens"]; ok {
		t.Errorf("max_completion_tokens should be deleted")
	}

	// DefaultThinkingTokenBudget (30) and Think (true) should reach provider
	if budget, ok := parsedReq["thinking_token_budget"].(float64); !ok || budget != 30 {
		t.Errorf("expected thinking_token_budget 30, got %v", parsedReq["thinking_token_budget"])
	}
	if think, ok := parsedReq["think"].(bool); !ok || !think {
		t.Errorf("expected think true, got %v", parsedReq["think"])
	}
}

func TestChatCompletionsProviderPoolAppliesGenerationPolicy(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames:       staticProviderTunnelFrames(`{"ok":true}`),
		tunnelServedTarget: "Ornith-1.0-35B",
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                         "ornith:35b",
		DefaultMaxTokens:           32768,
		MinMaxTokens:               32768,
		DefaultThinkingTokenBudget: 8192,
		Providers:                  map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}
	srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ornith:35b",
		"messages":[{"role":"user","content":"hello"}],
		"max_tokens":4096
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if len(fake.reqsSnapshot()) != 0 {
		t.Fatalf("provider-pool generation policy must stay on tunnel path, got %d SubmitRun calls", len(fake.reqsSnapshot()))
	}
	bodies := fake.tunnelBodiesSnapshot()
	if len(bodies) != 1 {
		t.Fatalf("expected 1 tunnel body, got %d", len(bodies))
	}
	var providerReq map[string]any
	if err := json.Unmarshal(bodies[0], &providerReq); err != nil {
		t.Fatalf("provider body JSON: %v body=%s", err, bodies[0])
	}
	if providerReq["model"] != "Ornith-1.0-35B" {
		t.Fatalf("served model rewrite not applied: %+v", providerReq["model"])
	}
	if providerReq["max_tokens"].(float64) != 32768 {
		t.Fatalf("max_tokens policy not applied: %+v", providerReq)
	}
	if providerReq["thinking_token_budget"].(float64) != 8192 {
		t.Fatalf("thinking_token_budget policy not applied: %+v", providerReq)
	}
}

func TestChatCompletionsProviderPoolThinkingPolicyOverridesStrictOutputDisable(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames:       staticProviderTunnelFrames(`{"ok":true}`),
		tunnelServedTarget: "Ornith-1.0-35B",
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                         "ornith:35b",
		DefaultThinkingTokenBudget: 8192,
		Providers:                  map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}
	srv := NewServer(config.EdgeOpenAIConf{StrictOutput: true}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ornith:35b",
		"messages":[{"role":"user","content":"hello"}]
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if len(fake.reqsSnapshot()) != 0 {
		t.Fatalf("provider-pool thinking policy must stay on tunnel path, got %d SubmitRun calls", len(fake.reqsSnapshot()))
	}
	bodies := fake.tunnelBodiesSnapshot()
	if len(bodies) != 1 {
		t.Fatalf("expected 1 tunnel body, got %d", len(bodies))
	}
	var providerReq map[string]any
	if err := json.Unmarshal(bodies[0], &providerReq); err != nil {
		t.Fatalf("provider body JSON: %v body=%s", err, bodies[0])
	}
	if providerReq["think"] != true {
		t.Fatalf("provider-pool thinking policy should keep thinking enabled under strict output: %+v", providerReq)
	}
	if providerReq["thinking_token_budget"].(float64) != 8192 {
		t.Fatalf("thinking_token_budget policy not applied: %+v", providerReq)
	}
}

func TestChatCompletionsProviderPoolProviderNativeThinkingDisablesDefaultBudgetInjection(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames:       staticProviderTunnelFrames(`{"ok":true}`),
		tunnelServedTarget: "Ornith-1.0-35B",
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                         "ornith:35b",
		DefaultThinkingTokenBudget: 8192,
		Providers:                  map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}
	srv := NewServer(config.EdgeOpenAIConf{StrictOutput: true}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ornith:35b",
		"messages":[{"role":"user","content":"hello"}],
		"chat_template_kwargs":{"enable_thinking":false}
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	bodies := fake.tunnelBodiesSnapshot()
	if len(bodies) != 1 {
		t.Fatalf("expected 1 tunnel body, got %d", len(bodies))
	}
	var providerReq map[string]any
	if err := json.Unmarshal(bodies[0], &providerReq); err != nil {
		t.Fatalf("provider body JSON: %v body=%s", err, bodies[0])
	}
	ctk, ok := providerReq["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("chat_template_kwargs not preserved as object: %+v", providerReq)
	}
	if ctk["enable_thinking"] != false {
		t.Fatalf("provider-native enable_thinking must stay false: %+v", ctk)
	}
	if _, ok := providerReq["think"]; ok {
		t.Fatalf("catalog policy must not inject top-level think over provider-native thinking: %+v", providerReq)
	}
	if _, ok := providerReq["thinking_token_budget"]; ok {
		t.Fatalf("catalog policy must not inject thinking_token_budget over provider-native thinking: %+v", providerReq)
	}
}

func TestChatCompletionsProviderPoolPreservesLargerGenerationPolicyValues(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames:       staticProviderTunnelFrames(`{"ok":true}`),
		tunnelServedTarget: "Ornith-1.0-35B",
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                         "ornith:35b",
		DefaultMaxTokens:           32768,
		MinMaxTokens:               32768,
		DefaultThinkingTokenBudget: 8192,
		Providers:                  map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}
	srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"ornith:35b",
		"messages":[{"role":"user","content":"hello"}],
		"max_tokens":40000,
		"thinking_token_budget":2048
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if len(fake.reqsSnapshot()) != 0 {
		t.Fatalf("provider-pool generation policy must stay on tunnel path, got %d SubmitRun calls", len(fake.reqsSnapshot()))
	}
	bodies := fake.tunnelBodiesSnapshot()
	if len(bodies) != 1 {
		t.Fatalf("expected 1 tunnel body, got %d", len(bodies))
	}
	var providerReq map[string]any
	if err := json.Unmarshal(bodies[0], &providerReq); err != nil {
		t.Fatalf("provider body JSON: %v body=%s", err, bodies[0])
	}
	if providerReq["max_tokens"].(float64) != 40000 {
		t.Fatalf("larger max_tokens should be preserved: %+v", providerReq)
	}
	if providerReq["thinking_token_budget"].(float64) != 2048 {
		t.Fatalf("explicit thinking_token_budget should be preserved: %+v", providerReq)
	}
}

// TestResponsesProviderPoolPassthroughOmitsThinkingPolicy verifies that raw
// passthrough forwards the caller body as-is under strict output: the
// normalized-path thinking policy is not injected into the provider body.
func TestResponsesProviderPoolPassthroughOmitsThinkingPolicy(t *testing.T) {
	fake := &providerFakeRunService{
		tunnelFrames:       staticProviderTunnelFrames(`{"ok":true}`),
		tunnelServedTarget: "Ornith-1.0-35B",
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                         "ornith:35b",
		DefaultThinkingTokenBudget: 8192,
		Providers:                  map[string]string{"prov-vllm": "Ornith-1.0-35B"},
	}}
	srv := NewServer(config.EdgeOpenAIConf{TimeoutSec: 5, StrictOutput: true}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"ornith:35b",
		"input":"hello"
	}`))
	w := httptest.NewRecorder()
	srv.handleResponses(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if len(fake.reqsSnapshot()) != 0 {
		t.Fatalf("provider-pool /v1/responses must not call SubmitRun, got %d calls", len(fake.reqsSnapshot()))
	}
	bodies := fake.tunnelBodiesSnapshot()
	if len(bodies) != 1 {
		t.Fatalf("expected 1 tunnel body, got %d", len(bodies))
	}
	var providerReq map[string]any
	if err := json.Unmarshal(bodies[0], &providerReq); err != nil {
		t.Fatalf("provider body JSON: %v body=%s", err, bodies[0])
	}
	if _, ok := providerReq["think"]; ok {
		t.Fatalf("passthrough must not inject think: %+v", providerReq)
	}
	if _, ok := providerReq["thinking_token_budget"]; ok {
		t.Fatalf("passthrough must not inject thinking_token_budget: %+v", providerReq)
	}
}

// TestResponsesProviderPoolStrictOutputNormalizesAgentResponse verifies that the
// provider-pool normalized path preserves the strict-output contract from
// PrepareRun into completeResponse. A prompt containing an XML completion
// contract must have its response wrapped in the expected block (SDD D02).
func TestResponsesProviderPoolStrictOutputNormalizesAgentResponse(t *testing.T) {
	fake := &providerFakeRunService{
		poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
	}
	fake.poolRunFrames = make(chan *iop.RunEvent, 3)
	fake.poolRunFrames <- &iop.RunEvent{Type: "reasoning_delta", Delta: "hidden"}
	fake.poolRunFrames <- &iop.RunEvent{Type: "delta", Delta: "plain answer"}
	fake.poolRunFrames <- &iop.RunEvent{Type: "complete", RunId: "run-pool-normalized"}
	close(fake.poolRunFrames)

	catalog := []config.ModelCatalogEntry{
		{ID: "ollama-strict-model", Providers: map[string]string{"prov-ollama": "served-ollama"}},
	}
	srv := NewServer(config.EdgeOpenAIConf{StrictOutput: true}, fake, nil)
	srv.SetModelCatalog(catalog)

	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"ollama-strict-model",
		"input":"Once you've completed the user's task, you must use the attempt_completion tool to present the result.\n\n<attempt_completion>\n<result>done</result>\n</attempt_completion>\n\nfinish"
	}`))
	w := httptest.NewRecorder()
	srv.handleResponses(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}

	// PrepareRun must have been invoked (normalized path).
	if !fake.poolPrepareRunCalledSnapshot() {
		t.Fatal("PrepareRun was not called on normalized path")
	}

	// The prepared request must carry strict_output and the injected contract
	// instruction.
	if len(fake.reqsSnapshot()) != 1 {
		t.Fatalf("expected one prepared req, got %d", len(fake.reqsSnapshot()))
	}
	reqPrep := fake.reqsSnapshot()[0]
	if reqPrep.Metadata["strict_output"] != "true" {
		t.Fatalf("strict_output metadata lost in prepared request: %+v", reqPrep.Metadata)
	}
	if !strings.Contains(reqPrep.Prompt, "output exactly one <attempt_completion> block") {
		t.Fatalf("strict contract instruction not injected into prepared prompt:\n%s", reqPrep.Prompt)
	}

	var resp responsesResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode: %v", err)
	}
	want := "<attempt_completion>\n<result>plain answer</result>\n</attempt_completion>"
	if resp.OutputText != want {
		t.Fatalf("output_text:\ngot  %q\nwant %q", resp.OutputText, want)
	}
}

func TestChatCompletionsAcceptsThinkControlFields(t *testing.T) {
	fake := &fakeRunService{
		events: make(chan *iop.RunEvent, 1),
	}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)

	// Case 1: All valid think-control fields provided
	reqBody := `{
		"model":"client-model",
		"messages":[{"role":"user","content":"hi"}],
		"think": true,
		"reasoning_effort": "medium",
		"thinking_token_budget": 1024,
		"include_reasoning": false
	}`
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(reqBody))
	w := httptest.NewRecorder()
	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if fake.req.Input["think"] != true {
		t.Fatalf("think: got %v, want true", fake.req.Input["think"])
	}
	if fake.req.Input["reasoning_effort"] != "medium" {
		t.Fatalf("reasoning_effort: got %v, want 'medium'", fake.req.Input["reasoning_effort"])
	}

	budgetVal := fake.req.Input["thinking_token_budget"]
	switch v := budgetVal.(type) {
	case float64:
		if v != 1024 {
			t.Fatalf("thinking_token_budget: got %f, want 1024", v)
		}
	case int:
		if v != 1024 {
			t.Fatalf("thinking_token_budget: got %d, want 1024", v)
		}
	default:
		t.Fatalf("thinking_token_budget: got %T, want number type", budgetVal)
	}

	if fake.req.Input["include_reasoning"] != false {
		t.Fatalf("include_reasoning: got %v, want false", fake.req.Input["include_reasoning"])
	}

	// Case 2: Omitted think-control fields should NOT add keys to runInput
	fake2 := &fakeRunService{
		events: make(chan *iop.RunEvent, 1),
	}
	fake2.events <- &iop.RunEvent{Type: "complete"}
	srv2 := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake2, nil)

	reqBody2 := `{
		"model":"client-model",
		"messages":[{"role":"user","content":"hi"}]
	}`
	req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(reqBody2))
	w2 := httptest.NewRecorder()
	srv2.routes().ServeHTTP(w2, req2)

	if w2.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w2.Code, w2.Body.String())
	}
	for _, key := range []string{"think", "reasoning_effort", "thinking_token_budget", "include_reasoning"} {
		if _, ok := fake2.req.Input[key]; ok {
			t.Fatalf("expected key %q to be omitted from runInput, but it was present", key)
		}
	}
}

func TestChatCompletionsRejectsInvalidThinkControlFields(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", Target: "llama"}, &fakeRunService{}, nil)
	cases := []struct {
		name string
		body string
	}{
		{"negative budget", `{"model":"m","messages":[{"role":"user","content":"hi"}],"thinking_token_budget":-1}`},
		{"bad effort", `{"model":"m","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"super_high"}`},
		{"empty effort", `{"model":"m","messages":[{"role":"user","content":"hi"}],"reasoning_effort":""}`},
		{"think=false and reasoning_effort other than none", `{"model":"m","messages":[{"role":"user","content":"hi"}],"think":false,"reasoning_effort":"high"}`},
		{"thinking_token_budget conflicts with think=false", `{"model":"m","messages":[{"role":"user","content":"hi"}],"think":false,"thinking_token_budget":100}`},
		{"thinking_token_budget conflicts with reasoning_effort=none", `{"model":"m","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"none","thinking_token_budget":100}`},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(tc.body))
			w := httptest.NewRecorder()
			srv.routes().ServeHTTP(w, req)
			if w.Code != http.StatusBadRequest {
				t.Fatalf("got %d want 400, body=%s", w.Code, w.Body.String())
			}
		})
	}
}

func TestChatCompletionsRejectsNumCtxOptionsWrapper(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", Target: "llama"}, &fakeRunService{}, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}],
		"options":{"num_ctx":8192}
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusBadRequest {
		t.Fatalf("got %d want 400, body=%s", w.Code, w.Body.String())
	}
}