// Package facility is where nodes are in the building, and what the building
// allows: the racks, rows, rooms and power distribution units a node belongs
// to, which of them an operator has taken out of service, and how much power
// each may draw.
//
// Placement reads it to keep replicas of one model off a single rack or PDU,
// to keep a replica's shards inside one domain when its profile asks, to place
// nothing in a domain held for maintenance, and to refuse a replica that would
// take a domain past its power cap. Everything here is optional: a node with
// no location labels and a cluster with no facility configuration place
// exactly as they did before.
package facility

import (
	"fmt"
	"slices"
	"strings"
	"time"
)

// Level is a kind of failure domain, from the largest to the smallest.
type Level string

const (
	Site Level = "site"
	Room Level = "room"
	Row  Level = "row"
	Rack Level = "rack"
	// PDU is a power distribution unit. A node may be on several, one per
	// feed of a dual-fed rack, so its label takes a comma separated list.
	PDU Level = "pdu"
)

// Levels are every level, the building's tree first and power after it.
var Levels = []Level{Site, Room, Row, Rack, PDU}

// LabelPrefix starts the node labels that say where a node is.
const LabelPrefix = "llmgw.io/"

// Label is the node label a level is read from: llmgw.io/rack.
func (l Level) Label() string { return LabelPrefix + string(l) }

// ParseLevel reads a level's name, as a profile or a command gives it.
func ParseLevel(s string) (Level, error) {
	l := Level(strings.ToLower(strings.TrimSpace(s)))
	if slices.Contains(Levels, l) {
		return l, nil
	}
	return "", fmt.Errorf("%q is not a facility level: use site, room, row, rack or pdu", s)
}

// Domain names one failure domain as level:name, rack:r12-07. Names are taken
// as the labels give them, so a rack is one domain across the whole cluster:
// two rooms that both have a rack r1 must label them apart.
type Domain string

// Of builds a domain from its level and name.
func Of(l Level, name string) Domain { return Domain(string(l) + ":" + name) }

// Level is the domain's level.
func (d Domain) Level() Level {
	l, _, _ := strings.Cut(string(d), ":")
	return Level(l)
}

// Name is the domain's name within its level.
func (d Domain) Name() string {
	_, n, _ := strings.Cut(string(d), ":")
	return n
}

// String is level:name; human text says "rack r12-07".
func (d Domain) String() string { return string(d) }

// Human is the domain as a sentence says it: rack r12-07.
func (d Domain) Human() string { return string(d.Level()) + " " + d.Name() }

// ParseDomain reads level:name, the way commands and configuration name a
// domain.
func ParseDomain(s string) (Domain, error) {
	l, n, ok := strings.Cut(strings.TrimSpace(s), ":")
	if !ok || strings.TrimSpace(n) == "" {
		return "", fmt.Errorf("%q is not a domain: name one as level:name, like rack:r12-07", s)
	}
	lv, err := ParseLevel(l)
	if err != nil {
		return "", err
	}
	return Of(lv, strings.TrimSpace(n)), nil
}

// Location is where a node is, read from its labels. Empty fields are not
// known, which places as it always did.
type Location struct {
	Site string   `json:"site,omitempty"`
	Room string   `json:"room,omitempty"`
	Row  string   `json:"row,omitempty"`
	Rack string   `json:"rack,omitempty"`
	PDUs []string `json:"pdus,omitempty"`
}

// LocationOf reads a node's labels.
func LocationOf(labels map[string]string) Location {
	get := func(l Level) string { return strings.TrimSpace(labels[l.Label()]) }
	loc := Location{Site: get(Site), Room: get(Room), Row: get(Row), Rack: get(Rack)}
	for _, p := range strings.Split(get(PDU), ",") {
		if p = strings.TrimSpace(p); p != "" && !slices.Contains(loc.PDUs, p) {
			loc.PDUs = append(loc.PDUs, p)
		}
	}
	return loc
}

// Empty reports whether nothing is known about where the node is.
func (l Location) Empty() bool {
	return l.Site == "" && l.Room == "" && l.Row == "" && l.Rack == "" && len(l.PDUs) == 0
}

// Domains are every domain the location is in, largest first.
func (l Location) Domains() []Domain {
	var out []Domain
	add := func(lv Level, name string) {
		if name != "" {
			out = append(out, Of(lv, name))
		}
	}
	add(Site, l.Site)
	add(Room, l.Room)
	add(Row, l.Row)
	add(Rack, l.Rack)
	for _, p := range l.PDUs {
		add(PDU, p)
	}
	return out
}

// At is the domain at one level, or empty. For PDUs it is the first.
func (l Location) At(lv Level) Domain {
	for _, d := range l.Domains() {
		if d.Level() == lv {
			return d
		}
	}
	return ""
}

// Short is the location in a few characters for a table: the rack, else the
// row, room or site, with the PDUs after a slash.
func (l Location) Short() string {
	s := ""
	switch {
	case l.Rack != "":
		s = l.Rack
	case l.Row != "":
		s = "row " + l.Row
	case l.Room != "":
		s = l.Room
	case l.Site != "":
		s = l.Site
	}
	if len(l.PDUs) > 0 {
		if s != "" {
			s += " "
		}
		s += "pdu " + strings.Join(l.PDUs, ",")
	}
	return s
}

// Hold takes a domain out of placement: maintenance, a cooling fault, work on
// a PDU. What runs there keeps running unless Drain is set.
type Hold struct {
	Domain Domain    `json:"domain"`
	Reason string    `json:"reason,omitempty"`
	Drain  bool      `json:"drain,omitempty"`
	By     string    `json:"by,omitempty"`
	At     time.Time `json:"at"`
}

// Budget is a domain's power cap.
type Budget struct {
	// CapWatts is what the domain may draw. For a dual-fed rack it is what
	// one feed can carry alone, since that is what the rack must survive.
	CapWatts float64
	// OtherWatts is drawn in the domain by things llmgw does not place on:
	// switches, storage, nodes outside the cluster.
	OtherWatts float64
}

// Policy is the facility's rules as placement reads them. The zero value
// holds nothing and caps nothing.
type Policy struct {
	// Holds are the domains taken out of placement.
	Holds map[Domain]Hold
	// Budgets are the domains with a power cap.
	Budgets map[Domain]Budget
	// Margin is the share of every cap kept free, 0.1 for ten percent.
	Margin float64
	// GPUWatts is a card's draw when its own power limit is not known.
	GPUWatts float64
	// CPUCoreWatts is what a core reserved by a CPU replica draws.
	CPUCoreWatts float64
	// NodeWatts is what a node draws beyond the GPUs and the cores given to
	// replicas: its idle draw, fans, memory and disks. Counted for every node
	// in a capped domain, busy or not.
	NodeWatts float64
	// MeasuredFresh is how old a node's power reading may be and still be
	// counted. Zero is two minutes.
	MeasuredFresh time.Duration
}

// Clone is a copy whose maps are its own.
func (p Policy) Clone() Policy {
	c := p
	if p.Holds != nil {
		c.Holds = make(map[Domain]Hold, len(p.Holds))
		for k, v := range p.Holds {
			c.Holds[k] = v
		}
	}
	if p.Budgets != nil {
		c.Budgets = make(map[Domain]Budget, len(p.Budgets))
		for k, v := range p.Budgets {
			c.Budgets[k] = v
		}
	}
	return c
}

// HeldIn is the first hold among a node's domains, largest first.
func (p Policy) HeldIn(domains []Domain) (Hold, bool) {
	for _, d := range domains {
		if h, ok := p.Holds[d]; ok {
			return h, true
		}
	}
	return Hold{}, false
}

// Limit is the most a domain may draw, its cap less the margin; zero with
// no cap.
func (p Policy) Limit(d Domain) float64 {
	b, ok := p.Budgets[d]
	if !ok || b.CapWatts <= 0 {
		return 0
	}
	m := p.Margin
	if m < 0 || m >= 1 {
		m = 0
	}
	return b.CapWatts * (1 - m)
}

// Fresh is how old a power reading may be.
func (p Policy) Fresh() time.Duration {
	if p.MeasuredFresh > 0 {
		return p.MeasuredFresh
	}
	return 2 * time.Minute
}

// KW renders watts as a sentence about a rack would.
func KW(w float64) string {
	if w >= 1000 {
		return fmt.Sprintf("%.1f kW", w/1000)
	}
	return fmt.Sprintf("%.0f W", w)
}