package autogen

import (
	"fmt"
	"math"
	"strings"
	"testing"
)

func TestResolveBackend_AutoPickAndExplicit(t *testing.T) {
	s := Settings{Backends: []BackendEntry{
		{ID: "llama-vk", Kind: "llama", Path: "C:/llama/llama-server.exe"},
		{ID: "vllm-1", Kind: "vllm", Path: "vllm", Default: true},
		{ID: "sd-1", Kind: "sd", Path: "sd-server"},
	}}

	// Auto-pick for the LLM class honours the ★ default (vllm here, not the
	// first-listed llama).
	if be := resolveBackend(s, nil, "llm"); be.ID != "vllm-1" {
		t.Errorf("auto-pick llm = %q, want vllm-1 (the class default)", be.ID)
	}
	// Explicit Override.Backend wins over the default.
	if be := resolveBackend(s, &Override{Backend: "llama-vk"}, "llm"); be.ID != "llama-vk" {
		t.Errorf("explicit backend = %q, want llama-vk", be.ID)
	}
	// A stale id falls back to auto-pick, not a broken zero value.
	if be := resolveBackend(s, &Override{Backend: "gone"}, "llm"); be.ID != "vllm-1" {
		t.Errorf("stale id = %q, want auto-pick vllm-1", be.ID)
	}
	// No entry of a class => zero value (caller uses the legacy exe).
	if be := resolveBackend(s, nil, "tts"); be.ID != "" {
		t.Errorf("no tts backend => %q, want empty", be.ID)
	}
	// First-of-class when none is marked default.
	s2 := Settings{Backends: []BackendEntry{
		{ID: "a", Kind: "llama", Path: "a"},
		{ID: "b", Kind: "llama", Path: "b"},
	}}
	if be := resolveBackend(s2, nil, "llm"); be.ID != "a" {
		t.Errorf("first-of-class = %q, want a", be.ID)
	}
}

func TestVllmCmdLines_Shape(t *testing.T) {
	s := Settings{Backends: []BackendEntry{{ID: "v", Kind: "vllm", Path: "C:/py/vllm.exe", Default: true}}}
	row := GgufRow{FullPath: `E:\Models\qwen\q.gguf`, SizeGB: 4}
	meta := Metadata{Architecture: "qwen3", ContextLength: 32768}
	be := resolveBackend(s, &Override{Ctx: 8192, VllmGpuUtil: 0.8, VllmTensorParallel: 2}, "llm")

	cmd := strings.Join(vllmCmdLines(s, row, &Override{Ctx: 8192, VllmGpuUtil: 0.8, VllmTensorParallel: 2}, "qwen3-q4", be, meta), " ")
	for _, want := range []string{
		"C:/py/vllm.exe", "serve E:/Models/qwen/q.gguf", "--port ${PORT}",
		"--served-model-name qwen3-q4", "--quantization gguf",
		"--gpu-memory-utilization 0.8", "--max-model-len 8192", "--tensor-parallel-size 2",
	} {
		if !strings.Contains(cmd, want) {
			t.Errorf("vllm cmd missing %q\ngot: %s", want, cmd)
		}
	}
	// A tokenizer is only emitted when set — never guessed from the model's folder.
	if strings.Contains(cmd, "--tokenizer") {
		t.Errorf("tokenizer emitted without an override: %s", cmd)
	}
	tok := strings.Join(vllmCmdLines(s, row, &Override{VllmTokenizer: "Qwen/Qwen3-8B"}, "m", be, meta), " ")
	if !strings.Contains(tok, "--tokenizer Qwen/Qwen3-8B") {
		t.Errorf("tokenizer override not emitted: %s", tok)
	}

	// With no GPU reading there is nothing to take a fraction of, so utilization
	// falls back to the flat default rather than to a number derived from half
	// the picture.
	withCard(t, 0, false)
	s.TargetVramGB = 24
	def := strings.Join(vllmCmdLines(s, row, nil, "m", be, meta), " ")
	if !strings.Contains(def, "--gpu-memory-utilization 0.9") {
		t.Errorf("no-GPU util should be the flat default: %s", def)
	}
}

// withCard stubs the one-shot card probe for the duration of a test, so sizing
// is exercised on machines with no GPU telemetry (CI) and with a card size we
// choose rather than whatever the host happens to have.
func withCard(t *testing.T, totalGB float64, ok bool) {
	t.Helper()
	prev := cachedTotalVramGB
	cachedTotalVramGB = func() (float64, bool) { return totalGB, ok }
	t.Cleanup(func() { cachedTotalVramGB = prev })
}

func TestVllmGpuUtil_DerivedFromBudget(t *testing.T) {
	withCard(t, 24, true)

	// The budget is the fraction of the card we hand vllm — not a flat 0.90 that
	// ignores the operator asking for a smaller share.
	if got, _ := vllmGpuUtil(Settings{TargetVramGB: 12}, nil, 12); round2(got) != 0.5 {
		t.Errorf("12GB of a 24GB card = %g, want 0.5", got)
	}
	// A per-model cap wins over the fleet budget, same precedence as llama sizing.
	if got, _ := vllmGpuUtil(Settings{TargetVramGB: 12}, &Override{VramTargetGB: 6}, 6); round2(got) != 0.25 {
		t.Errorf("per-model 6GB cap = %g, want 0.25", got)
	}
	// Clamped at both ends: a tiny budget still needs room for weights, and a
	// budget at/over the card size must leave the desktop something.
	if got, _ := vllmGpuUtil(Settings{TargetVramGB: 0.5}, nil, 0.5); got != vllmMinGpuUtil {
		t.Errorf("tiny budget = %g, want the floor %g", got, vllmMinGpuUtil)
	}
	if got, _ := vllmGpuUtil(Settings{TargetVramGB: 32}, nil, 32); got != vllmMaxGpuUtil {
		t.Errorf("oversized budget = %g, want the ceiling %g", got, vllmMaxGpuUtil)
	}
	// An explicit knob is passed through untouched, clamps included.
	if got, _ := vllmGpuUtil(Settings{TargetVramGB: 12}, &Override{VllmGpuUtil: 0.99}, 12); got != 0.99 {
		t.Errorf("pinned util = %g, want 0.99", got)
	}
}

func TestVllmMaxModelLen_SizedToBudget(t *testing.T) {
	// A model whose trained window is far larger than the budget can hold: vllm
	// allocates its KV pool from --max-model-len up front, so emitting the
	// trained length is a refused startup, not a large context.
	meta := Metadata{
		Architecture: "qwen3", ContextLength: 262144, BlockCount: 48,
		HeadCount: 32, HeadCountKv: 8, EmbeddingLength: 4096,
		KeyLength: 128, ValueLength: 128,
	}
	row := GgufRow{FullPath: "/m/q.gguf", SizeGB: 8}
	s := Settings{TargetVramGB: 16}

	ctx, _ := vllmMaxModelLen(s, nil, row, meta)
	if ctx <= 0 || ctx >= int(meta.ContextLength) {
		t.Fatalf("ctx = %d, want a budget-sized window below the trained %d", ctx, meta.ContextLength)
	}
	if ctx%4096 != 0 {
		t.Errorf("ctx = %d, want a 4096 multiple", ctx)
	}
	// The same model on a bigger budget must not shrink.
	big, _ := vllmMaxModelLen(Settings{TargetVramGB: 48}, nil, row, meta)
	if big < ctx {
		t.Errorf("bigger budget gave a smaller window: %d < %d", big, ctx)
	}
	// Never exceed what the model was trained for, however much VRAM there is.
	if huge, _ := vllmMaxModelLen(Settings{TargetVramGB: 4096}, nil, row, meta); huge != int(meta.ContextLength) {
		t.Errorf("huge budget = %d, want the trained %d", huge, meta.ContextLength)
	}
	// A pinned ctx is the user stating a requirement; sizing must not override it.
	if pinned, _ := vllmMaxModelLen(s, &Override{Ctx: 200000}, row, meta); pinned != 200000 {
		t.Errorf("pinned ctx = %d, want 200000", pinned)
	}
	// Weights alone over budget: vllm has no CPU offload, so say so rather than
	// emitting a window that implies the model fits.
	tiny, note := vllmMaxModelLen(Settings{TargetVramGB: 4}, nil, row, meta)
	if tiny != 4096 || !strings.Contains(note, "exceed") {
		t.Errorf("over-budget model = (%d, %q), want (4096, an exceeds-budget note)", tiny, note)
	}
}

// A small model whose whole trained window fits with room to spare must be
// granted only what it needs. vllm preallocates its full share, so handing it
// the budget makes a 0.5B model hold 22.8GB and the router evict everything
// else to admit it.
func TestVllmFootprint_SmallModelTakesOnlyWhatItNeeds(t *testing.T) {
	withCard(t, 24, true)
	// Qwen2.5-0.5B-Instruct as read from its config.json.
	meta := Metadata{
		Architecture: "qwen2", ContextLength: 32768, BlockCount: 24,
		HeadCount: 14, HeadCountKv: 2, EmbeddingLength: 896,
		KeyLength: 64, ValueLength: 64,
	}
	row := GgufRow{FullPath: "/m/qwen2.5-0.5b", IsHF: true, SizeGB: 0.92}
	s := Settings{TargetVramGB: 22.8, Backends: []BackendEntry{{ID: "v", Kind: "vllm", Path: "vllm", Default: true}}}

	ctx, _ := vllmMaxModelLen(s, nil, row, meta)
	if ctx != 32768 {
		t.Fatalf("ctx = %d, want the trained 32768", ctx)
	}
	m := GetKvCostModel(meta, "f16", "f16")
	kv := KvReserveGB(ctx, m.SlopeGB, m.ConstGB)*(1+vllmKvMarginFrac) + vllmKvMarginGB
	need := row.SizeGB + kv + vllmOverheadGB
	foot := vllmFootprintGB(s, nil, row, meta, ctx)
	if foot != need || foot > 4 {
		t.Fatalf("footprint = %.2fGB, want weights+KV+overhead = %.2fGB (a few GB, not the budget)", foot, need)
	}
	util, _ := vllmGpuUtil(s, nil, foot)
	// Rounded UP: printed at 2 decimals, the share must still cover the need.
	if round2(util)*24 < need {
		t.Errorf("util %g of 24GB = %.2fGB, under the %.2fGB needed", util, round2(util)*24, need)
	}

	var b strings.Builder
	var emitted []string
	emitVllmModel(&b, s, row, nil, "q", resolveBackend(s, nil, "llm"), meta, &emitted)
	out := b.String()
	if !strings.Contains(out, fmt.Sprintf("--gpu-memory-utilization %g", round2(util))) {
		t.Errorf("emitted util != %g:\n%s", round2(util), out)
	}
	if !strings.Contains(out, fmt.Sprintf("estVramGB: %g", round2(util*24))) {
		t.Errorf("estVramGB must charge the granted share %g:\n%s", round2(util*24), out)
	}
	if want := fmt.Sprintf("--kv-cache-memory-bytes %d", int64(math.Ceil(kv*gib))); !strings.Contains(out, want) {
		t.Errorf("KV pool not pinned, want %q:\n%s", want, out)
	}

	// A model whose window was sized TO the budget is charged the budget plus the
	// KV margin: the pinned pool is held whatever the window was sized against.
	big := Metadata{
		Architecture: "qwen3", ContextLength: 262144, BlockCount: 48,
		HeadCount: 32, HeadCountKv: 8, EmbeddingLength: 4096,
		KeyLength: 128, ValueLength: 128,
	}
	bigRow := GgufRow{FullPath: "/m/q.gguf", SizeGB: 8}
	bs := Settings{TargetVramGB: 16}
	bctx, _ := vllmMaxModelLen(bs, nil, bigRow, big)
	if got := vllmFootprintGB(bs, nil, bigRow, big, bctx); got < 15 || got > 17.5 {
		t.Errorf("budget-sized model footprint = %.2fGB, want the 16GB budget plus the KV margin", got)
	}
}

// Issue #93: vllm sized KV as utilization x card minus its own measured
// overhead, so a low overhead guess left MiniCPM without room for one sequence.
// The pool is pinned instead, and the two ways a user already owns the memory
// (a utilization pin, or the flags in extra args) turn the pin off.
func TestVllmKvCachePin(t *testing.T) {
	withCard(t, 24, true)
	meta := Metadata{
		Architecture: "qwen3", ContextLength: 32768, BlockCount: 36,
		HeadCount: 32, HeadCountKv: 8, EmbeddingLength: 4096,
		KeyLength: 128, ValueLength: 128,
	}
	row := GgufRow{FullPath: "/m/q", IsHF: true, SizeGB: 8}
	s := Settings{TargetVramGB: 22, Backends: []BackendEntry{{ID: "v", Kind: "vllm", Path: "vllm", Default: true}}}
	be := resolveBackend(s, nil, "llm")
	cmd := func(ov *Override) string { return strings.Join(vllmCmdLines(s, row, ov, "m", be, meta), " ") }

	auto := cmd(nil)
	if !strings.Contains(auto, "--kv-cache-memory-bytes ") || !strings.Contains(auto, "--gpu-memory-utilization ") {
		t.Errorf("auto: want the KV pin AND a utilization (vllm's startup free check reads it): %s", auto)
	}
	pinned := cmd(&Override{VllmGpuUtil: 0.8})
	if strings.Contains(pinned, "--kv-cache-memory-bytes") || !strings.Contains(pinned, "--gpu-memory-utilization 0.8") {
		t.Errorf("a pinned utilization must leave KV to vllm: %s", pinned)
	}
	for _, extra := range []string{"--gpu-memory-utilization 0.85", "--kv-cache-memory-bytes 4000000000"} {
		got := cmd(&Override{ExtraArgs: extra})
		if strings.Count(got, "--gpu-memory-utilization")+strings.Count(got, "--kv-cache-memory-bytes") != 1 {
			t.Errorf("extra args %q own the memory, want only theirs: %s", extra, got)
		}
	}
	// No KV cost model: nothing to pin from, vllm sizes KV itself.
	if bare := strings.Join(vllmCmdLines(s, row, nil, "m", be, Metadata{Architecture: "qwen3"}), " "); strings.Contains(bare, "--kv-cache-memory-bytes") {
		t.Errorf("no KV model, yet a pin: %s", bare)
	}
}

func TestEmitVllmModel_SkipsSplitShards(t *testing.T) {
	withCard(t, 24, true)
	s := Settings{TargetVramGB: 16, Backends: []BackendEntry{{ID: "v", Kind: "vllm", Path: "vllm", Default: true}}}
	be := resolveBackend(s, nil, "llm")
	meta := Metadata{Architecture: "qwen3", ContextLength: 32768}

	// Discovery represents a split set by shard 1 alone. llama.cpp opens the
	// siblings itself; vllm would load a fifth of the weights.
	var b strings.Builder
	var emitted []string
	split := GgufRow{FullPath: `E:\Models\big\big-00001-of-00005.gguf`, FileName: "big.gguf", SizeGB: 8}
	emitVllmModel(&b, s, split, nil, "big", be, meta, &emitted)
	if len(emitted) != 0 {
		t.Errorf("split model was emitted as %v, want skipped", emitted)
	}
	if out := b.String(); !strings.Contains(out, "skipped") || strings.Contains(out, "cmd:") {
		t.Errorf("want a skip comment and no command, got:\n%s", out)
	}

	// A single-file model still emits normally.
	b.Reset()
	emitVllmModel(&b, s, GgufRow{FullPath: `E:\Models\q\q.gguf`, SizeGB: 4}, nil, "q", be, meta, &emitted)
	if len(emitted) != 1 || !strings.Contains(b.String(), "cmd:") {
		t.Errorf("single-file model should emit, got %v:\n%s", emitted, b.String())
	}
}

func TestVllmToolParser_PickAndOverride(t *testing.T) {
	cases := []struct {
		name, arch string
		ov         *Override
		want       string
	}{
		{"qwen3-8b", "qwen3", nil, "hermes"},
		{"qwen3-coder-30b", "qwen3_moe", nil, "qwen3_coder"},
		{"qwen3.6-27b", "qwen3_5", nil, "qwen3_coder"},
		{"llama-3.1-8b", "llama", nil, "llama3_json"},
		{"mystery", "somearch", nil, ""},
		{"qwen3-8b", "qwen3", &Override{VllmToolParser: "none"}, ""},
		{"mystery", "somearch", &Override{VllmToolParser: "minicpm5"}, "minicpm5"},
	}
	for _, c := range cases {
		got, _ := vllmToolParser(c.ov, Metadata{Architecture: c.arch}, c.name)
		if got != c.want {
			t.Errorf("%s (%s) => %q, want %q", c.name, c.arch, got, c.want)
		}
	}
}

func TestVllmCmdLines_ToolParserFlags(t *testing.T) {
	s := Settings{Backends: []BackendEntry{{ID: "v", Kind: "vllm", Path: "vllm", Default: true}}}
	row := GgufRow{FullPath: "/m/q", SizeGB: 4, IsHF: true}
	meta := Metadata{Architecture: "qwen3", ContextLength: 32768}
	be := resolveBackend(s, nil, "llm")

	cmd := strings.Join(vllmCmdLines(s, row, &Override{Ctx: 8192}, "q", be, meta), " ")
	if !strings.Contains(cmd, "--enable-auto-tool-choice --tool-call-parser hermes") {
		t.Errorf("auto parser missing: %s", cmd)
	}
	// A parser already in extra args (the pre-knob workaround) is not doubled.
	ov := &Override{Ctx: 8192, ExtraArgs: "--enable-auto-tool-choice --tool-call-parser minicpm5"}
	cmd = strings.Join(vllmCmdLines(s, row, ov, "q", be, meta), " ")
	if n := strings.Count(cmd, "--tool-call-parser"); n != 1 {
		t.Errorf("--tool-call-parser appears %d times: %s", n, cmd)
	}
	// Unknown arch: no flag, and the emitted entry does not advertise tools.
	var b strings.Builder
	var emitted []string
	emitVllmModel(&b, s, row, &Override{Ctx: 8192}, "m", be, Metadata{Architecture: "somearch"}, &emitted)
	if strings.Contains(b.String(), "--tool-call-parser") || strings.Contains(b.String(), "tools: true") {
		t.Errorf("unknown arch got a parser:\n%s", b.String())
	}
	b.Reset()
	emitVllmModel(&b, s, row, &Override{Ctx: 8192}, "q", be, meta, &emitted)
	if !strings.Contains(b.String(), "tools: true") {
		t.Errorf("mapped arch not advertised as tools-capable:\n%s", b.String())
	}
}