package inventory

import (
	"testing"
	"time"

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

func TestOperatorLabelsWinOverTheAgentsAndSurviveItsHeartbeats(t *testing.T) {
	s := NewStore()
	// Set before the node registers: labelling a rack's machines comes first.
	s.SetOperatorLabels("n1", map[string]string{"llmgw.io/rack": "r2"})
	s.Upsert(Node{ID: "n1", Labels: map[string]string{"llmgw.io/rack": "r1", "role": "gpu"}})

	n, _ := s.Get("n1")
	if n.Labels["llmgw.io/rack"] != "r2" || n.Labels["role"] != "gpu" {
		t.Fatalf("labels %v: the operator's rack over the agent's, the agent's role kept", n.Labels)
	}
	// A heartbeat hands back the node it read.
	s.Upsert(n)
	s.Upsert(n)
	if n, _ := s.Get("n1"); n.Labels["llmgw.io/rack"] != "r2" || n.AgentLabels["llmgw.io/rack"] != "r1" {
		t.Fatalf("after heartbeats: labels %v, agent's %v", n.Labels, n.AgentLabels)
	}

	s.SetOperatorLabels("n1", nil)
	if n, _ := s.Get("n1"); n.Labels["llmgw.io/rack"] != "r1" {
		t.Errorf("the operator's label removed, rack is %q; want the agent's r1 back", n.Labels["llmgw.io/rack"])
	}
}

func TestAnOperatorLabelOnANodeWithNoAgentLabelsIsNotTakenForTheAgents(t *testing.T) {
	s := NewStore()
	s.Upsert(Node{ID: "n1"})
	s.SetOperatorLabels("n1", map[string]string{"llmgw.io/rack": "r2"})
	n, _ := s.Get("n1")
	s.Upsert(n)
	s.SetOperatorLabels("n1", nil)
	if n, _ := s.Get("n1"); n.Labels["llmgw.io/rack"] != "" {
		t.Errorf("removed, but the node still carries rack %q", n.Labels["llmgw.io/rack"])
	}
}

func TestADomainsPowerIsTheLargerOfReservedAndMeasuredNodeByNode(t *testing.T) {
	s := NewStore()
	gpus := []discovery.GPU{{Index: 0}, {Index: 1}}
	rack := map[string]string{"llmgw.io/rack": "r1"}
	now := time.Now()
	// a: two cards placed at a 700 W limit, drawing 300 each: 1400 reserved.
	s.Upsert(Node{ID: "a", Labels: rack, GPUs: gpus})
	s.Upsert(Node{ID: "b", Labels: rack, GPUs: gpus})
	setLoad(s, "a", &Load{At: now, GPUs: []GPUSensor{{Index: 0, PowerLimitW: 700, PowerW: 300}, {Index: 1, PowerLimitW: 700, PowerW: 300}}})
	// b: nothing placed, but drawing 500 W.
	setLoad(s, "b", &Load{At: now, GPUs: []GPUSensor{{Index: 0, PowerW: 250}, {Index: 1, PowerW: 250}}})
	s.Allocate("a", Allocation{ShardID: "x", GPUIndices: []int{0, 1}})
	s.SetFacility(facility.Policy{NodeWatts: 100, Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 5000, OtherWatts: 50}}})

	u := s.PowerUse(now)["rack:r1"]
	if u.Reserved != 50+100+1400+100+0 {
		t.Errorf("reserved %v", u.Reserved)
	}
	if u.Measured != 50+100+600+100+500 {
		t.Errorf("measured %v", u.Measured)
	}
	if u.Used != 50+100+1400+100+500 {
		t.Errorf("used %v: a's reservation and b's reading", u.Used)
	}
	if len(u.Nodes) != 2 || u.Limit != 5000 {
		t.Errorf("nodes %v, limit %v", u.Nodes, u.Limit)
	}
}

func setLoad(s *Store, id string, l *Load) {
	n, _ := s.Get(id)
	n.Load = l
	s.Upsert(n)
}