package modelplan

import (
	"strings"
	"testing"
)

const gib = int64(1) << 30

// A 14B GGUF: 48 layers, 8 KV heads of 128 — 192 KB a token at f16.
var shape14B = Shape{KVBytesPerToken: 48 * 8 * 256 * 2, KVHeads: 8, Params: 14_800_000_000}

func labels(opts []Option) []string {
	var out []string
	for _, o := range opts {
		out = append(out, o.Artifact+" "+o.Label())
	}
	return out
}

func TestEachFileIsMatchedToTheCardsThatHoldIt(t *testing.T) {
	nodes := []PoolNode{
		{ID: "gpu1", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 1, HostMemGB: 64, HostFreeGB: 40},
		{ID: "cpu1", HostMemGB: 32, HostFreeGB: 30},
	}
	arts := []Artifact{
		{Name: "q4.gguf", Format: "gguf", Bytes: 9 * gib},
		{Name: "q8.gguf", Format: "gguf", Bytes: 15 * gib},
		{Name: "f16.gguf", Format: "gguf", Bytes: 29 * gib},
	}
	opts := Options(arts, shape14B, nodes, 16384)
	got := strings.Join(labels(opts), "; ")
	// The best: the better precision on the one card, free now.
	if opts[0].Artifact != "q8.gguf" || opts[0].CPU || !opts[0].FreeNow {
		t.Errorf("first: %+v (all: %s)", opts[0], got)
	}
	// f16 fits no card, only a CPU node with the RAM.
	for _, o := range opts {
		if o.Artifact == "f16.gguf" && !o.CPU {
			t.Errorf("f16 placed on a card: %+v", o)
		}
	}
	if !strings.Contains(got, "f16.gguf CPU") {
		t.Errorf("f16 not offered the CPU: %s", got)
	}
	// The card busy: still the card first, said to be busy, and the CPU,
	// free, after it — for a fallback.
	nodes[0].FreeGPUs = 0
	opts = Options(arts, shape14B, nodes, 16384)
	if opts[0].CPU || opts[0].FreeNow || opts[0].Artifact != "q8.gguf" {
		t.Errorf("with the card busy: %+v", opts[0])
	}
	cpuAt := -1
	for i, o := range opts {
		if o.CPU && o.FreeNow {
			cpuAt = i
			break
		}
	}
	if cpuAt < 0 {
		t.Errorf("no free CPU way: %s", strings.Join(labels(opts), "; "))
	}
}

func TestALlamaCppModelSpreadsOverANodesCardsByLayer(t *testing.T) {
	nodes := []PoolNode{{ID: "a", Class: "l4-24g", GPUMemGB: 24, GPUs: 4, FreeGPUs: 4}}
	opts := Options([]Artifact{{Name: "70b-q4.gguf", Format: "gguf", Bytes: 40 * gib}},
		Shape{KVBytesPerToken: 80 * 8 * 256 * 2}, nodes, 8192)
	if len(opts) != 1 || opts[0].GPUsPerNode != 3 || opts[0].Nodes != 1 {
		t.Fatalf("%+v", opts)
	}
}

func TestOriginalWeightsSpanNodesOnlyOverAFabric(t *testing.T) {
	var nodes []PoolNode
	for _, id := range []string{"a", "b"} {
		nodes = append(nodes, PoolNode{ID: id, Class: "h100-80g", GPUMemGB: 80, GPUs: 8, FreeGPUs: 8, Fabric: "ib-0"})
	}
	llama70 := Shape{KVBytesPerToken: 80 * 8 * 256 * 2, KVHeads: 8, Params: 70e9}
	opts := Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 140 * gib}}, llama70, nodes, 8192)
	if len(opts) != 1 || opts[0].GPUsPerNode != 4 || opts[0].Nodes != 1 || opts[0].Engine != "vllm" {
		t.Fatalf("70B: %+v", opts)
	}
	big := Shape{KVBytesPerToken: 126 * 8 * 256 * 2, KVHeads: 8, Params: 405e9}
	opts = Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 810 * gib}}, big, nodes, 8192)
	if len(opts) != 1 || opts[0].Nodes != 2 || opts[0].GPUsPerNode != 8 || opts[0].Fabric != "ib-0" || !opts[0].FreeNow {
		t.Fatalf("405B: %+v", opts)
	}
	// No fabric, no replica across the two.
	for i := range nodes {
		nodes[i].Fabric = ""
	}
	if opts = Options([]Artifact{{Name: "original", Format: "safetensors", Bytes: 810 * gib}}, big, nodes, 8192); len(opts) != 0 {
		t.Fatalf("spanned nodes with no fabric: %+v", opts)
	}
}

func TestQuantisationsAreReadFromNames(t *testing.T) {
	for name, want := range map[string]string{
		"Qwen2.5-7B-Instruct-Q4_K_M.gguf":                       "Q4_K_M",
		"qwen2.5-coder-14b-instruct-q4_k_m-00001-of-00002.gguf": "Q4_K_M",
		"phi-4-IQ3_XS.gguf":                                     "IQ3_XS",
		"model-f16.gguf":                                        "F16",
		"JiRackUltra_1b.Q6_K.gguf":                              "Q6_K",
	} {
		if got := QuantOf(name); got != want {
			t.Errorf("%s: %s", name, got)
		}
	}
}

func TestOnTheCPUAMiddlingQuantisationIsBestAndOnACardEightBits(t *testing.T) {
	nodes := []PoolNode{{ID: "c", HostMemGB: 256, HostFreeGB: 200}}
	arts := []Artifact{
		{Name: "q4", Format: "gguf", Bytes: 9 * gib}, {Name: "q5", Format: "gguf", Bytes: 10 * gib},
		{Name: "q8", Format: "gguf", Bytes: 15 * gib}, {Name: "f16", Format: "gguf", Bytes: 28 * gib},
	}
	if o := Options(arts, shape14B, nodes, 16384); o[0].Artifact != "q5" {
		t.Errorf("CPU: %s first (%s)", o[0].Artifact, strings.Join(labels(o), "; "))
	}
	nodes = []PoolNode{{ID: "g", Class: "a100-80g", GPUMemGB: 80, GPUs: 1, FreeGPUs: 1}}
	if o := Options(arts, shape14B, nodes, 16384); o[0].Artifact != "q8" {
		t.Errorf("GPU: %s first", o[0].Artifact)
	}
}