package autogen

import (
	"path/filepath"
	"strings"
	"testing"
)

// Real output from llama-server b10837-vulkan on a box with a discrete card and
// an APU's integrated adapter. Two things matter here and both are load-bearing:
// the integrated adapter is listed (and advertises 15.8GB, well over
// minInferenceVramGB, so an eligibility floor alone will not exclude it), and
// the discrete card sorts FIRST, which telemetry need not agree with.
const llamaListDevicesOut = `ggml_vulkan: Found 2 Vulkan devices:
Available devices:
  Vulkan0: AMD Radeon RX 7900 XTX (24560 MiB, 23748 MiB free)
  Vulkan1: AMD Radeon(TM) Graphics (16186 MiB, 15377 MiB free)
`

func TestAutogen_ParseBackendDevices_Llama(t *testing.T) {
	devs := parseBackendDevices(llamaListDevicesOut)
	if len(devs) != 2 {
		t.Fatalf("parsed %d devices, want 2: %+v", len(devs), devs)
	}
	if devs[0].ID != "Vulkan0" || devs[1].ID != "Vulkan1" {
		t.Fatalf("ids = %q,%q, want Vulkan0,Vulkan1", devs[0].ID, devs[1].ID)
	}
	if devs[0].Name != "AMD Radeon RX 7900 XTX" {
		t.Fatalf("name = %q", devs[0].Name)
	}
	// 24560 MiB / 1024, matching gpuSetFromStats' MemTotalMB conversion so the
	// two figures are comparable without a unit change at the match site.
	if got := devs[0].TotalGB; got < 23.9 || got > 24.1 {
		t.Fatalf("TotalGB = %v, want ~23.98", got)
	}
	// The banner line ("Found 2 Vulkan devices:") has no MiB group and must not
	// become a device.
	for _, d := range devs {
		if d.ID == "ggml_vulkan" {
			t.Fatalf("banner parsed as a device: %+v", devs)
		}
	}
}

func TestAutogen_ParseBackendDevices_SdServer(t *testing.T) {
	// sd-server documents "one 'name<TAB>description' per line". It reports no
	// size, so these match on name alone.
	out := "cuda0\tNVIDIA GeForce RTX 4070 Ti SUPER\ncuda1\tNVIDIA GeForce RTX 3060\ncpu\tCPU\n"
	devs := parseBackendDevices(out)
	if len(devs) != 2 {
		t.Fatalf("parsed %d devices, want 2 (cpu dropped): %+v", len(devs), devs)
	}
	if devs[0].ID != "cuda0" || devs[1].ID != "cuda1" {
		t.Fatalf("ids = %q,%q", devs[0].ID, devs[1].ID)
	}
	if devs[0].TotalGB != 0 {
		t.Fatalf("TotalGB = %v, want 0 (unreported)", devs[0].TotalGB)
	}
}

// The case the whole file exists for: telemetry and the backend disagree about
// which card is device 0. Issue #4's reporter has a 3060 beside a 4070 Ti SUPER;
// give telemetry the 3060 first and the backend the faster card first, which is
// exactly what CUDA_DEVICE_ORDER=FASTEST_FIRST produces.
func TestAutogen_BackendIDs_ReversedEnumeration(t *testing.T) {
	set := GpuSet{
		{Index: 0, Name: "NVIDIA GeForce RTX 3060", TotalGB: 12, FreeGB: 11},
		{Index: 1, Name: "NVIDIA GeForce RTX 4070 Ti SUPER", TotalGB: 16, FreeGB: 15},
	}
	devs := []BackendDevice{
		{ID: "CUDA0", Name: "NVIDIA GeForce RTX 4070 Ti SUPER", TotalGB: 16},
		{ID: "CUDA1", Name: "NVIDIA GeForce RTX 3060", TotalGB: 12},
	}
	got := set.BackendIDs(devs)
	if len(got) != 2 || got[0] != "CUDA1" || got[1] != "CUDA0" {
		t.Fatalf("BackendIDs = %v, want [CUDA1 CUDA0]: the ids must follow OUR order, not the backend's", got)
	}
}

// An adapter we filtered out (an iGPU under minInferenceVramGB, or a card the
// user excluded) is still in the backend's list. Naming devices is what keeps it
// out of the split, so the mapping must skip it rather than shift onto it.
func TestAutogen_BackendIDs_SkipsUnplannedAdapter(t *testing.T) {
	set := GpuSet{{Index: 1, Name: "AMD Radeon RX 7900 XTX", TotalGB: 23.98, FreeGB: 23}}
	devs := parseBackendDevices(llamaListDevicesOut)
	got := set.BackendIDs(devs)
	if len(got) != 1 || got[0] != "Vulkan0" {
		t.Fatalf("BackendIDs = %v, want [Vulkan0]", got)
	}
}

// Refusal is the safety property: a set we cannot fully place must produce no
// device list at all, so the caller falls back to unnamed placement instead of
// emitting one that silently drops a card the sizer budgeted for.
func TestAutogen_BackendIDs_RefusesPartialMatch(t *testing.T) {
	set := GpuSet{
		{Index: 0, Name: "AMD Radeon RX 7900 XTX", TotalGB: 23.98, FreeGB: 23},
		{Index: 1, Name: "NVIDIA GeForce RTX 3060", TotalGB: 12, FreeGB: 11},
	}
	if got := set.BackendIDs(parseBackendDevices(llamaListDevicesOut)); got != nil {
		t.Fatalf("BackendIDs = %v, want nil when a planned card is not in the backend's list", got)
	}
	if got := set.BackendIDs(nil); got != nil {
		t.Fatalf("BackendIDs with no listing = %v, want nil", got)
	}
}

// Same name, same capacity, twice. Either assignment is correct, so the mapping
// must succeed rather than refuse: two identical cards are interchangeable.
func TestAutogen_BackendIDs_IdenticalCards(t *testing.T) {
	set := GpuSet{
		{Index: 0, Name: "NVIDIA GeForce RTX 3090", TotalGB: 24, FreeGB: 23},
		{Index: 1, Name: "NVIDIA GeForce RTX 3090", TotalGB: 24, FreeGB: 23},
	}
	devs := []BackendDevice{
		{ID: "CUDA0", Name: "NVIDIA GeForce RTX 3090", TotalGB: 24},
		{ID: "CUDA1", Name: "NVIDIA GeForce RTX 3090", TotalGB: 24},
	}
	got := set.BackendIDs(devs)
	if len(got) != 2 || got[0] == got[1] {
		t.Fatalf("BackendIDs = %v, want two distinct ids", got)
	}
}

// Vendor prefixes, trademark marks and brand words differ between telemetry and
// ggml for the same silicon. The model designation is what has to survive.
func TestAutogen_DeviceNamesMatch(t *testing.T) {
	match := [][2]string{
		{"AMD Radeon(TM) Graphics", "Radeon Graphics"},
		{"NVIDIA GeForce RTX 4070 Ti SUPER", "GeForce RTX 4070 Ti SUPER"},
		{"AMD Radeon RX 7900 XTX", "Radeon RX 7900 XTX"},
	}
	for _, p := range match {
		if !deviceNamesMatch(p[0], p[1]) {
			t.Errorf("deviceNamesMatch(%q, %q) = false, want true", p[0], p[1])
		}
	}
	// A 4070 beside a 4070 Ti is the case containment must NOT collapse, since
	// picking the wrong one hands the split to the wrong card.
	differ := [][2]string{
		{"NVIDIA GeForce RTX 4070 Ti SUPER", "NVIDIA GeForce RTX 3060"},
		{"AMD Radeon RX 7900 XTX", "AMD Radeon(TM) Graphics"},
	}
	for _, p := range differ {
		if deviceNamesMatch(p[0], p[1]) {
			t.Errorf("deviceNamesMatch(%q, %q) = true, want false", p[0], p[1])
		}
	}
}

// DeviceFlagFor's refusal paths, which are the ones that decide whether a
// launch keeps the argv that shipped. The probing path needs a real backend
// binary and is covered by parseBackendDevices + BackendIDs above.
func TestAutogen_DeviceFlagFor_Refusals(t *testing.T) {
	two := GpuSet{
		{Index: 0, Name: "NVIDIA GeForce RTX 3060", TotalGB: 12, FreeGB: 11},
		{Index: 3, Name: "NVIDIA GeForce RTX 4070 Ti SUPER", TotalGB: 16, FreeGB: 15},
	}
	// No probe can help a set that isn't split.
	if devs, order := DeviceFlagFor("llama-server", GpuSet{two[0]}, 0); devs != "" || order != nil {
		t.Fatalf("single-device DeviceFlagFor = %q,%v, want \"\",nil", devs, order)
	}
	// A main index that names no device in the set: nothing to pin it to.
	if devs, order := DeviceFlagFor("llama-server", two, 1); devs != "" || order != nil {
		t.Fatalf("unknown main index = %q,%v, want \"\",nil", devs, order)
	}
	// A binary that cannot be run at all degrades to unnamed placement rather
	// than to an error.
	if devs, order := DeviceFlagFor(filepath.Join(t.TempDir(), "not-a-backend"), two, 3); devs != "" || order != nil {
		t.Fatalf("missing exe = %q,%v, want \"\",nil", devs, order)
	}
}

// The placement rule the whole ordering rests on: measured on a Vulkan build,
// -sm layer puts the non-splittable output weight on the device listed LAST and
// ignores --main-gpu. So the main device has to be last, and everything else
// keeps its relative order (the split vector is only meaningful positionally).
func TestAutogen_MainLastOrder(t *testing.T) {
	three := GpuSet{
		{Index: 0, Name: "a", TotalGB: 12, FreeGB: 11},
		{Index: 3, Name: "b", TotalGB: 16, FreeGB: 15},
		{Index: 7, Name: "c", TotalGB: 24, FreeGB: 23},
	}
	// Main is telemetry index 3, which is POSITION 1. Only the named form can
	// say so, and it must end up last.
	got := three.MainLastOrder(3)
	if len(got) != 3 || got[0] != 0 || got[1] != 2 || got[2] != 1 {
		t.Fatalf("MainLastOrder(3) = %v, want [0 2 1]", got)
	}
	// Already last: an identity permutation, not a rotation.
	if got := three.MainLastOrder(7); len(got) != 3 || got[0] != 0 || got[1] != 1 || got[2] != 2 {
		t.Fatalf("MainLastOrder(7) = %v, want [0 1 2]", got)
	}
	if got := three.MainLastOrder(9); got != nil {
		t.Fatalf("MainLastOrder of an absent index = %v, want nil", got)
	}
	// The split must travel with the list or each card gets the other's ratio.
	if got := PermuteSplit([]float64{0.2, 0.3, 0.5}, []int{0, 2, 1}); got[0] != 0.2 || got[1] != 0.5 || got[2] != 0.3 {
		t.Fatalf("PermuteSplit = %v, want [0.2 0.5 0.3]", got)
	}
	// A length that cannot be mapped is left alone rather than rearranged.
	in := []float64{0.5, 0.5}
	if got := PermuteSplit(in, []int{0, 2, 1}); len(got) != 2 || got[0] != 0.5 {
		t.Fatalf("PermuteSplit with a mismatched order = %v, want the input back", got)
	}
}

// The probe has to enumerate the way the LAUNCH will. Every multi-GPU config we
// emit carries CUDA_DEVICE_ORDER=PCI_BUS_ID, while the CUDA runtime defaults to
// FASTEST_FIRST: a listing read under the inherited environment numbers a
// mismatched pair the other way round, so "CUDA0" handed back to a bus-ordered
// process names the OTHER card. That reversal is the exact failure --device
// exists to prevent, and it is silent.
func TestAutogen_ProbeEnv_PinsBusOrder(t *testing.T) {
	var found bool
	for _, kv := range probeEnv() {
		if kv == cudaOrderEnv {
			found = true
		}
	}
	if !found {
		t.Fatalf("probe environment does not pin %s", cudaOrderEnv)
	}
}

// A single-GPU box whose BACKEND also sees one GPU is the population that never
// asked for any of this. It must get no probe, no --device and no env block:
// byte-identical to what shipped.
func TestAutogen_SingleGpu_EmitsNothingNew(t *testing.T) {
	one := GpuSet{{Index: 0, Name: "NVIDIA GeForce RTX 4090", TotalGB: 24, FreeGB: 23}}
	if devs, order := DeviceFlagFor(filepath.Join(t.TempDir(), "nope"), one, 0); devs != "" || order != nil {
		t.Fatalf("single-GPU DeviceFlagFor = %q,%v, want no flags", devs, order)
	}
	var b strings.Builder
	writeSingleDeviceEnv(&b, Settings{Gpus: one}, filepath.Join(t.TempDir(), "nope"))
	if b.String() != "" {
		t.Fatalf("single-GPU writeSingleDeviceEnv emitted %q, want nothing", b.String())
	}
}

// The other single-GPU box: telemetry counts one card, the backend counts two.
// An integrated adapter has no dedicated VRAM to plan against so the GPU set
// drops it, but ROCm and Vulkan list it anyway, advertising shared system
// memory as VRAM. Left unpinned, sd-server's own placement planner sizes against
// a card whose kernels the build does not contain. Pin the real one.
func TestAutogen_SingleGpu_PinsWhenBackendSeesMore(t *testing.T) {
	exe := filepath.Join(t.TempDir(), "sd-server-no-such-file")
	seedBackendDevCache(t, exe, []BackendDevice{
		{ID: "ROCm0", Name: "AMD Radeon RX 7900 XTX", TotalGB: 23.98},
		{ID: "ROCm1", Name: "AMD Radeon(TM) Graphics", TotalGB: 12.17},
	})

	one := GpuSet{{Index: 0, Name: "AMD Radeon RX 7900 XTX", TotalGB: 23.95, FreeGB: 23}}
	var b strings.Builder
	writeSingleDeviceEnv(&b, Settings{Gpus: one}, exe)

	got := b.String()
	if !strings.Contains(got, "HIP_VISIBLE_DEVICES=0") {
		t.Fatalf("writeSingleDeviceEnv = %q, want a HIP pin to device 0", got)
	}
	// CUDA_DEVICE_ORDER means nothing to HIP and would be noise in the config.
	if strings.Contains(got, "CUDA_DEVICE_ORDER") {
		t.Fatalf("writeSingleDeviceEnv = %q, want no CUDA ordering for a ROCm pin", got)
	}
}

// seedBackendDevCache plants a listing for an exe that will never run, so the
// mapping is assertable without a backend binary. The cache key is the plain
// path when the file does not exist (see ListBackendDevices).
func seedBackendDevCache(t *testing.T, exe string, devs []BackendDevice) {
	t.Helper()
	backendDevMu.Lock()
	backendDevCache[exe] = devs
	backendDevMu.Unlock()
	t.Cleanup(func() {
		backendDevMu.Lock()
		delete(backendDevCache, exe)
		backendDevMu.Unlock()
	})
}

// A generate reaches up to five distinct backend binaries. The per-probe window
// is deliberately generous, so without a shared budget a box where every backend
// is wedged multiplies it into a startup that visibly hangs.
func TestAutogen_ProbeBudget_BoundsTheWholeGenerate(t *testing.T) {
	probeBudgetMu.Lock()
	saved := probeBudgetLeft
	probeBudgetMu.Unlock()
	t.Cleanup(func() {
		probeBudgetMu.Lock()
		probeBudgetLeft = saved
		probeBudgetMu.Unlock()
	})

	if d := takeProbeBudget(); d != backendProbeTimeout {
		t.Fatalf("first probe window = %v, want the full %v", d, backendProbeTimeout)
	}
	spendProbeBudget(backendProbeTimeout)
	if d := takeProbeBudget(); d <= 0 || d > backendProbeBudget-backendProbeTimeout {
		t.Fatalf("second probe window = %v, want what is left of %v", d, backendProbeBudget)
	}
	spendProbeBudget(backendProbeBudget)
	if d := takeProbeBudget(); d != 0 {
		t.Fatalf("exhausted probe window = %v, want 0", d)
	}
	if _, err := probeBackendDevices(filepath.Join(t.TempDir(), "nope")); err == nil {
		t.Fatal("probe with no budget returned no error")
	}
}

// A hand-written --tensor-split is positional against the list the backend would
// have enumerated on its own. --device replaces that list in main-last order, so
// emitting both would silently point the user's ratio at a different pair.
func TestAutogen_PinsOwnSplit(t *testing.T) {
	if pinsOwnSplit(nil) {
		t.Fatal("no override pins nothing")
	}
	if pinsOwnSplit(&Override{}) {
		t.Fatal("blank tensorSplit pins nothing")
	}
	if pinsOwnSplit(&Override{TensorSplit: "  "}) {
		t.Fatal("whitespace tensorSplit pins nothing")
	}
	if !pinsOwnSplit(&Override{TensorSplit: "3,1"}) {
		t.Fatal("a hand-written ratio must suppress --device")
	}
}