package placement

import (
	"slices"
	"strings"
	"testing"

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

// migNode is an A100 80GB node as discovery reports it with card 0 cut into
// two 3g.40gb instances and card 1 whole.
func migNode(id string) inventory.Node {
	card := discovery.GPU{Name: "NVIDIA A100-SXM4-80GB", MemoryTotalMB: 81920, ComputeCapability: "8.0"}
	nd := inventory.Node{ID: id, State: inventory.Ready}
	for dev, uuid := range []string{"MIG-c6d4f1ef", "MIG-0f7c3a2e"} {
		g := card
		g.Index, g.UUID, g.MemoryTotalMB = discovery.MIGIndex(0, dev), uuid, 40192
		g.Name += " MIG 3g.40gb"
		g.MIG = &discovery.MIGDevice{ParentIndex: 0, ParentMemoryTotalMB: 81920, Device: dev, Profile: "3g.40gb"}
		nd.GPUs = append(nd.GPUs, g)
	}
	whole := card
	whole.Index, whole.UUID = 1, "GPU-whole"
	// NVLink to the card in MIG mode, as topo -m still reports it: nothing
	// may take that for a link an instance can use.
	whole.NVLinkPeers = []int{0}
	nd.GPUs = append(nd.GPUs, whole)
	return nd
}

func TestAModelIsPlacedOnAMIGInstanceOfItsClass(t *testing.T) {
	st := store(migNode("n1"))
	s := New(st)
	p := Profile{Name: "mig", GPUClass: "a100-80g-mig-3g.40gb", GPUsPerNode: 1}

	res := s.Place([]Profile{p}, nil)
	if res.Placement == nil {
		t.Fatalf("no placement:\n%s", res.Explain())
	}
	got := res.Placement.Nodes[0].GPUIndices
	if !slices.Equal(got, []int{1000}) {
		t.Fatalf("placed on %v, want MIG instance 1000", got)
	}
	if res.Placement.Pool != "none/a100-80g-mig-3g.40gb" {
		t.Errorf("pool %q", res.Placement.Pool)
	}

	// The next replica takes the other instance, and the one after finds
	// none: the whole card beside them is another class.
	st.Allocate("n1", inventory.Allocation{ShardID: "s1", ReplicaID: "r1", GPUIndices: got})
	res = s.Place([]Profile{p}, nil)
	if res.Placement == nil || !slices.Equal(res.Placement.Nodes[0].GPUIndices, []int{1001}) {
		t.Fatalf("second replica: %+v\n%s", res.Placement, res.Explain())
	}
	st.Allocate("n1", inventory.Allocation{ShardID: "s2", ReplicaID: "r2", GPUIndices: []int{1001}})
	if res = s.Place([]Profile{p}, nil); res.Placement != nil {
		t.Errorf("a third replica was placed on %v", res.Placement.Nodes[0].GPUIndices)
	}

	// A whole-card profile takes the whole card, never an instance.
	res = s.Place([]Profile{{Name: "whole", GPUClass: "a100-80g", GPUsPerNode: 1}}, nil)
	if res.Placement == nil || !slices.Equal(res.Placement.Nodes[0].GPUIndices, []int{1}) {
		t.Errorf("whole-card profile: %+v\n%s", res.Placement, res.Explain())
	}
}

func TestMoreThanOneMIGInstanceForAReplicaIsRefused(t *testing.T) {
	// Two instances are free on one node, which is what a tp 2 profile would
	// otherwise be given — and fail at start, since a process opens only one
	// MIG device and instances have no NVLink between them.
	s := New(store(migNode("n1"), migNode("n2")))
	for _, p := range []Profile{
		{Name: "tp2", GPUClass: "a100-80g-mig-3g.40gb", GPUsPerNode: 2},
		{Name: "two-nodes", GPUClass: "a100-80g-mig-3g.40gb", GPUsPerNode: 1, Nodes: 2},
	} {
		res := s.Place([]Profile{p}, nil)
		if res.Placement != nil {
			t.Errorf("%s was placed: %+v", p.Name, res.Placement.Nodes)
			continue
		}
		if !strings.Contains(res.Explain(), "a replica runs on exactly one") {
			t.Errorf("%s refused without saying why:\n%s", p.Name, res.Explain())
		}
	}
}