package placement

import (
	"fmt"
	"math/rand/v2"
	"strings"
	"testing"
	"time"

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

// at puts a node in a rack, and on PDUs when given.
func at(n inventory.Node, rack string, pdus ...string) inventory.Node {
	if n.Labels == nil {
		n.Labels = map[string]string{}
	}
	n.Labels[facility.Rack.Label()] = rack
	if len(pdus) > 0 {
		n.Labels[facility.PDU.Label()] = strings.Join(pdus, ",")
	}
	return n
}

// withPowerLimit has a node's cards report a power limit, and draw, as a
// heartbeat would.
func withPowerLimit(n inventory.Node, limit, draw float64) inventory.Node {
	l := &inventory.Load{At: time.Now()}
	for _, g := range n.GPUs {
		l.GPUs = append(l.GPUs, inventory.GPUSensor{Index: g.Index, PowerLimitW: limit, PowerW: draw})
	}
	n.Load = l
	return n
}

func onePerNode(name string, gpus int) Profile {
	return Profile{Name: name, GPUClass: "h100-80g", GPUsPerNode: gpus, Nodes: 1}
}

func TestNothingIsPlacedInAHeldRackAndTheRejectionSaysWhy(t *testing.T) {
	st := store(at(node("a1", "h100-80g", 8), "r1"), at(node("b1", "h100-80g", 8), "r2"))
	st.SetFacility(facility.Policy{Holds: map[facility.Domain]facility.Hold{
		"rack:r2": {Domain: "rack:r2", Reason: "PDU A replacement"},
	}})
	st.Allocate("a1", inventory.Allocation{ShardID: "x", GPUIndices: []int{0, 1, 2, 3, 4, 5, 6, 7}})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil)
	if res.Placement != nil {
		t.Fatalf("placed on %s, in a held rack", res.Placement.Nodes[0].NodeID)
	}
	if !strings.Contains(res.Explain(), "b1: rack r2 is held: PDU A replacement") {
		t.Errorf("the rejection does not name the hold:\n%s", res.Explain())
	}
}

func TestAHeldPDUHoldsEveryRackOnIt(t *testing.T) {
	st := store(at(node("a1", "h100-80g", 8), "r1", "p1", "p2"), at(node("b1", "h100-80g", 8), "r2", "p3"))
	st.SetFacility(facility.Policy{Holds: map[facility.Domain]facility.Hold{"pdu:p2": {Domain: "pdu:p2"}}})

	res := New(st).Place([]Profile{onePerNode("tp8", 8)}, nil)
	if res.Placement == nil || res.Placement.Nodes[0].NodeID != "b1" {
		t.Fatalf("want b1, the node on no held PDU:\n%s", res.Explain())
	}
}

func TestASecondReplicaGoesToAnotherRackBeforeAFullerNodeInTheSameOne(t *testing.T) {
	// Bin packing alone would choose a2, the partly used node beside the
	// first replica; a rack apart is worth an emptier node.
	st := store(
		at(node("a1", "h100-80g", 8), "r1"),
		at(node("a2", "h100-80g", 8), "r1"),
		at(node("b1", "h100-80g", 8), "r2"),
	)
	st.Allocate("a2", inventory.Allocation{ShardID: "other", GPUIndices: []int{0, 1, 2, 3}})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, []string{"a1"})
	if res.Placement == nil {
		t.Fatalf("not placed:\n%s", res.Explain())
	}
	if got := res.Placement.Nodes[0].NodeID; got != "b1" {
		t.Errorf("placed on %s, in the rack of the first replica; want b1 in another", got)
	}
}

func TestReplicasAreSpreadAcrossPDUsAsWellAsRacks(t *testing.T) {
	st := store(
		at(node("a1", "h100-80g", 8), "r1", "p1"),
		at(node("b1", "h100-80g", 8), "r2", "p1"),
		at(node("c1", "h100-80g", 8), "r3", "p2"),
	)
	st.Allocate("b1", inventory.Allocation{ShardID: "other", GPUIndices: []int{0, 1, 2, 3}})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, []string{"a1"})
	if res.Placement == nil || res.Placement.Nodes[0].NodeID != "c1" {
		t.Fatalf("want c1, on another rack and another PDU:\n%s", res.Explain())
	}
}

func TestSpreadingAcrossRacksNeverLeavesAReplicaUnplaced(t *testing.T) {
	st := store(at(node("a1", "h100-80g", 8), "r1"), at(node("a2", "h100-80g", 8), "r1"))

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, []string{"a1"})
	if res.Placement == nil || res.Placement.Nodes[0].NodeID != "a2" {
		t.Fatalf("the rack's other node should still be used:\n%s", res.Explain())
	}
}

func TestNodesWithoutALocationPlaceAsTheyAlwaysDid(t *testing.T) {
	// Bin packing picks the partly used node; nothing about racks changes it.
	st := store(node("n1", "h100-80g", 8), node("n2", "h100-80g", 8))
	st.Allocate("n2", inventory.Allocation{ShardID: "other", GPUIndices: []int{0, 1}})
	st.SetFacility(facility.Policy{GPUWatts: 700, Margin: 0.1})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil)
	if res.Placement == nil || res.Placement.Nodes[0].NodeID != "n2" {
		t.Fatalf("want n2, the bin-packed choice:\n%s", res.Explain())
	}
}

func TestAReplicaThatWouldTakeARackPastItsCapGoesToAnotherRack(t *testing.T) {
	// Rack r1: 8 cards at 700 W, four of them placed: 2.8 kW used of a
	// 4.5 kW limit (5 kW less 10%). Four more need 2.8 kW, which does not fit.
	st := store(
		withPowerLimit(at(node("a1", "h100-80g", 8), "r1"), 700, 0),
		withPowerLimit(at(node("b1", "h100-80g", 8), "r2"), 700, 0),
	)
	st.Allocate("a1", inventory.Allocation{ShardID: "other", GPUIndices: []int{0, 1, 2, 3}})
	st.SetFacility(facility.Policy{Margin: 0.1, Budgets: map[facility.Domain]facility.Budget{
		"rack:r1": {CapWatts: 5000},
	}})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil)
	if res.Placement == nil || res.Placement.Nodes[0].NodeID != "b1" {
		t.Fatalf("want b1, in the rack with room:\n%s", res.Explain())
	}
	want := "a1: rack r1: 2.8 kW of a 4.5 kW limit (5.0 kW cap less 10%) in use, the shard needs ~2.8 kW"
	if !strings.Contains(res.Explain(), want) {
		t.Errorf("the rejection should give the figures, %q:\n%s", want, res.Explain())
	}
}

func TestAProfilesDeclaredDrawIsWhatIsCounted(t *testing.T) {
	st := store(withPowerLimit(at(node("a1", "h100-80g", 8), "r1"), 700, 0))
	st.SetFacility(facility.Policy{Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 2000}}})

	derived := onePerNode("tp4", 4) // 4 × 700 W: too much for 2 kW
	if res := New(st).Place([]Profile{derived}, nil); res.Placement != nil {
		t.Fatalf("2.8 kW was placed under a 2 kW cap")
	}
	declared := derived
	declared.PowerWatts = 1500
	res := New(st).Place([]Profile{declared}, nil)
	if res.Placement == nil {
		t.Fatalf("a profile declaring 1.5 kW was refused under a 2 kW cap:\n%s", res.Explain())
	}
	if res.Placement.Nodes[0].PowerW != 1500 {
		t.Errorf("the shard carries %v W to record, want the declared 1500", res.Placement.Nodes[0].PowerW)
	}
}

func TestWhatTheSensorsReadCountsWhenItIsMoreThanWasReserved(t *testing.T) {
	// Nothing placed, but the cards draw 400 W each: 3.2 kW measured, and
	// another 2.8 kW would pass a 5 kW cap only on the reservations.
	st := store(withPowerLimit(at(node("a1", "h100-80g", 8), "r1"), 700, 400))
	st.SetFacility(facility.Policy{Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 5000}}})

	res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil)
	if res.Placement != nil {
		t.Fatalf("placed against reservations alone, with 3.2 kW measured")
	}
	if !strings.Contains(res.Explain(), "3.2 kW of a 5.0 kW cap in use") {
		t.Errorf("the measured draw should be what is in use:\n%s", res.Explain())
	}
}

func TestAStaleReadingIsNotCounted(t *testing.T) {
	n := withPowerLimit(at(node("a1", "h100-80g", 8), "r1"), 700, 400)
	n.Load.At = time.Now().Add(-10 * time.Minute)
	st := store(n)
	st.SetFacility(facility.Policy{Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 5000}}})

	if res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil); res.Placement == nil {
		t.Fatalf("a ten-minute-old reading kept a replica out:\n%s", res.Explain())
	}
}

func TestADrawNobodyKnowsNeverBlocksPlacement(t *testing.T) {
	// A cap, but no card reports a limit and no figure is configured: the
	// shard's draw is unknown, which is not a reason to refuse it.
	st := store(at(node("a1", "h100-80g", 8), "r1"))
	st.SetFacility(facility.Policy{Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 100}}})

	if res := New(st).Place([]Profile{onePerNode("tp8", 8)}, nil); res.Placement == nil {
		t.Fatalf("an unknown draw was refused:\n%s", res.Explain())
	}
}

func TestTheOtherEquipmentInARackCountsAgainstItsCap(t *testing.T) {
	st := store(at(node("a1", "h100-80g", 8), "r1"))
	st.SetFacility(facility.Policy{GPUWatts: 500, Budgets: map[facility.Domain]facility.Budget{
		"rack:r1": {CapWatts: 5000, OtherWatts: 3500},
	}})

	if res := New(st).Place([]Profile{onePerNode("tp4", 4)}, nil); res.Placement != nil {
		t.Fatal("2 kW was placed where switches already take 3.5 of 5")
	}
}

func TestAMultiNodeReplicasShardsAreCountedTogetherAgainstTheirRack(t *testing.T) {
	// Each node alone fits r1's budget; two together do not, so the second
	// shard goes to the rack with room — on the same fabric, as always.
	st := store(
		withPowerLimit(at(node("a1", "h100-80g", 8, "ib"), "r1"), 500, 0),
		withPowerLimit(at(node("a2", "h100-80g", 8, "ib"), "r1"), 500, 0),
		withPowerLimit(at(node("b1", "h100-80g", 8, "ib"), "r2"), 500, 0),
	)
	st.SetFacility(facility.Policy{Budgets: map[facility.Domain]facility.Budget{"rack:r1": {CapWatts: 6000}}})

	res := New(st).Place([]Profile{{Name: "2x8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2}}, nil)
	if res.Placement == nil {
		t.Fatalf("not placed:\n%s", res.Explain())
	}
	racks := map[string]int{}
	for _, a := range res.Placement.Nodes {
		n, _ := st.Get(a.NodeID)
		racks[n.Location().Rack]++
	}
	if racks["r1"] != 1 || racks["r2"] != 1 {
		t.Errorf("shards by rack %v: two in r1 draw 8 kW against a 6 kW cap", racks)
	}
	if !strings.Contains(res.Explain(), "with this replica's other shards") {
		t.Errorf("the node passed over should say why:\n%s", res.Explain())
	}
}

func TestAProfileKeptWithinARackNeverSpansTwo(t *testing.T) {
	// The emptiest-first order would take a1 and b1; the profile wants one
	// rack, and only r2 has two nodes.
	st := store(
		at(node("a1", "h100-80g", 8, "ib"), "r1"),
		at(node("b1", "h100-80g", 8, "ib"), "r2"),
		at(node("b2", "h100-80g", 8, "ib"), "r2"),
		node("x1", "h100-80g", 8, "ib"),
	)
	p := Profile{Name: "2x8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2, KeepWithin: facility.Rack}

	res := New(st).Place([]Profile{p}, nil)
	if res.Placement == nil {
		t.Fatalf("not placed:\n%s", res.Explain())
	}
	for _, a := range res.Placement.Nodes {
		if a.NodeID != "b1" && a.NodeID != "b2" {
			t.Errorf("a shard on %s, outside rack r2", a.NodeID)
		}
	}
	explain := res.Explain()
	if !strings.Contains(explain, "x1: has no llmgw.io/rack label, and the profile keeps its nodes within one rack") {
		t.Errorf("the unlabelled node should be named:\n%s", explain)
	}
	if !strings.Contains(explain, "needs 2 nodes within one rack, rack r1 has 1") {
		t.Errorf("the rack too small should be named:\n%s", explain)
	}
}

func TestAProfileKeptWithinARackIsNotPlacedAcrossRacks(t *testing.T) {
	st := store(at(node("a1", "h100-80g", 8, "ib"), "r1"), at(node("b1", "h100-80g", 8, "ib"), "r2"))
	p := Profile{Name: "2x8", GPUClass: "h100-80g", GPUsPerNode: 8, Nodes: 2, KeepWithin: facility.Rack}
	if res := New(st).Place([]Profile{p}, nil); res.Placement != nil {
		t.Fatal("a replica kept within one rack was placed across two")
	}
}

// TestNoCapIsEverExceeded fills random clusters until nothing more fits, and
// checks every capped domain stayed within its limit — and that a cluster
// whose draw nobody knows is filled as if it had no caps at all.
func TestNoCapIsEverExceeded(t *testing.T) {
	rng := rand.New(rand.NewPCG(3, 4))
	for iter := range 200 {
		st := inventory.NewStore()
		known := rng.IntN(4) != 0
		cards := 0
		for i := range 2 + rng.IntN(6) {
			n := at(node(fmt.Sprintf("n%d", i), "h100-80g", 1+rng.IntN(8), "ib"),
				fmt.Sprintf("r%d", rng.IntN(3)), fmt.Sprintf("p%d", rng.IntN(2)))
			if known {
				n = withPowerLimit(n, float64(300+rng.IntN(500)), 0)
			}
			cards += len(n.GPUs)
			st.Upsert(n)
		}
		pol := facility.Policy{Margin: float64(rng.IntN(20)) / 100, Budgets: map[facility.Domain]facility.Budget{}}
		for _, d := range []facility.Domain{"rack:r0", "rack:r1", "rack:r2", "pdu:p0", "pdu:p1"} {
			if rng.IntN(2) == 0 {
				pol.Budgets[d] = facility.Budget{CapWatts: float64(1000 + rng.IntN(8000))}
			}
		}
		st.SetFacility(pol)

		for k := range 60 {
			p := Profile{Name: "gen", GPUClass: "h100-80g", GPUsPerNode: 1, Nodes: 1 + rng.IntN(2)}
			res := New(st).Place([]Profile{p}, nil)
			if res.Placement == nil {
				continue
			}
			for _, a := range res.Placement.Nodes {
				st.Allocate(a.NodeID, inventory.Allocation{ShardID: fmt.Sprintf("s%d-%d", k, a.Rank), GPUIndices: a.GPUIndices})
			}
		}
		for d, u := range st.PowerUse(time.Now()) {
			if u.Limit > 0 && u.Used > u.Limit+1e-6 {
				t.Fatalf("iteration %d: %s uses %.0f W of a %.0f W limit", iter, d, u.Used, u.Limit)
			}
		}
		if !known {
			free := 0
			for _, n := range st.Nodes() {
				free += len(st.FreeGPUs(n.ID))
			}
			if free != 0 {
				t.Fatalf("iteration %d: %d of %d cards left free with no draw known", iter, free, cards)
			}
		}
	}
}