package openai

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

	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"

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

func TestGeminiIngressAuthenticatesAndStreamsThroughChatRoute(t *testing.T) {
	fake := &fakeRunService{events: bufferedRunEvents(
		&iop.RunEvent{Type: "delta", Delta: "hello"},
		&iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 3, OutputTokens: 2}, Metadata: map[string]string{"finish_reason": "length"}},
	)}
	srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama", Target: "provider-model"}, fake, nil)
	body := `{
		"systemInstruction":{"role":"user","parts":[{"text":"be concise"}]},
        "contents":[{"role":"user","parts":[{"text":"say hello"}]}],
        "generationConfig":{"candidateCount":1,"maxOutputTokens":16,"temperature":0.2,"topK":8,"topP":0.9,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
    }`
	req := httptest.NewRequest(http.MethodPost, "/gemini/gemini-direct/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", strings.NewReader(body))
	req.Header.Set("X-Goog-Api-Key", "iop-principal")
	w := httptest.NewRecorder()
	srv.routes().ServeHTTP(w, req)

	if w.Code != http.StatusOK || w.Header().Get("Content-Type") != "text/event-stream" {
		t.Fatalf("status/header: %d %q body=%s", w.Code, w.Header().Get("Content-Type"), w.Body.String())
	}
	if fake.req.ModelGroupKey != "gemini-direct" || fake.req.Target != "provider-model" {
		t.Fatalf("route bypassed Chat admission: %+v", fake.req)
	}
	if !strings.Contains(fake.req.Prompt, "system: be concise") || !strings.Contains(fake.req.Prompt, "user: say hello") {
		t.Fatalf("prompt conversion mismatch: %q", fake.req.Prompt)
	}
	options := fake.req.Input["options"].(map[string]any)
	if options["max_tokens"] != 16 {
		t.Fatalf("generation config mismatch: input=%+v", fake.req.Input)
	}
	if _, exists := options["top_k"]; exists {
		t.Fatalf("deprecated Gemini sampling option reached Chat upstream: input=%+v", fake.req.Input)
	}
	response := w.Body.String()
	for _, want := range []string{`"text":"hello"`, `"finishReason":"MAX_TOKENS"`} {
		if !strings.Contains(response, want) {
			t.Fatalf("missing %s in %s", want, response)
		}
	}
	if strings.Contains(response, "iop-principal") {
		t.Fatal("principal token leaked to response")
	}
}

func TestGeminiIngressRejectsAuthenticationAndShapeBeforeDispatch(t *testing.T) {
	base := `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`
	for _, tc := range []struct {
		name, path, body, bearer, key string
		status                        int
	}{
		{"missing key", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "", "", 401},
		{"conflicting auth", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", base, "Bearer other", "iop-principal", 401},
		{"wrong query", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent", base, "", "iop-principal", 400},
		{"duplicate member", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[],"contents":[]}`, "", "iop-principal", 400},
		{"two candidates", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"candidateCount":2}}`, "", "iop-principal", 400},
		{"invalid thinking budget", "/gemini/r/v1beta/models/gemini-3.6-flash:streamGenerateContent?alt=sse", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"generationConfig":{"thinkingConfig":{"thinkingBudget":-2}}}`, "", "iop-principal", 400},
	} {
		t.Run(tc.name, func(t *testing.T) {
			fake := &fakeRunService{}
			srv := NewServer(config.EdgeOpenAIConf{BearerToken: "iop-principal", Adapter: "ollama"}, fake, nil)
			req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
			if tc.bearer != "" {
				req.Header.Set("Authorization", tc.bearer)
			}
			if tc.key != "" {
				req.Header.Set("X-Goog-Api-Key", tc.key)
			}
			w := httptest.NewRecorder()
			srv.routes().ServeHTTP(w, req)
			if w.Code != tc.status {
				t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
			}
			var failure geminiErrorResponse
			if json.Unmarshal(w.Body.Bytes(), &failure) != nil || failure.Error.Code != tc.status {
				t.Fatalf("not a Gemini error: %s", w.Body.String())
			}
			if fake.req.ModelGroupKey != "" {
				t.Fatalf("unexpected dispatch: %+v", fake.req)
			}
		})
	}
}

func TestGeminiRequestBridgePreservesToolsAndThoughtSignature(t *testing.T) {
	body := []byte(`{
      "contents":[
        {"role":"model","parts":[{"text":"thinking","thought":true},{"functionCall":{"name":"lookup","args":{"q":"x"}},"thoughtSignature":"opaque"}]},
        {"role":"user","parts":[{"functionResponse":{"name":"lookup","response":{"value":1}}}]},
        {"role":"user","parts":[{"text":"continue"}]}
      ],
      "tools":[{"functionDeclarations":[{"name":"lookup","description":"find","parametersJsonSchema":{"type":"object"}}]}],
      "toolConfig":{"functionCallingConfig":{"mode":"ANY"}}
    }`)
	converted, err := prepareGeminiChatBridge(body, "preset-gemini-hybrid")
	if err != nil {
		t.Fatal(err)
	}
	text := string(converted)
	for _, want := range []string{`"model":"preset-gemini-hybrid"`, `"reasoning_content":"thinking"`, `"thought_signature":"opaque"`, `"tool_call_id":"gemini_call_0_1"`, `"tool_choice":"required"`} {
		if !strings.Contains(text, want) {
			t.Fatalf("missing %s in %s", want, text)
		}
	}
	if strings.Contains(text, `"tool_name"`) {
		t.Fatalf("nonstandard tool result field reached Chat bridge: %s", text)
	}
}

func TestGeminiRequestBridgeMatchesExplicitToolCallIDsOutOfOrder(t *testing.T) {
	converted, err := prepareGeminiChatBridge([]byte(`{
	  "contents":[
	    {"role":"model","parts":[
	      {"functionCall":{"id":"call_first","name":"lookup","args":{"q":"first"}}},
	      {"functionCall":{"id":"call_second","name":"lookup","args":{"q":"second"}}}
	    ]},
	    {"role":"user","parts":[
	      {"functionResponse":{"id":"call_second","name":"lookup","response":{"value":2}}},
	      {"functionResponse":{"id":"call_first","name":"lookup","response":{"value":1}}}
	    ]}
	  ]
	}`), "gemini-route")
	if err != nil {
		t.Fatal(err)
	}
	var body struct {
		Messages []map[string]any `json:"messages"`
	}
	if err := json.Unmarshal(converted, &body); err != nil {
		t.Fatal(err)
	}
	if len(body.Messages) != 3 {
		t.Fatalf("messages=%+v", body.Messages)
	}
	if body.Messages[1]["tool_call_id"] != "call_second" || body.Messages[2]["tool_call_id"] != "call_first" {
		t.Fatalf("explicit response order was not preserved: %+v", body.Messages)
	}
	for _, message := range body.Messages[1:] {
		if _, exists := message["tool_name"]; exists {
			t.Fatalf("nonstandard tool_name field survived: %+v", message)
		}
	}
}

func TestGeminiRequestBridgeKeepsIDLessFIFOCompatibility(t *testing.T) {
	converted, err := prepareGeminiChatBridge([]byte(`{
	  "contents":[
	    {"role":"model","parts":[
	      {"functionCall":{"name":"lookup","args":{"q":"first"}}},
	      {"functionCall":{"name":"lookup","args":{"q":"second"}}}
	    ]},
	    {"role":"user","parts":[
	      {"functionResponse":{"name":"lookup","response":{"value":1}}},
	      {"functionResponse":{"name":"lookup","response":{"value":2}}}
	    ]}
	  ]
	}`), "gemini-route")
	if err != nil {
		t.Fatal(err)
	}
	var body struct {
		Messages []map[string]any `json:"messages"`
	}
	if err := json.Unmarshal(converted, &body); err != nil {
		t.Fatal(err)
	}
	if len(body.Messages) != 3 || body.Messages[1]["tool_call_id"] != "gemini_call_0_0" || body.Messages[2]["tool_call_id"] != "gemini_call_0_1" {
		t.Fatalf("ID-less FIFO changed: %+v", body.Messages)
	}
}

func TestGeminiRequestBridgeRejectsInvalidToolCallIdentity(t *testing.T) {
	for _, tc := range []struct {
		name string
		body string
	}{
		{
			name: "duplicate id",
			body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}},{"functionCall":{"id":"call_1","name":"b","args":{}}}]}]}`,
		},
		{
			name: "mismatched name",
			body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}}]},{"role":"user","parts":[{"functionResponse":{"id":"call_1","name":"b","response":{}}}]}]}`,
		},
		{
			name: "missing response id for explicit call",
			body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"a","args":{}}}]},{"role":"user","parts":[{"functionResponse":{"name":"a","response":{}}}]}]}`,
		},
		{
			name: "invalid id",
			body: `{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call id","name":"a","args":{}}}]}]}`,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := prepareGeminiChatBridge([]byte(tc.body), "gemini-route"); err == nil {
				t.Fatal("invalid identity must fail")
			}
		})
	}
}

func TestGeminiRequestBridgeUsesProviderNativeThinkingEnvelope(t *testing.T) {
	converted, err := prepareGeminiChatBridge([]byte(`{
	  "contents":[{"role":"user","parts":[{"text":"hello"}]}],
	  "generationConfig":{"temperature":1,"topK":50,"topP":1,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":-1}}
	}`), "gemini-3.6-flash")
	if err != nil {
		t.Fatal(err)
	}
	var body map[string]any
	if json.Unmarshal(converted, &body) != nil {
		t.Fatal("converted body is invalid")
	}
	for _, forbidden := range []string{"temperature", "top_k", "top_p", "think", "include_reasoning", "thinking_token_budget"} {
		if _, exists := body[forbidden]; exists {
			t.Fatalf("unsupported field %q in converted body: %s", forbidden, converted)
		}
	}
	extra := body["extra_body"].(map[string]any)
	google := extra["google"].(map[string]any)
	thinking := google["thinking_config"].(map[string]any)
	if thinking["include_thoughts"] != true || thinking["thinking_budget"] != float64(-1) {
		t.Fatalf("thinking envelope mismatch: %+v", thinking)
	}
}

func TestGeminiRequestBridgeAcceptsOfficialAgyPlannerStructuredOutput(t *testing.T) {
	converted, err := prepareGeminiChatBridge([]byte(`{
	  "contents":[{"role":"user","parts":[{"text":"plan"}]}],
	  "generationConfig":{
	    "responseMimeType":"application/json",
	    "responseJsonSchema":{"type":"OBJECT","properties":{"steps":{"type":"ARRAY"}},"required":["steps"]}
	  },
	  "tools":[{"functionDeclarations":[{
	    "name":"read_file","parameters":{"type":"OBJECT","properties":{"path":{"type":"STRING"}}},
	    "responseJsonSchema":{"type":"object"}
	  }]}]
	}`), "gemini-3.6-flash")
	if err != nil {
		t.Fatal(err)
	}
	var body map[string]any
	if json.Unmarshal(converted, &body) != nil {
		t.Fatal("converted body is invalid")
	}
	format := body["response_format"].(map[string]any)
	jsonSchema := format["json_schema"].(map[string]any)
	schema := jsonSchema["schema"].(map[string]any)
	if format["type"] != "json_schema" || jsonSchema["name"] != "agy_response" || jsonSchema["strict"] != true || schema["type"] != "object" {
		t.Fatalf("structured output mismatch: %+v", format)
	}
	tools := body["tools"].([]any)
	function := tools[0].(map[string]any)["function"].(map[string]any)
	if function["name"] != "read_file" {
		t.Fatalf("function conversion mismatch: %+v", function)
	}
}

func TestGeminiRequestBridgeRejectsConflictingStructuredSchemas(t *testing.T) {
	for _, body := range []string{
		`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"generationConfig":{"responseMimeType":"application/json","responseSchema":{"type":"object"},"responseJsonSchema":{"type":"object"}}}`,
		`{"contents":[{"role":"user","parts":[{"text":"x"}]}],"tools":[{"functionDeclarations":[{"name":"f","parameters":{"type":"object"},"parametersJsonSchema":{"type":"object"}}]}]}`,
	} {
		if _, err := prepareGeminiChatBridge([]byte(body), "gemini-3.6-flash"); err == nil {
			t.Fatal("conflicting schema alternatives must fail")
		}
	}
}

func TestGeminiStreamBridgeEmitsBoundedToolOnce(t *testing.T) {
	w := httptest.NewRecorder()
	stream := newGeminiBridgeStream(w, "gemini-3.6-flash")
	input := "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_provider_1\",\"function\":{\"name\":\"lookup\",\"arguments\":\"{\\\"q\\\":\"}}]}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"x\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
		"data: [DONE]\n\n"
	if err := stream.Feed([]byte(input)); err != nil {
		t.Fatal(err)
	}
	if strings.Count(w.Body.String(), `"functionCall"`) != 1 || !strings.Contains(w.Body.String(), `"finishReason":"STOP"`) {
		t.Fatalf("tool terminal mismatch: %s", w.Body.String())
	}
	if !strings.Contains(w.Body.String(), `"id":"call_provider_1"`) {
		t.Fatalf("provider tool call id was not preserved: %s", w.Body.String())
	}
	usageWriter := httptest.NewRecorder()
	usageStream := newGeminiBridgeStream(usageWriter, "m")
	if err := usageStream.Feed([]byte("data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5,\"prompt_tokens_details\":{\"cached_tokens\":1},\"completion_tokens_details\":{\"reasoning_tokens\":1}}}\n\ndata: [DONE]\n\n")); err != nil {
		t.Fatal(err)
	}
	for _, want := range []string{`"promptTokenCount":3`, `"candidatesTokenCount":2`, `"cachedContentTokenCount":1`, `"thoughtsTokenCount":1`, `"totalTokenCount":5`} {
		if !strings.Contains(usageWriter.Body.String(), want) {
			t.Fatalf("missing usage %s in %s", want, usageWriter.Body.String())
		}
	}
	oversize := newGeminiBridgeStream(httptest.NewRecorder(), "m")
	state := &geminiBridgeToolState{name: "lookup"}
	state.arguments.WriteString(strings.Repeat("x", geminiToolArgumentLimit))
	oversize.tools[0] = state
	chunk := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"x"}}]}}]}` + "\n\n"
	if err := oversize.Feed([]byte(chunk)); err == nil {
		t.Fatal("oversize tool arguments must fail")
	}
}

func TestGeminiRejectionObservationIsClassificationOnly(t *testing.T) {
	core, observed := observer.New(zap.InfoLevel)
	secret := "secret-fixture-must-not-be-logged"

	authServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal"}, &fakeRunService{}, zap.New(core))
	authReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"`+secret+`"}]}]}`))
	authServer.routes().ServeHTTP(httptest.NewRecorder(), authReq)

	managedServer := managedCredentialMigrationServer(t, &providerFakeRunService{})
	managedServer.logger = zap.New(core)
	managedReq := httptest.NewRequest(http.MethodPost, "/gemini/public-route/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`))
	managedReq.Header.Set("X-Goog-Api-Key", "managed-iop-token")
	managedReq.Header.Set(legacyProviderCredentialHeader, secret)
	managedServer.routes().ServeHTTP(httptest.NewRecorder(), managedReq)

	candidate := anthropicTestCandidate(t, "openai")
	candidate.ActualModel = "served-chat"
	providerBody := []byte(`{"error":{"message":"` + secret + `"}}`)
	fake := &providerFakeRunService{
		poolDispatchPath:      string(edgeservice.ProviderPoolPathTunnel),
		poolSelectedCandidate: candidate,
		tunnelFrames:          anthropicTunnelFrames(http.StatusBadRequest, "application/json", providerBody),
	}
	providerServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal"}, fake, zap.New(core))
	providerServer.SetModelCatalog([]config.ModelCatalogEntry{{ID: "r", Providers: map[string]string{"chat": "served-chat"}}})
	providerReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`))
	providerReq.Header.Set("X-Goog-Api-Key", "principal")
	providerServer.routes().ServeHTTP(httptest.NewRecorder(), providerReq)

	localFailureServer := NewServer(config.EdgeOpenAIConf{BearerToken: "principal", Adapter: "ollama", Target: "provider-model"}, &fakeRunService{submitErrAfter: 1}, zap.New(core))
	localFailureReq := httptest.NewRequest(http.MethodPost, "/gemini/r/v1beta/models/m:streamGenerateContent?alt=sse", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"`+secret+`"}]}]}`))
	localFailureReq.Header.Set("X-Goog-Api-Key", "principal")
	localFailureServer.routes().ServeHTTP(httptest.NewRecorder(), localFailureReq)

	entries := observed.FilterMessage(geminiRejectionLogMessage).All()
	if len(entries) != 3 {
		t.Fatalf("rejection observations=%d logs=%+v", len(entries), observed.All())
	}
	wantClasses := []string{string(geminiRejectionPreIngress), string(geminiRejectionPreIngress), string(geminiRejectionProviderHTTP)}
	wantStatuses := []int64{http.StatusUnauthorized, http.StatusBadRequest, http.StatusBadRequest}
	for index, entry := range entries {
		context := entry.ContextMap()
		if len(context) != 4 || context["surface"] != "gemini" || context["bridge"] != "chat" || context["rejection_class"] != wantClasses[index] || context["http_status"] != wantStatuses[index] {
			t.Fatalf("unsafe or incomplete observation: %+v", context)
		}
		if strings.Contains(fmt.Sprint(context), secret) {
			t.Fatalf("secret leaked to observation: %+v", context)
		}
	}
}