package modelplan

import (
	"strings"
	"testing"
)

func h100Node(nodes, gpus int, fabric string) Hardware {
	return Hardware{Class: "h100-80g", GPUMemGB: 80, GPUsPerNode: gpus, Nodes: nodes, Fabric: fabric}
}

func TestASmallModelFitsOneCard(t *testing.T) {
	fits := Fits(coder3B, []Hardware{h100Node(4, 8, "ib-a")}, 32768, []Precision{BF16})
	if len(fits) == 0 {
		t.Fatal("a 3B was reported as not fitting an H100")
	}
	// Best first, and the best answer for a small model is one card.
	if fits[0].TP != 1 {
		t.Errorf("best fit uses %d GPUs; a 3B needs one", fits[0].TP)
	}
}

func TestA70BNeedsTwoCardsAtFullPrecision(t *testing.T) {
	// The configuration people run every day, and the one an earlier overhead
	// model wrongly called impossible by reserving slack twice.
	fits := Fits(llama70B, []Hardware{h100Node(4, 8, "ib-a")}, 8192, []Precision{BF16})
	if len(fits) == 0 {
		t.Fatal("70B bf16 reported as not fitting eight H100s")
	}
	if fits[0].TP != 2 {
		t.Errorf("best fit is tp=%d; 70B at bf16 on two 80 GB cards is routine", fits[0].TP)
	}
	for _, f := range fits {
		if f.TP == 1 {
			t.Errorf("70B bf16 reported as fitting a single card: %s", f.Describe())
		}
	}
}

func TestContextChangesHowManyCardsAModelNeeds(t *testing.T) {
	// The fact that makes this worth computing rather than looking up. The same
	// 70B at the same precision needs two cards for an 8k pool and four for
	// 32k, because the KV cache is what grows — and "70B needs two H100s" is
	// repeated everywhere without the context it was true for.
	hw := []Hardware{h100Node(4, 8, "ib-a")}
	short := Fits(llama70B, hw, 8192, []Precision{BF16})
	long := Fits(llama70B, hw, 32768, []Precision{BF16})
	if len(short) == 0 || len(long) == 0 {
		t.Fatal("expected fits at both context sizes")
	}
	if long[0].TP <= short[0].TP {
		t.Errorf("tp at 32k is %d and at 8k is %d; a larger context should need at least as much",
			long[0].TP, short[0].TP)
	}
}

func TestQuantisationOpensUpSmallerHardware(t *testing.T) {
	// The answer someone with one spare card needs: not "no", but "not at this
	// precision".
	one := []Hardware{h100Node(1, 1, "")}
	if full := Fits(llama70B, one, 8192, []Precision{BF16}); len(full) != 0 {
		t.Errorf("70B bf16 fitted a single card: %s", full[0].Describe())
	}
	if quant := Fits(llama70B, one, 8192, []Precision{Q4GGUF}); len(quant) == 0 {
		t.Error("70B at q4 did not fit a single 80 GB card, which it should")
	}
}

func TestATensorGroupNeverSpansFabrics(t *testing.T) {
	// Two nodes with no fabric are two islands. A multi-node replica placed
	// across them would start, pass health checks and hang on its first
	// all-reduce, which is worse than refusing to plan it.
	noFabric := Hardware{Class: "h100-80g", GPUMemGB: 80, GPUsPerNode: 2, Nodes: 4}
	for _, f := range Fits(llama70B, []Hardware{noFabric}, 8192, []Precision{BF16}) {
		if f.Nodes > 1 {
			t.Errorf("planned %d nodes with no fabric: %s", f.Nodes, f.Describe())
		}
	}
	withFabric := noFabric
	withFabric.Fabric = "ib-a"
	var spanning bool
	for _, f := range Fits(llama70B, []Hardware{withFabric}, 8192, []Precision{BF16}) {
		if f.Nodes > 1 {
			spanning = true
		}
	}
	if !spanning {
		t.Error("a fabric was available and no multi-node fit was offered")
	}
}

func TestTensorWidthCannotExceedTheKVHeads(t *testing.T) {
	// Grouped-query attention leaves far fewer KV heads than attention heads,
	// and an engine refuses a width that cannot divide them. Offering tp=16 for
	// a model with 8 KV heads is a plan that fails at launch.
	for _, f := range Fits(llama70B, []Hardware{h100Node(8, 8, "ib-a")}, 8192, []Precision{Q4GGUF}) {
		if f.TP > llama70B.KVHeads {
			t.Errorf("tp=%d exceeds the model's %d KV heads", f.TP, llama70B.KVHeads)
		}
	}
}

func TestOnlyPowersOfTwoAreOffered(t *testing.T) {
	for _, f := range Fits(llama70B, []Hardware{h100Node(4, 6, "ib-a")}, 8192, []Precision{Q4GGUF}) {
		if f.TP&(f.TP-1) != 0 {
			t.Errorf("tp=%d is not a power of two", f.TP)
		}
	}
}

func TestACardIsNeverPlannedToTheBrim(t *testing.T) {
	// A plan that fills a card completely is one that loads on a good day.
	// This matches vLLM's own 0.90 default deliberately: a planner that
	// disagrees with the engine it plans for is worse than none.
	for _, f := range Fits(coder14B, []Hardware{h100Node(1, 8, "")}, 32768, []Precision{BF16}) {
		if f.PerGPUGB > 80*usableFraction+0.01 {
			t.Errorf("planned %.1f GB on an 80 GB card: %s", f.PerGPUGB, f.Describe())
		}
		if f.Headroom < 0 {
			t.Errorf("negative headroom: %s", f.Describe())
		}
	}
}

func TestACPUNodeCanServeASmallModel(t *testing.T) {
	// How this cluster began, and a reasonable way to try llmgw before
	// committing cards to it.
	cpu := Hardware{HostMemGB: 32}
	fits := Fits(coder3B, []Hardware{cpu}, 32768, []Precision{Q4GGUF})
	if len(fits) == 0 {
		t.Fatal("a 3B at q4 did not fit a 32 GB CPU node")
	}
	if !fits[0].CPUOnly {
		t.Error("the fit is not marked CPU-only")
	}
	if !strings.Contains(fits[0].Describe(), "CPU") {
		t.Errorf("the description does not say it is CPU: %q", fits[0].Describe())
	}
	// And a 70B does not, whatever the quantisation.
	if big := Fits(llama70B, []Hardware{cpu}, 8192, []Precision{Q4GGUF}); len(big) != 0 {
		t.Errorf("a 70B was reported as fitting 32 GB of RAM: %s", big[0].Describe())
	}
}

func TestNothingFitsIsAnEmptyAnswerNotAWrongOne(t *testing.T) {
	tiny := Hardware{Class: "t4-16g", GPUMemGB: 16, GPUsPerNode: 1, Nodes: 1}
	if fits := Fits(llama70B, []Hardware{tiny}, 32768, []Precision{BF16}); len(fits) != 0 {
		t.Errorf("70B bf16 fitted a single 16 GB card: %s", fits[0].Describe())
	}
}

func TestIdenticalHardwareOnTwoFabricsIsNotOfferedTwice(t *testing.T) {
	// A single-node replica does not span anything, so which fabric its node is
	// attached to changes nothing about whether the model fits. Listing it once
	// per fabric is noise: on a real fleet a 32B produced twenty-six rows of
	// which thirteen were duplicates.
	hw := []Hardware{h100Node(2, 4, "ib-a"), h100Node(2, 4, "ib-b")}
	fits := Fits(coder14B, hw, 8192, []Precision{BF16})

	seen := map[int]int{}
	for _, f := range fits {
		if f.Nodes == 1 {
			seen[f.TP]++
		}
	}
	for tp, n := range seen {
		if n > 1 {
			t.Errorf("single-node tp=%d offered %d times", tp, n)
		}
	}
}

func TestAMultiNodeFitIsStillOfferedPerFabric(t *testing.T) {
	// The other half: a replica that spans nodes is fabric-specific, and which
	// fabric has room is exactly what the operator needs told.
	hw := []Hardware{h100Node(2, 2, "ib-a"), h100Node(2, 2, "ib-b")}
	fabrics := map[string]bool{}
	for _, f := range Fits(llama70B, hw, 8192, []Precision{Q4GGUF}) {
		if f.Nodes > 1 {
			fabrics[f.Fabric] = true
		}
	}
	if len(fabrics) < 2 {
		t.Errorf("multi-node fits were offered for %d fabrics, want both: %v", len(fabrics), fabrics)
	}
}