package modelplan

import (
	"math"
	"testing"
)

// Real models, with the figures measured on this cluster rather than derived
// from a formula — the point of testing arithmetic is that it agrees with the
// world, and a test built from the same formula would agree with itself.
var (
	coder3B = Model{
		Name: "Qwen2.5-Coder-3B", Params: 3_090_000_000,
		Layers: 36, KVHeads: 2, HeadDim: 128, MaxContext: 32768,
	}
	coder14B = Model{
		Name: "Qwen2.5-Coder-14B", Params: 14_770_000_000,
		Layers: 48, KVHeads: 8, HeadDim: 128, MaxContext: 32768,
	}
	llama70B = Model{
		Name: "Llama-3.3-70B", Params: 70_600_000_000,
		Layers: 80, KVHeads: 8, HeadDim: 128, MaxContext: 131072,
	}
)

func closeTo(t *testing.T, what string, got, want, tolerance float64) {
	t.Helper()
	if math.Abs(got-want) > tolerance {
		t.Errorf("%s = %.2f, want %.2f (±%.2f)", what, got, want, tolerance)
	}
}

func TestKVCachePerTokenMatchesTheArchitecture(t *testing.T) {
	// 192 KiB a token for the 14B is the number that made sense of what this
	// cluster measured: 8192 of context across four llama.cpp slots is 2048
	// each, and six gigabytes of cache at the model's full 32k.
	if got := coder14B.KVBytesPerToken(BF16); got != 192*1024 {
		t.Errorf("14B KV/token = %d bytes, want %d", got, 192*1024)
	}
	// Grouped-query attention is why this is KVHeads and not the attention
	// head count. The 14B has 40 attention heads and 8 KV heads, and paying for
	// 40 would overstate the cache fivefold.
	if got := coder3B.KVBytesPerToken(BF16); got != 36*1024 {
		t.Errorf("3B KV/token = %d bytes, want %d", got, 36*1024)
	}
	if got := llama70B.KVBytesPerToken(BF16); got != 320*1024 {
		t.Errorf("70B KV/token = %d bytes, want %d", got, 320*1024)
	}
}

func TestQuantisedWeightsMatchTheFilesOnDisk(t *testing.T) {
	// Calibrated against the GGUFs downloaded to this cluster, so the
	// bytes-per-parameter figures are measurements rather than folklore.
	// In GiB, which is what llama-bench itself reported for these files —
	// "14.62 GiB" and "8.37 GiB". The download sizes quoted in decimal GB are
	// the same bytes in different units, and mixing the two is how a plan ends
	// up seven percent optimistic.
	for _, tc := range []struct {
		p    Precision
		want float64 // GiB, as llama.cpp accounts for it
	}{
		{Q8GGUF, 14.62},
		{Q4GGUF, 8.37},
	} {
		got := GB(coder14B.Needs(tc.p, 0).Weights)
		closeTo(t, "14B "+tc.p.Name+" weights", got, tc.want, 0.35)
	}
	// And full precision, which is two bytes a parameter by definition.
	closeTo(t, "14B bf16 weights", GB(coder14B.Needs(BF16, 0).Weights), 27.5, 0.6)
}

func TestContextDominatesAtLongWindows(t *testing.T) {
	// The reason KV cache is worth computing rather than hand-waving: at 128k
	// of pooled context a 70B's cache is a substantial fraction of its weights,
	// and a plan that ignores it fits a model that then will not load.
	short := llama70B.Needs(BF16, 4096)
	long := llama70B.Needs(BF16, 131072)
	if short.Weights != long.Weights {
		t.Fatal("weights should not depend on context")
	}
	if GB(long.KVCache) < 35 {
		t.Errorf("70B KV at 128k = %.1f GB, expected tens of GB", GB(long.KVCache))
	}
	if long.KVCache <= short.KVCache*8 {
		t.Error("KV cache should scale with the context pool")
	}
}

func TestTensorParallelismDividesWeightsAndCacheButNotOverhead(t *testing.T) {
	// Each rank holds its share of the weights and of the KV heads, so both
	// divide. The runtime does not: every process pays its own CUDA context and
	// activation buffers, which is why eight-way parallelism does not make a
	// model eight times cheaper per card.
	d := llama70B.Needs(BF16, 32768)

	one := d.PerGPU(1)
	eight := d.PerGPU(8)
	if eight >= one {
		t.Fatal("splitting across GPUs should reduce the per-GPU demand")
	}
	// Strictly more than an eighth, because the overhead is not shared.
	if eight <= one/8 {
		t.Errorf("per-GPU at tp=8 is %d, at or below a clean eighth of %d: overhead is being divided",
			eight, one)
	}
	if d.PerGPU(0) != d.PerGPU(1) {
		t.Error("a nonsensical GPU count should be treated as one, not divided by zero")
	}
}

func TestA70BNeedsMoreThanOneCard(t *testing.T) {
	// The question someone with hardware actually asks, and the answer that is
	// easy to get wrong by hand: 70B at bf16 does not fit an 80 GB card, and
	// does fit two.
	d := llama70B.Needs(BF16, 32768)
	const h100 = 80.0

	if GB(d.PerGPU(1)) <= h100 {
		t.Errorf("70B bf16 reported as fitting one 80 GB card (%.1f GB)", GB(d.PerGPU(1)))
	}
	if GB(d.PerGPU(2)) > h100 {
		t.Errorf("70B bf16 across two cards is %.1f GB each, which should fit 80", GB(d.PerGPU(2)))
	}
}

func TestOverheadIsChargedEvenWithNoContext(t *testing.T) {
	// A plan that forgets the runtime fits a model into exactly its weights and
	// then watches it fail to load.
	d := coder3B.Needs(BF16, 0)
	if d.Overhead <= 0 {
		t.Error("no overhead was charged")
	}
	if d.Total() <= d.Weights {
		t.Error("the total is not more than the weights alone")
	}
}

func TestPrecisionByNameRefusesWhatItDoesNotKnow(t *testing.T) {
	if _, err := PrecisionByName("bf16"); err != nil {
		t.Errorf("bf16 was refused: %v", err)
	}
	_, err := PrecisionByName("fp3")
	if err == nil {
		t.Fatal("an unknown precision was accepted")
	}
	// The error lists what would work, because the caller mistyped something
	// and a bare refusal makes them go looking.
	if got := err.Error(); !contains(got, "q4_k_m") || !contains(got, "bf16") {
		t.Errorf("the error does not say what is known: %q", got)
	}
}

func contains(s, sub string) bool {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return true
		}
	}
	return false
}