package adminapi

import (
	"reflect"
	"testing"

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

// migNode has card 0 cut into two 3g.40gb instances and card 1 whole.
func migNode(id string) inventory.Node {
	nd := inventory.Node{ID: id, State: inventory.Ready, AgentVersion: "v1", MemoryTotalMB: 512 * 1024}
	for dev := range 2 {
		nd.GPUs = append(nd.GPUs, discovery.GPU{
			Index: discovery.MIGIndex(0, dev), Name: "NVIDIA A100-SXM4-80GB MIG 3g.40gb", UUID: "MIG-x",
			MemoryTotalMB: 40192, MemoryUsedMB: 37,
			MIG: &discovery.MIGDevice{ParentIndex: 0, ParentMemoryTotalMB: 81920, Device: dev, Profile: "3g.40gb"},
		})
	}
	nd.GPUs = append(nd.GPUs, discovery.GPU{Index: 1, Name: "NVIDIA A100-SXM4-80GB", MemoryTotalMB: 81920})
	return nd
}

func TestANodeWithACardInMIGModeCountsItsDevicesByClass(t *testing.T) {
	a, ctl := setup(t, testModels(), migNode("a100"))
	ctl.Inventory().Allocate("a100", inventory.Allocation{ShardID: "s", ReplicaID: "r", GPUIndices: []int{1000}})

	n := a.Nodes()[0]
	want := []ClassGPUs{
		{Class: "a100-80g-mig-3g.40gb", Total: 2, Free: 1, MIG: true},
		{Class: "a100-80g", Total: 1, Free: 1},
	}
	if !reflect.DeepEqual(n.ClassGPUs, want) {
		t.Errorf("by class: %+v\nwant %+v", n.ClassGPUs, want)
	}
	// Each device as people are shown it: the instance under its card, with
	// its own memory, and the card its sensors are read under.
	g := n.GPUs[1]
	if g.Label != "0/mig1" || g.Card != 0 || g.MIGProfile != "3g.40gb" || g.MemoryTotalMB != 40192 ||
		g.Class != "a100-80g-mig-3g.40gb" {
		t.Errorf("instance view: %+v", g)
	}
	if w := n.GPUs[2]; w.Label != "1" || w.Card != 1 || w.MIGProfile != "" {
		t.Errorf("whole card view: %+v", w)
	}

	// The fleet summary attributes each device to its own class, where it
	// once put every card of a node under the node's only class.
	s := a.Summary()
	got := map[string][2]int{}
	for _, c := range s.GPUs {
		got[c.Class] = [2]int{c.Total, c.Used}
	}
	if !reflect.DeepEqual(got, map[string][2]int{"a100-80g-mig-3g.40gb": {2, 1}, "a100-80g": {1, 0}}) {
		t.Errorf("summary by class: %v", got)
	}

	// And a plan sees both classes, each with its own count.
	pool := a.PoolNodes()
	if len(pool) != 1 || len(pool[0].Kinds) != 2 || pool[0].Kinds[0].GPUMemGB != 40 ||
		pool[0].Kinds[0].FreeGPUs != 1 || pool[0].Kinds[1].GPUMemGB != 80 {
		t.Errorf("pool: %+v", pool)
	}
}