package placement

import (
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

// cpuProfile is what the catalog produces for a llama.cpp model: no GPU class,
// because there is no card to ask for.
func cpuProfile(name string) Profile {
	return Profile{Name: name, Engine: "llamacpp"}
}

func TestCPUProfilePlacesOnANodeWithNoGPUs(t *testing.T) {
	// The whole point: a machine with no cards is useless to every other engine
	// and perfectly fine for this one.
	s := New(store(node("n1", "", 0)))

	res := s.Place([]Profile{cpuProfile("cpu")}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if len(res.Placement.Nodes) != 1 || res.Placement.Nodes[0].NodeID != "n1" {
		t.Fatalf("placed on %+v", res.Placement.Nodes)
	}
	if len(res.Placement.Nodes[0].GPUIndices) != 0 {
		t.Errorf("a CPU replica was given GPUs: %v", res.Placement.Nodes[0].GPUIndices)
	}
	if res.Placement.Pool != "none/cpu" {
		t.Errorf("pool %q", res.Placement.Pool)
	}
}

func TestCPUReplicasSpreadRatherThanStack(t *testing.T) {
	// There is no allocation to stop two CPU replicas sharing a machine's
	// cores, so placement has to do it. Bin packing, which is right for GPUs,
	// would put every replica on the same node.
	st := store(node("n1", "", 0), node("n2", "", 0), node("n3", "", 0))
	s := New(st)

	seen := map[string]int{}
	for i := range 3 {
		res := s.Place([]Profile{cpuProfile("cpu")}, nil)
		if res.Placement == nil {
			t.Fatalf("iteration %d:\n%s", i, res.Explain())
		}
		id := res.Placement.Nodes[0].NodeID
		seen[id]++
		// Record it the way the reconciler would, so the next decision sees it.
		st.Allocate(id, inventory.Allocation{ShardID: id + "-sh", ReplicaID: id + "-rep"})
	}
	if len(seen) != 3 {
		t.Errorf("three replicas landed on %d nodes: %v", len(seen), seen)
	}
}

func TestCPUProfileIgnoresGPUsItDoesNotNeed(t *testing.T) {
	// A GPU node is still a valid host, and taking it must not consume cards
	// that a GPU profile would need.
	st := store(node("gpu", "h100-80g", 8))
	s := New(st)

	res := s.Place([]Profile{cpuProfile("cpu")}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if n := len(res.Placement.Nodes[0].GPUIndices); n != 0 {
		t.Errorf("a CPU replica claimed %d GPUs", n)
	}
	if free := st.FreeGPUs("gpu"); len(free) != 8 {
		t.Errorf("GPUs were consumed: %d free", len(free))
	}
}

func TestGPUProfileStillNeedsItsClass(t *testing.T) {
	// The CPU path must not become a way for a GPU profile to be placed on a
	// machine that cannot run it.
	s := New(store(node("n1", "", 0)))

	res := s.Place([]Profile{{Name: "tp1", GPUClass: "h100-80g", GPUsPerNode: 1}}, nil)
	if res.Placement != nil {
		t.Fatal("a GPU profile was placed on a node with no GPUs")
	}
	if len(res.Rejections) == 0 {
		t.Error("the rejection should say why")
	}
}

func TestCPUProfileRespectsAvoid(t *testing.T) {
	// Spreading replicas of one model across failure domains matters the same
	// way it does for GPU replicas.
	s := New(store(node("n1", "", 0), node("n2", "", 0)))

	res := s.Place([]Profile{cpuProfile("cpu")}, []string{"n1"})
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	if got := res.Placement.Nodes[0].NodeID; got != "n2" {
		t.Errorf("placed on %q despite being told to avoid it", got)
	}
}