package placement

import (
	"fmt"
	"slices"
	"sort"
	"time"

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

// declaredShardWatts is the profile's declared draw for one node's shard.
func (p Profile) declaredShardWatts() float64 {
	if p.PowerWatts <= 0 {
		return 0
	}
	return p.PowerWatts / float64(max(p.Nodes, 1))
}

// shardWatts is what one shard of this profile is estimated to draw on a
// candidate: the declared figure, or its cards at the highest power limit
// among the node's free ones of the class (the cards are not chosen yet, and
// guessing low is the mistake that trips a breaker), or its cores.
func (s *Scheduler) shardWatts(p Profile, c candidate) float64 {
	if w := p.declaredShardWatts(); w > 0 {
		return w
	}
	pol := s.inv.Facility()
	if p.CPUOnly() {
		return float64(p.CPUCores) * pol.CPUCoreWatts
	}
	per := 0.0
	for _, g := range c.free {
		per = max(per, inventory.GPUWatts(c.node, g.Index, pol))
	}
	if per == 0 {
		per = pol.GPUWatts
	}
	return per * float64(p.GPUsPerNode)
}

// budget is the power budgets as one placement decision spends them: the
// domains' use when it began, and what the shards it has chosen so far add.
type budget struct {
	pol   facility.Policy
	use   map[facility.Domain]inventory.PowerUse
	added map[facility.Domain]float64
}

// newBudget reads the domains' use, once per decision. Without budgets there
// is nothing to read and every check passes.
func (s *Scheduler) newBudget() *budget {
	pol := s.inv.Facility()
	b := &budget{pol: pol, added: map[facility.Domain]float64{}}
	if len(pol.Budgets) > 0 {
		b.use = s.inv.PowerUse(s.clock())
	}
	return b
}

// check says why a shard drawing need watts on a node would take one of the
// node's domains past its limit; empty when it fits them all. A shard whose
// draw is not known (no declared figure, no card limit, no policy figure)
// needs nothing: missing data never blocks placement.
func (b *budget) check(n inventory.Node, need float64) string {
	if b.use == nil {
		return ""
	}
	for _, d := range n.Domains() {
		limit := b.pol.Limit(d)
		if limit <= 0 {
			continue
		}
		used := b.use[d].Used + b.added[d]
		if used+need > limit {
			capW := b.pol.Budgets[d].CapWatts
			of := facility.KW(limit) + " cap"
			if limit < capW {
				of = fmt.Sprintf("%s limit (%s cap less %.0f%%)", facility.KW(limit), facility.KW(capW), 100-100*limit/capW)
			}
			return fmt.Sprintf("%s: %s of a %s in use, the shard needs ~%s",
				d.Human(), facility.KW(used), of, facility.KW(need))
		}
	}
	return ""
}

// add counts a chosen shard against the node's domains.
func (b *budget) add(n inventory.Node, need float64) {
	for _, d := range n.Domains() {
		b.added[d] += need
	}
}

// clock is the scheduler's time, for the freshness of power readings.
func (s *Scheduler) clock() time.Time {
	if s.now != nil {
		return s.now()
	}
	return time.Now()
}

// heldReason is a hold as a rejection says it.
func heldReason(h facility.Hold) string {
	why := h.Domain.Human() + " is held"
	if h.Reason != "" {
		why += ": " + h.Reason
	}
	return why
}

// spreadLevels are the domains replicas of one model are kept apart in:
// the rack, which a top-of-rack switch or a door left open takes out whole,
// and the PDU, which a breaker does.
var spreadLevels = []facility.Level{facility.Rack, facility.PDU}

// spreadAcrossDomains moves to the front the candidates sharing the fewest
// racks and PDUs with the nodes the model already runs on, keeping the order
// given among equals.
//
// A preference and never a filter, as spreading across nodes is: a rack
// holding a model's only free cards still takes its second replica, which
// survives an engine crash if not a breaker. Nodes with no location share
// nothing and lose nothing, so a cluster without labels places as before.
func (s *Scheduler) spreadAcrossDomains(candidates []candidate, avoid []string) {
	if len(avoid) == 0 {
		return
	}
	taken := map[facility.Domain]bool{}
	for _, id := range avoid {
		if n, ok := s.inv.Get(id); ok {
			for _, d := range n.Domains() {
				if slices.Contains(spreadLevels, d.Level()) {
					taken[d] = true
				}
			}
		}
	}
	if len(taken) == 0 {
		return
	}
	shared := func(c candidate) int {
		k := 0
		for _, d := range c.node.Domains() {
			if taken[d] {
				k++
			}
		}
		return k
	}
	sort.SliceStable(candidates, func(i, j int) bool { return shared(candidates[i]) < shared(candidates[j]) })
}

// withinGroups divides a multi-node replica's candidates by the domain its
// profile keeps it within, largest first, and rejects the nodes that do not
// say which they are in. Without the constraint it is all of them, as one.
func withinGroups(p Profile, pool string, cands []candidate, rejections *[]Rejection) [][]candidate {
	if p.KeepWithin == "" {
		return [][]candidate{cands}
	}
	by := map[facility.Domain][]candidate{}
	for _, c := range cands {
		d := c.node.Location().At(p.KeepWithin)
		if d == "" {
			*rejections = append(*rejections, Rejection{Profile: p.Name, Pool: pool, Node: c.node.ID,
				Reason: fmt.Sprintf("has no %s label, and the profile keeps its nodes within one %s",
					p.KeepWithin.Label(), p.KeepWithin)})
			continue
		}
		by[d] = append(by[d], c)
	}
	keys := make([]facility.Domain, 0, len(by))
	for d := range by {
		keys = append(keys, d)
	}
	sort.Slice(keys, func(i, j int) bool {
		if len(by[keys[i]]) != len(by[keys[j]]) {
			return len(by[keys[i]]) > len(by[keys[j]])
		}
		return keys[i] < keys[j]
	})
	out := make([][]candidate, 0, len(keys))
	for _, d := range keys {
		if len(by[d]) < p.Nodes {
			*rejections = append(*rejections, Rejection{Profile: p.Name, Pool: pool,
				Reason: fmt.Sprintf("needs %d nodes within one %s, %s has %d", p.Nodes, p.KeepWithin, d.Human(), len(by[d]))})
			continue
		}
		out = append(out, by[d])
	}
	return out
}