package modelplan

import "testing"

func TestKVHeadsFallBackToAttentionHeads(t *testing.T) {
	// A config without num_key_value_heads is multi-head attention, where every
	// attention head keeps its own cache. Defaulting to zero instead would make
	// the KV cache vanish from the calculation entirely.
	m, err := hfConfig{
		NumHiddenLayers: 32, NumAttentionHeads: 32, HiddenSize: 4096,
		MaxPositions: 4096,
	}.toModel("org/old-model")
	if err != nil {
		t.Fatal(err)
	}
	if m.KVHeads != 32 {
		t.Errorf("KVHeads = %d, want the attention head count", m.KVHeads)
	}
	if m.KVBytesPerToken(BF16) == 0 {
		t.Error("the KV cache came out as nothing")
	}
}

func TestHeadDimIsDerivedFromAttentionHeadsNotKVHeads(t *testing.T) {
	// The hidden size is split across the query heads. Dividing by the KV heads
	// instead inflates head_dim by the grouping factor — eight, on a typical
	// model — and the error lands straight on the cache figure.
	m, err := hfConfig{
		NumHiddenLayers: 48, NumAttentionHeads: 40, NumKeyValueHeads: 8,
		HiddenSize: 5120, MaxPositions: 32768,
	}.toModel("Qwen/Qwen2.5-Coder-14B-Instruct")
	if err != nil {
		t.Fatal(err)
	}
	if m.HeadDim != 128 {
		t.Errorf("HeadDim = %d, want 5120/40 = 128", m.HeadDim)
	}
	// Which reproduces the figure measured on this cluster.
	if got := m.KVBytesPerToken(BF16); got != 192*1024 {
		t.Errorf("KV/token = %d, want %d", got, 192*1024)
	}
}

func TestAnExplicitHeadDimWins(t *testing.T) {
	// Newer configs state it, and some models do not have hidden_size divisible
	// by the head count — deriving it there is simply wrong.
	m, err := hfConfig{
		NumHiddenLayers: 32, NumAttentionHeads: 32, NumKeyValueHeads: 8,
		HiddenSize: 4096, HeadDim: 256, MaxPositions: 8192,
	}.toModel("org/model")
	if err != nil {
		t.Fatal(err)
	}
	if m.HeadDim != 256 {
		t.Errorf("HeadDim = %d; the explicit value should win over 4096/32", m.HeadDim)
	}
}

func TestAConfigWithoutLayersIsRefusedClearly(t *testing.T) {
	// Pointing this at an embedding or vision repository is an easy mistake,
	// and a plan built from zeros would be confidently meaningless.
	_, err := hfConfig{MaxPositions: 512}.toModel("org/not-a-text-model")
	if err == nil {
		t.Fatal("a config with no layers was accepted")
	}
	if !contains(err.Error(), "num_hidden_layers") {
		t.Errorf("the error does not name what is missing: %v", err)
	}
}

func TestTheModelNameIsTheRepositoryLeaf(t *testing.T) {
	m, _ := hfConfig{NumHiddenLayers: 1, NumAttentionHeads: 1, HiddenSize: 64}.
		toModel("Qwen/Qwen2.5-Coder-14B-Instruct")
	if m.Name != "Qwen2.5-Coder-14B-Instruct" {
		t.Errorf("Name = %q", m.Name)
	}
}