package openai

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

	"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)
		}
	}
}

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 TestGeminiStreamBridgeEmitsBoundedToolOnce(t *testing.T) {
	w := httptest.NewRecorder()
	stream := newGeminiBridgeStream(w, "gemini-3.6-flash")
	input := "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"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())
	}
	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")
	}
}