package openai

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

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

type fakeRunService struct {
	req        edgeservice.SubmitRunRequest
	ollamaReq  edgeservice.OllamaAPIRequest
	ollamaResp edgeservice.OllamaAPIView
	events     chan *iop.RunEvent
}

func (s *fakeRunService) SubmitRun(_ context.Context, req edgeservice.SubmitRunRequest) (*edgeservice.RunHandle, error) {
	s.req = req
	return &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{
			RunID:      "run-test",
			Target:     req.Target,
			TimeoutSec: 5,
		},
		RunStream: edgeservice.RunStream{
			Events:     s.events,
			NodeEvents: make(chan *iop.EdgeNodeEvent),
		},
	}, nil
}

func (s *fakeRunService) OllamaAPI(_ context.Context, req edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error) {
	s.ollamaReq = req
	if s.ollamaResp.StatusCode == 0 {
		s.ollamaResp.StatusCode = http.StatusOK
	}
	return s.ollamaResp, nil
}

func TestChatCompletionsDispatchesConfiguredOllamaTarget(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "hello"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: " world"}
	fake.events <- &iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 2, OutputTokens: 2}}

	srv := NewServer(config.EdgeOpenAIConf{
		Adapter:    "ollama",
		Target:     "llama-fixed",
		SessionID:  "cline",
		TimeoutSec: 15,
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"client-model",
		"messages":[{"role":"system","content":"brief"},{"role":"user","content":"say 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 fake.req.Adapter != "ollama" || fake.req.Target != "llama-fixed" {
		t.Fatalf("dispatch target mismatch: %+v", fake.req)
	}
	if fake.req.SessionID != "cline" || fake.req.TimeoutSec != 15 {
		t.Fatalf("execution config mismatch: %+v", fake.req)
	}
	if !strings.Contains(fake.req.Prompt, "system: brief") || !strings.Contains(fake.req.Prompt, "user: say hello") {
		t.Fatalf("prompt did not include messages: %q", fake.req.Prompt)
	}

	var resp chatCompletionResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if resp.Choices[0].Message.Content != "hello world" {
		t.Fatalf("content: got %q", resp.Choices[0].Message.Content)
	}
	if resp.Usage == nil || resp.Usage.TotalTokens != 4 {
		t.Fatalf("usage: %+v", resp.Usage)
	}
}

func TestChatCompletionsUsesRequestModelWhenNoConfiguredTarget(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}]
	}`))
	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 fake.req.Target != "from-request" {
		t.Fatalf("target: got %q", fake.req.Target)
	}
}

func TestChatCompletionsPassesOllamaOptions(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"from-request",
		"messages":[{"role":"user","content":"hi"}],
		"options":{"temperature":0.2,"top_p":0.9,"num_predict":32,"stop":["END"]},
		"keep_alive":"10m",
		"think":false,
		"format":"json",
		"tools":[{"type":"function","function":{"name":"lookup"}}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	options, ok := fake.req.Input["options"].(map[string]any)
	if !ok {
		t.Fatalf("options not passed: %+v", fake.req.Input)
	}
	if options["temperature"].(float64) != 0.2 || options["top_p"].(float64) != 0.9 || options["num_predict"].(float64) != 32 {
		t.Fatalf("unexpected options: %+v", options)
	}
	if fake.req.Input["keep_alive"] != "10m" || fake.req.Input["think"] != false || fake.req.Input["format"] != "json" {
		t.Fatalf("top-level ollama fields not passed: %+v", fake.req.Input)
	}
	if tools, ok := fake.req.Input["tools"].([]any); !ok || len(tools) != 1 {
		t.Fatalf("tools not passed: %+v", fake.req.Input["tools"])
	}
}

func TestChatCompletionsStreamsSSE(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "hi"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"stream-model",
		"stream":true,
		"messages":[{"role":"user","content":"hi"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	body := w.Body.String()
	if !strings.Contains(body, `"content":"hi"`) || !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("unexpected SSE body:\n%s", body)
	}
}

func TestChatCompletionsReturnsReasoningContentSeparately(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 5)}
	fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "thinking "}
	fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "more"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "answer"}
	fake.events <- &iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 2, OutputTokens: 2}}

	srv := NewServer(config.EdgeOpenAIConf{
		Adapter: "ollama",
		Target:  "llama-fixed",
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"client-model",
		"messages":[{"role":"user","content":"hi"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

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

	var resp chatCompletionResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if resp.Choices[0].Message.Content != "answer" {
		t.Fatalf("content: got %q, expected \"answer\"", resp.Choices[0].Message.Content)
	}
	if resp.Choices[0].Message.ReasoningContent != "thinking more" {
		t.Fatalf("reasoning_content: got %q, expected \"thinking more\"", resp.Choices[0].Message.ReasoningContent)
	}
}

func TestChatCompletionsStreamsReasoningSSE(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 4)}
	fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "think"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "hi"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"stream-model",
		"stream":true,
		"messages":[{"role":"user","content":"hi"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	body := w.Body.String()
	if !strings.Contains(body, `"reasoning_content":"think"`) || !strings.Contains(body, `"content":"hi"`) || !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("unexpected SSE body:\n%s", body)
	}
}

func TestChatCompletionsStrictOutputNormalizesXMLStyleAgentResponse(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "---\n<thinking>hidden</thinking>\n\n<attempt_completion>\n<result>ok</result>\n</attempt completion>\n```\ntrailing noise"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", StrictOutput: true}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"agent-model",
		"messages":[{"role":"user","content":"finish"}]
	}`))
	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 fake.req.Input["think"] != false {
		t.Fatalf("strict output should disable thinking by default: %+v", fake.req.Input)
	}
	var resp chatCompletionResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	want := "<attempt_completion>\n<result>ok</result>\n</attempt_completion>"
	if resp.Choices[0].Message.Content != want {
		t.Fatalf("content:\ngot  %q\nwant %q", resp.Choices[0].Message.Content, want)
	}
	if resp.Choices[0].Message.ReasoningContent != "" {
		t.Fatalf("strict output should drop reasoning content: %q", resp.Choices[0].Message.ReasoningContent)
	}
}

func TestChatCompletionsStrictOutputWrapsPlainTextWhenPromptDeclaresCompletionTool(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "plain answer"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", StrictOutput: true}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"agent-model",
		"messages":[
			{"role":"system","content":"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>"},
			{"role":"user","content":"finish"}
		]
	}`))
	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 !strings.Contains(fake.req.Prompt, "output exactly one <attempt_completion> block") || !strings.Contains(fake.req.Prompt, "actual user-facing response") {
		t.Fatalf("strict contract instruction not injected:\n%s", fake.req.Prompt)
	}
	var resp chatCompletionResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	want := "<attempt_completion>\n<result>plain answer</result>\n</attempt_completion>"
	if resp.Choices[0].Message.Content != want {
		t.Fatalf("content:\ngot  %q\nwant %q", resp.Choices[0].Message.Content, want)
	}
}

func TestChatCompletionsStrictOutputBuffersStreamingResponse(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 4)}
	fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "hidden"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "---\n<think>noise</think><read_file><path>a.go</path></read file>"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", StrictOutput: true, StrictStreamBuffer: true}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"agent-model",
		"stream":true,
		"messages":[{"role":"user","content":"read"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	body := w.Body.String()
	if strings.Contains(body, "reasoning_content") || strings.Contains(body, "<think>") || strings.Contains(body, "</read file>") {
		t.Fatalf("strict stream leaked unstable content:\n%s", body)
	}
	if !strings.Contains(body, "read_file") || !strings.Contains(body, "a.go") || !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("unexpected strict SSE body:\n%s", body)
	}
}

func TestChatCompletionsStrictOutputStreamsLiveContent(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 5)}
	fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: "hidden"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "<attempt_completion>\n"}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "<result>ok</result>\n</attempt_completion>"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", StrictOutput: true}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"agent-model",
		"stream":true,
		"messages":[{"role":"user","content":"finish"}]
	}`))
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	body := w.Body.String()
	if strings.Contains(body, "reasoning_content") || strings.Contains(body, "hidden") {
		t.Fatalf("strict stream leaked reasoning content:\n%s", body)
	}
	if strings.Count(body, `"content":`) != 2 || !strings.Contains(body, "attempt_completion") || !strings.Contains(body, "data: [DONE]") {
		t.Fatalf("strict stream should pass live content deltas:\n%s", body)
	}
}

func TestModelsUsesConfiguredModelsOrTarget(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{Target: "fallback-model"}, &fakeRunService{}, nil)
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	w := httptest.NewRecorder()
	srv.handleModels(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status: got %d", w.Code)
	}
	if !strings.Contains(w.Body.String(), "fallback-model") {
		t.Fatalf("expected target model, got %s", w.Body.String())
	}

	srv = NewServer(config.EdgeOpenAIConf{Models: []string{"a", "b"}}, &fakeRunService{}, nil)
	w = httptest.NewRecorder()
	srv.handleModels(w, req)
	if !strings.Contains(w.Body.String(), `"id":"a"`) || !strings.Contains(w.Body.String(), `"id":"b"`) {
		t.Fatalf("expected configured models, got %s", w.Body.String())
	}
}

func TestOllamaAPIPassthrough(t *testing.T) {
	fake := &fakeRunService{
		ollamaResp: edgeservice.OllamaAPIView{
			StatusCode:  http.StatusAccepted,
			ContentType: "application/json",
			Body:        `{"ok":true}`,
		},
	}
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", TimeoutSec: 15}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/api/show", strings.NewReader(`{"model":"gemma4:26b"}`))
	w := httptest.NewRecorder()

	srv.handleOllamaAPI(w, req)

	if w.Code != http.StatusAccepted {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	if fake.ollamaReq.Adapter != "ollama" || fake.ollamaReq.Method != http.MethodPost || fake.ollamaReq.Path != "/api/show" {
		t.Fatalf("passthrough req mismatch: %+v", fake.ollamaReq)
	}
	if fake.ollamaReq.Body != `{"model":"gemma4:26b"}` || fake.ollamaReq.TimeoutSec != 15 {
		t.Fatalf("passthrough body/timeout mismatch: %+v", fake.ollamaReq)
	}
	if w.Body.String() != `{"ok":true}` {
		t.Fatalf("body: got %s", w.Body.String())
	}
}

func TestCollectRunResultTimesOut(t *testing.T) {
	handle := &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{TimeoutSec: 1},
		RunStream: edgeservice.RunStream{
			Events:     make(chan *iop.RunEvent),
			NodeEvents: make(chan *iop.EdgeNodeEvent),
		},
	}
	ctx, cancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
	defer cancel()
	_, _, _, err := collectRunResult(ctx, handle)
	if err == nil {
		t.Fatal("expected timeout error")
	}
}