package openai

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

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

// fakeRunResultWithTimeout implements edgeservice.RunResult with a directly
// configurable WaitTimeout, so tests can exercise the run-timeout cancel path
// without waiting out the real DefaultTimeoutSec-derived duration.
type fakeRunResultWithTimeout struct {
	dispatch    edgeservice.RunDispatch
	stream      edgeservice.RunStream
	waitTimeout time.Duration
}

func (f *fakeRunResultWithTimeout) Dispatch() edgeservice.RunDispatch { return f.dispatch }
func (f *fakeRunResultWithTimeout) Stream() edgeservice.RunStream     { return f.stream }
func (f *fakeRunResultWithTimeout) Close()                            {}
func (f *fakeRunResultWithTimeout) WaitTimeout() time.Duration        { return f.waitTimeout }

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.Stream(), handle.WaitTimeout())
	if err == nil {
		t.Fatal("expected timeout error")
	}
}
func TestChatCompletionContextCancelSendsCancelRun(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent)}
	srv := NewServer(config.EdgeOpenAIConf{
		Adapter:    "ollama",
		Target:     "llama-fixed",
		SessionID:  "cline",
		TimeoutSec: 15,
	}, fake, nil)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"client-model",
		"messages":[{"role":"user","content":"hi"}]
	}`)).WithContext(ctx)
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	if w.Code != http.StatusRequestTimeout {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	calls := fake.cancelCallsSnapshot()
	if len(calls) != 1 {
		t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
	}
	if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
		t.Fatalf("unexpected cancel request: %+v", calls[0])
	}
}
func TestResponsesContextCancelSendsCancelRun(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent)}
	srv := NewServer(config.EdgeOpenAIConf{
		Adapter:    "ollama",
		Target:     "llama-fixed",
		SessionID:  "cline",
		TimeoutSec: 15,
	}, fake, nil)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"client-model",
		"input":"say hello"
	}`)).WithContext(ctx)
	w := httptest.NewRecorder()

	srv.handleResponses(w, req)

	if w.Code != http.StatusRequestTimeout {
		t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
	}
	calls := fake.cancelCallsSnapshot()
	if len(calls) != 1 {
		t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
	}
	if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
		t.Fatalf("unexpected cancel request: %+v", calls[0])
	}
}
func TestStreamChatCompletionContextCancelSendsCancelRun(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent)}
	srv := NewServer(config.EdgeOpenAIConf{
		Adapter:    "ollama",
		Target:     "llama-fixed",
		SessionID:  "cline",
		TimeoutSec: 15,
	}, fake, nil)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"client-model",
		"stream":true,
		"messages":[{"role":"user","content":"hi"}]
	}`)).WithContext(ctx)
	w := httptest.NewRecorder()

	srv.handleChatCompletions(w, req)

	calls := fake.cancelCallsSnapshot()
	if len(calls) != 1 {
		t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
	}
	if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
		t.Fatalf("unexpected cancel request: %+v", calls[0])
	}
}
func TestStreamChatCompletionTimeoutSendsCancelRun(t *testing.T) {
	fake := &fakeRunService{}
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)

	handle := &fakeRunResultWithTimeout{
		dispatch: edgeservice.RunDispatch{
			RunID:     "run-timeout",
			NodeID:    "node-1",
			Adapter:   "ollama",
			Target:    "llama3",
			SessionID: "cline",
		},
		stream: edgeservice.RunStream{
			Events:     make(chan *iop.RunEvent),
			NodeEvents: make(chan *iop.EdgeNodeEvent),
		},
		waitTimeout: 50 * time.Millisecond,
	}

	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	w := httptest.NewRecorder()

	srv.streamChatCompletion(w, newTestChatDispatchContext(req, chatCompletionRequest{Model: "llama3"}), handle)

	calls := fake.cancelCallsSnapshot()
	if len(calls) != 1 {
		t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
	}
	if calls[0].RunID != "run-timeout" || calls[0].NodeRef != "node-1" {
		t.Fatalf("unexpected cancel request: %+v", calls[0])
	}
	if !strings.Contains(w.Body.String(), "run timed out") {
		t.Fatalf("expected timeout SSE error, got %s", w.Body.String())
	}
}
func TestChatCompletionCompleteDoesNotSendCancelRun(t *testing.T) {
	fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
	fake.events <- &iop.RunEvent{Type: "delta", Delta: "hello"}
	fake.events <- &iop.RunEvent{Type: "complete"}

	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())
	}
	if calls := fake.cancelCallsSnapshot(); len(calls) != 0 {
		t.Fatalf("expected no CancelRun calls for a normal completion, got %d: %+v", len(calls), calls)
	}
}
func TestModelsExposesCatalogRouteModels(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "llama3"},
			{Model: "model-b", Adapter: "vllm", Target: "qwen"},
		},
	}, &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)
	}
	body := w.Body.String()
	if !strings.Contains(body, `"id":"model-a"`) || !strings.Contains(body, `"id":"model-b"`) {
		t.Fatalf("expected catalog model IDs, got %s", body)
	}
}
func TestModelsSkipsRoutesWithEmptyTarget(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-with-target", Adapter: "ollama", Target: "llama3"},
			{Model: "model-no-target", Adapter: "ollama", Target: ""},
		},
	}, &fakeRunService{}, nil)
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	w := httptest.NewRecorder()
	srv.handleModels(w, req)
	body := w.Body.String()
	if !strings.Contains(body, `"id":"model-with-target"`) {
		t.Fatalf("expected model-with-target in response, got %s", body)
	}
	if strings.Contains(body, "model-no-target") {
		t.Fatalf("model-no-target should be skipped (no target), got %s", body)
	}
}
func TestModelsRouteCatalogWinsOverModelsField(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{
		Models: []string{"legacy-model"},
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "catalog-model", Adapter: "ollama", Target: "llama3"},
		},
	}, &fakeRunService{}, nil)
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	w := httptest.NewRecorder()
	srv.handleModels(w, req)
	body := w.Body.String()
	if !strings.Contains(body, `"id":"catalog-model"`) {
		t.Fatalf("expected catalog model in response, got %s", body)
	}
	if strings.Contains(body, "legacy-model") {
		t.Fatalf("legacy model should be hidden when catalog is set, got %s", body)
	}
}
func TestChatCompletionsRouteCatalogDispatches(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{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "llama3"},
			{Model: "model-b", Adapter: "vllm", Target: "qwen"},
		},
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"model-a",
		"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.Adapter != "ollama" || fake.req.Target != "llama3" {
		t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
	}
}
func TestChatCompletionsRouteCatalogDispatchesModelB(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{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "llama3"},
			{Model: "model-b", Adapter: "ollama", Target: "qwen", MaxQueue: 5, QueueTimeoutMS: 2000},
		},
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"metadata":{},
		"model":"model-b",
		"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.Adapter != "ollama" || fake.req.Target != "qwen" {
		t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
	}
	if fake.req.ModelGroupKey != "model-b" {
		t.Fatalf("model group key: got %q, want model-b", fake.req.ModelGroupKey)
	}
	if fake.req.MaxQueue != 5 || fake.req.QueueTimeoutMS != 2000 {
		t.Fatalf("queue policy mismatch: MaxQueue=%d, QueueTimeoutMS=%d", fake.req.MaxQueue, fake.req.QueueTimeoutMS)
	}
}
func TestChatCompletionsCatalogMissFallsToLegacyTarget(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",
		Target:  "legacy-target",
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "llama3"},
		},
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"unknown-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())
	}
	if fake.req.Target != "legacy-target" {
		t.Fatalf("expected legacy-target fallback, got %q", fake.req.Target)
	}
}
func TestChatCompletionsEmptyModelReturns400(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, &fakeRunService{}, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
		"model":"",
		"messages":[{"role":"user","content":"hi"}]
	}`))
	w := httptest.NewRecorder()
	srv.handleChatCompletions(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 for empty model, got %d body=%s", w.Code, w.Body.String())
	}
}
func TestResponsesRouteCatalogDispatchesRoute(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{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "qwen"},
		},
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"model-a",
		"input":"say hello"
	}`))
	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.Adapter != "ollama" || fake.req.Target != "qwen" {
		t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
	}
	if fake.req.ModelGroupKey != "model-a" {
		t.Fatalf("model group key: got %q, want model-a", fake.req.ModelGroupKey)
	}
}
func TestResponsesRouteCatalogDispatchesQueuePolicy(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{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "model-a", Adapter: "ollama", Target: "route-target", MaxQueue: 3, QueueTimeoutMS: 1500},
		},
	}, fake, nil)
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
		"model":"model-a",
		"input":"test"
	}`))
	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.Target != "route-target" {
		t.Fatalf("expected route-target, got %q", fake.req.Target)
	}
	if fake.req.ModelGroupKey != "model-a" {
		t.Fatalf("model group key: got %q, want model-a", fake.req.ModelGroupKey)
	}
	if fake.req.MaxQueue != 3 || fake.req.QueueTimeoutMS != 1500 {
		t.Fatalf("queue policy mismatch: MaxQueue=%d, QueueTimeoutMS=%d", fake.req.MaxQueue, fake.req.QueueTimeoutMS)
	}
}

func TestResolveRouteDispatchPreservesWorkspaceRequired(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "codex", Adapter: "cli", Target: "codex", WorkspaceRequired: true},
			{Model: "llama3", Adapter: "ollama", Target: "llama3:8b"},
		},
	}, &fakeRunService{}, nil)

	dispatch, ok := srv.resolveRouteDispatch("codex")
	if !ok {
		t.Fatal("expected dispatch to succeed for codex route")
	}
	if !dispatch.WorkspaceRequired {
		t.Fatalf("expected workspace_required=true for codex route, got false")
	}

	dispatch, ok = srv.resolveRouteDispatch("llama3")
	if !ok {
		t.Fatal("expected dispatch to succeed for llama3 route")
	}
	if dispatch.WorkspaceRequired {
		t.Fatalf("expected workspace_required=false for llama3 route, got true")
	}
}
func TestResolveRouteDispatchFallbackWorkspaceRequiredFalse(t *testing.T) {
	srv := NewServer(config.EdgeOpenAIConf{
		Adapter: "ollama",
		Target:  "llama3",
		ModelRoutes: []config.OpenAIRouteEntry{
			{Model: "codex", Adapter: "cli", Target: "codex", WorkspaceRequired: true},
		},
	}, &fakeRunService{}, nil)

	dispatch, ok := srv.resolveRouteDispatch("unknown-model")
	if !ok {
		t.Fatal("expected fallback dispatch to succeed")
	}
	if dispatch.WorkspaceRequired {
		t.Fatalf("expected workspace_required=false for fallback route, got true")
	}
}
func TestCollectRunResultFailsWhenEventStreamCloses(t *testing.T) {
	events := make(chan *iop.RunEvent)
	close(events)
	handle := &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{TimeoutSec: 60},
		RunStream: edgeservice.RunStream{
			Events:     events,
			NodeEvents: make(chan *iop.EdgeNodeEvent),
		},
	}

	_, _, _, _, _, _, err := collectRunResult(context.Background(), handle.Stream(), handle.WaitTimeout())
	if err == nil {
		t.Fatal("expected closed stream error")
	}
	if !strings.Contains(err.Error(), "run stream closed") {
		t.Fatalf("expected run stream closed error, got %v", err)
	}
}