package inventory

import (
	"maps"
	"slices"
	"time"

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

// Location is where the node is, from its labels.
func (n Node) Location() facility.Location { return facility.LocationOf(n.Labels) }

// Domains are the failure domains the node is in, largest first.
func (n Node) Domains() []facility.Domain { return n.Location().Domains() }

// mergeLabels lays the operator's labels over the agent's. A copy, so
// neither is written through.
func mergeLabels(agent, operator map[string]string) map[string]string {
	if len(operator) == 0 {
		return maps.Clone(agent)
	}
	out := make(map[string]string, len(agent)+len(operator))
	maps.Copy(out, agent)
	maps.Copy(out, operator)
	return out
}

// SetFacility puts the facility's holds and power budgets in force for
// every placement from now on.
func (s *Store) SetFacility(p facility.Policy) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.facility = p.Clone()
}

// Facility is the policy in force; a copy, so a caller may change it.
func (s *Store) Facility() facility.Policy {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.facility.Clone()
}

// SetOperatorLabels replaces the labels set through the control plane for a
// node, which may not have registered yet: they are laid over its agent's
// when it does.
func (s *Store) SetOperatorLabels(id string, labels map[string]string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.operatorLabels == nil {
		s.operatorLabels = map[string]map[string]string{}
	}
	if len(labels) == 0 {
		delete(s.operatorLabels, id)
	} else {
		s.operatorLabels[id] = maps.Clone(labels)
	}
	if n, ok := s.nodes[id]; ok {
		n.Labels = mergeLabels(n.AgentLabels, s.operatorLabels[id])
	}
}

// OperatorLabels are the labels set through the control plane for a node.
func (s *Store) OperatorLabels(id string) map[string]string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return maps.Clone(s.operatorLabels[id])
}

// PowerUse is one domain's draw as placement counts it.
type PowerUse struct {
	Domain facility.Domain
	// Nodes are the domain's nodes, sorted.
	Nodes []string
	// Reserved is what the replicas placed in the domain are estimated to
	// draw, with each node's own draw and the domain's other equipment.
	Reserved float64
	// Measured is what the nodes' sensors read, with the same additions,
	// for the nodes that read anything; the others count what is reserved.
	Measured float64
	// Used is what placement counts against the cap: per node, the larger
	// of reserved and measured, so an estimate that was too low is
	// corrected by the room it stops offering rather than by a breaker.
	Used float64
	// Limit is the cap less the margin; zero is no cap.
	Limit  float64
	Budget facility.Budget
}

// PowerUse is every domain any node is in, or that has a budget, by domain.
func (s *Store) PowerUse(now time.Time) map[facility.Domain]PowerUse {
	s.mu.RLock()
	defer s.mu.RUnlock()
	pol := s.facility
	out := map[facility.Domain]PowerUse{}
	for d, b := range pol.Budgets {
		out[d] = PowerUse{Domain: d, Budget: b, Limit: pol.Limit(d),
			Reserved: b.OtherWatts, Measured: b.OtherWatts, Used: b.OtherWatts}
	}
	ids := slices.Sorted(maps.Keys(s.nodes))
	for _, id := range ids {
		n := s.nodes[id]
		reserved := 0.0
		for _, a := range s.allocs[id] {
			reserved += AllocationWatts(*n, a, pol)
		}
		measured, known := MeasuredWatts(*n, now, pol.Fresh())
		used := reserved
		if known && measured > used {
			used = measured
		}
		if !known {
			measured = reserved
		}
		for _, d := range n.Domains() {
			u, ok := out[d]
			if !ok {
				u = PowerUse{Domain: d}
			}
			u.Nodes = append(u.Nodes, id)
			u.Reserved += pol.NodeWatts + reserved
			u.Measured += pol.NodeWatts + measured
			u.Used += pol.NodeWatts + used
			out[d] = u
		}
	}
	return out
}

// AllocationWatts is what one shard is estimated to draw: what its profile
// declared, or else each of its cards at the power limit the card reports
// (the policy's figure for a card that reports none) and each of its cores
// at the policy's figure.
func AllocationWatts(n Node, a Allocation, pol facility.Policy) float64 {
	if a.PowerW > 0 {
		return a.PowerW
	}
	w := float64(a.CPUCores) * pol.CPUCoreWatts
	for _, i := range a.GPUIndices {
		w += GPUWatts(n, i, pol)
	}
	return w
}

// GPUWatts is one card's estimated draw: the power limit it reports, else
// the policy's figure for a card.
func GPUWatts(n Node, index int, pol facility.Policy) float64 {
	if n.Load != nil {
		for _, g := range n.Load.GPUs {
			if g.Index == index && g.PowerLimitW > 0 {
				return g.PowerLimitW
			}
		}
	}
	return pol.GPUWatts
}

// MeasuredWatts is what the node's sensors read at its last heartbeat: its
// cards and its CPU packages. Known is false when the reading is older than
// fresh, or the node has no power sensor at all, which a VM never has.
func MeasuredWatts(n Node, now time.Time, fresh time.Duration) (w float64, known bool) {
	if n.Load == nil || now.Sub(n.Load.At) > fresh {
		return 0, false
	}
	for _, g := range n.Load.GPUs {
		if g.PowerW > 0 {
			w += g.PowerW
			known = true
		}
	}
	if n.Load.CPUPowerW > 0 {
		w += n.Load.CPUPowerW
		known = true
	}
	return w, known
}