package router

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

func TestAReplySaysWhatHardwareAnsweredIt(t *testing.T) {
	// qwen2.5-coder-14b had one replica on cuda01's card and one on a CPU
	// node, and the chat page could not tell which had answered: the same
	// model, ten times faster or slower, for no visible reason.
	for pool, want := range map[string]string{
		"none/6000-24g": "6000-24g",
		"none/cpu":      "cpu",
		"":              "",
	} {
		e := newFakeEngine(t, "m")
		rt := newTestRouter(t, baseConfig())
		rt.SetEndpoints(t.Context(), []endpoint.Info{
			{ID: "r", URL: e.srv.URL, Models: []string{"m"}, Pool: pool, Weight: 1},
		})
		waitForEndpoint(t, rt)

		rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
		if rec.Code != http.StatusOK {
			t.Fatalf("%q: status %d: %s", pool, rec.Code, rec.Body.String())
		}
		if got := rec.Header().Get(hardwareHeader); got != want {
			t.Errorf("pool %q: hardware %q, want %q", pool, got, want)
		}
	}
}

func TestThePickerSaysWhereEachModelRuns(t *testing.T) {
	gpu, cpu := newFakeEngine(t, "qwen2.5-coder-7b-instruct"), newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, user := chatRouter(t, gpu.srv.URL)
	rt.SetEndpoints(t.Context(), []endpoint.Info{
		{ID: "g", URL: gpu.srv.URL, Models: []string{"qwen2.5-coder-7b-instruct"}, Pool: "none/6000-24g", Weight: 1},
		{ID: "c", URL: cpu.srv.URL, Models: []string{"coder"}, Pool: "none/cpu", Weight: 1},
	})
	waitForEndpoint(t, rt)

	got := rt.modelsFor(httptest.NewRequest("GET", "/chat", nil), user)
	if len(got) != 1 || got[0].Runs() != "GPU + CPU" {
		t.Fatalf("picker %+v; want one model running on GPU + CPU, its alias's replica included", got)
	}
}