package adminapi

import (
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
	"sort"
	"time"
)

// Summary is the whole fleet in one answer.
//
// A per-node listing is the right tool for four nodes and the wrong one for two
// hundred: the questions that matter at that size — how full am I, what is
// broken, what is holding cards it is not using — are aggregates, and reading
// them off a long table is a chore that gets skipped. Everything here is
// derived from the same state the tables show, so it can never disagree with
// them; it is a different question asked of the same facts.
type Summary struct {
	// Version is the release the answering control plane runs: not the
	// binary on its disk, which a manual install changes underneath it.
	Version string     `json:"version,omitempty"`
	Nodes   NodeTotals `json:"nodes"`
	GPUs    []ClassUse `json:"gpus"`
	// Fabrics matter because a multi-node replica may never span two of them,
	// so free capacity in the wrong fabric is not capacity for the model that
	// needs it.
	Fabrics  []FabricUse  `json:"fabrics,omitempty"`
	Replicas ReplicaTotal `json:"replicas"`
	Models   ModelTotals  `json:"models"`
	// Unplaceable names models the scheduler currently cannot satisfy. This is
	// the line that turns "something is wrong" into a place to look.
	Unplaceable []string `json:"unplaceable,omitempty"`
	// Parked are models that failed every time they started and are no
	// longer retried, with the last error.
	Parked []ParkedModel `json:"parked,omitempty"`
	// IdleReplicas are serving replicas nothing has asked for in a while, which
	// on a fleet with more models than GPUs is the first place to find room.
	IdleReplicas []IdleReplica `json:"idle_replicas,omitempty"`
	// Store is the members of the store's cluster when it is replicated, and
	// StoreError why it could not be read; both empty for a single database.
	Store      []StoreMember `json:"store,omitempty"`
	StoreError string        `json:"store_error,omitempty"`
	// ControlPlanes is the cluster's control planes when there is more than
	// one: which leads, which stand by, and the floating address in front.
	ControlPlanes ControlPlanes `json:"control_planes,omitzero"`
}

// StoreMember is one member of a replicated store, as its Patroni reports it.
type StoreMember struct {
	Name  string `json:"name"`
	Host  string `json:"host"`
	Role  string `json:"role"`
	State string `json:"state"`
	// Lag is bytes behind the leader; -1 unknown, 0 for the leader.
	Lag int64 `json:"lag"`
}

// ControlPlanes is every control plane of a cluster that has several, as
// each answered when asked.
type ControlPlanes struct {
	Members []ControlPlane `json:"members,omitempty"`
	// Floating is the address keepalived keeps on the leader (10.0.0.161),
	// and FloatingUp whether it answered as a leader when asked; empty when
	// nothing floats.
	Floating   string `json:"floating,omitempty"`
	FloatingUp bool   `json:"floating_up,omitempty"`
}

// ControlPlane is one control plane.
type ControlPlane struct {
	Name    string `json:"name"`
	Address string `json:"address"`
	// State is "leading", "standing by" or "not answering".
	State string `json:"state"`
	// This is the control plane that answered the summary.
	This bool `json:"this,omitempty"`
}

// NodeTotals counts nodes by the states that need different responses.
type NodeTotals struct {
	Total       int `json:"total"`
	Ready       int `json:"ready"`
	Cordoned    int `json:"cordoned"`
	Draining    int `json:"draining"`
	Unreachable int `json:"unreachable"`
}

// ClassUse is how much of one kind of card is in use.
type ClassUse struct {
	Class string `json:"class"`
	Total int    `json:"total"`
	Used  int    `json:"used"`
	// Unhealthy GPUs are neither used nor available, and counting them as free
	// is how a cluster looks like it has room it does not have.
	Unhealthy int `json:"unhealthy"`
}

// Free is what can still be allocated.
func (c ClassUse) Free() int { return c.Total - c.Used - c.Unhealthy }

// FabricUse is capacity within one fabric.
type FabricUse struct {
	Fabric string `json:"fabric"`
	Nodes  int    `json:"nodes"`
	// ModelsPresent is how many nodes here hold weights at all, and Models how
	// many distinct models they hold between them.
	//
	// Per fabric because that is the unit at which the answer is actionable.
	// Fabrics are not connected to one another, so weights in one are not
	// weights in another: a model seeded into a single fabric leaves the rest
	// of the cluster unable to serve it without first moving tens of gigabytes
	// over the management network, and nothing said so.
	NodesWithWeights int `json:"nodes_with_weights"`
	Models           int `json:"models_present"`
	// LargestFreeBlock is the most nodes on this fabric that are entirely free,
	// which is the number that decides whether a multi-node replica can land.
	LargestFreeBlock int `json:"largest_free_block"`
}

// ReplicaTotal counts replicas by state.
type ReplicaTotal struct {
	Ready    int `json:"ready"`
	Starting int `json:"starting"`
	Draining int `json:"draining"`
	Failed   int `json:"failed"`
}

// ModelTotals counts models by whether they have what they asked for.
type ModelTotals struct {
	Total     int `json:"total"`
	Satisfied int `json:"satisfied"`
	Short     int `json:"short"`
}

// IdleReplica is a replica holding hardware nothing is asking for.
type IdleReplica struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	IdleSec int64  `json:"idle_seconds"`
	GPUs    int    `json:"gpus"`
}

// idleThreshold is how long a replica must go unasked-for to be worth naming.
//
// Minutes rather than seconds, because a replica that flaps in and out of this
// list is noise, and because unloading one throws away its KV cache — the thing
// prefix routing spent its effort filling.
const idleThreshold = 10 * time.Minute

// Summary builds the fleet summary.
func (a *API) Summary() Summary {
	s := Summary{Version: version.Version}
	if a.storeStatus != nil {
		members, err := a.storeStatus()
		s.Store = members
		if err != nil {
			s.StoreError = err.Error()
		}
	}
	if a.controlPlanes != nil {
		s.ControlPlanes = a.controlPlanes()
	}

	nodes := a.Nodes()
	byClass := map[string]*ClassUse{}
	s.Nodes.Total = len(nodes)
	for _, n := range nodes {
		switch n.State {
		case "ready":
			s.Nodes.Ready++
		case "cordoned":
			s.Nodes.Cordoned++
		case "draining":
			s.Nodes.Draining++
		default:
			s.Nodes.Unreachable++
		}
		for _, c := range n.Classes {
			if byClass[c] == nil {
				byClass[c] = &ClassUse{Class: c}
			}
		}
		// Each class its own devices, which a node with a card in MIG mode
		// needs: its instances are one class or several, and any whole card
		// beside them another.
		for _, c := range n.ClassGPUs {
			u := byClass[c.Class]
			u.Total += c.Total
			u.Unhealthy += c.Unhealthy
			u.Used += c.Total - c.Free - c.Unhealthy
		}
	}
	for _, u := range byClass {
		s.GPUs = append(s.GPUs, *u)
	}
	sort.Slice(s.GPUs, func(i, j int) bool { return s.GPUs[i].Class < s.GPUs[j].Class })

	byFabric := map[string]*FabricUse{}
	for _, n := range nodes {
		for _, f := range n.Fabrics {
			if byFabric[f] == nil {
				byFabric[f] = &FabricUse{Fabric: f}
			}
			byFabric[f].Nodes++
			if len(n.Models) > 0 {
				byFabric[f].NodesWithWeights++
			}
			// Entirely free, which is what a multi-node replica needs: a node
			// with one card in use cannot contribute a full-width shard.
			if n.State == "ready" && n.GPUFree == n.GPUTotal && n.GPUTotal > 0 {
				byFabric[f].LargestFreeBlock++
			}
		}
	}
	// Distinct models per fabric, counted across its nodes: one copy on a
	// shared filer and twenty local copies are the same answer to "can this
	// fabric serve it".
	perFabricModels := map[string]map[string]struct{}{}
	for _, n := range nodes {
		for _, f := range n.Fabrics {
			if perFabricModels[f] == nil {
				perFabricModels[f] = map[string]struct{}{}
			}
			for _, m := range n.Models {
				perFabricModels[f][m] = struct{}{}
			}
		}
	}
	for name, f := range byFabric {
		f.Models = len(perFabricModels[name])
		s.Fabrics = append(s.Fabrics, *f)
	}
	sort.Slice(s.Fabrics, func(i, j int) bool { return s.Fabrics[i].Fabric < s.Fabrics[j].Fabric })

	for _, r := range a.Replicas("") {
		switch r.State {
		case "ready":
			s.Replicas.Ready++
		case "draining":
			s.Replicas.Draining++
		case "failed":
			s.Replicas.Failed++
		default:
			s.Replicas.Starting++
		}
		if r.State == "ready" && r.IdleSec >= int64(idleThreshold.Seconds()) {
			gpus := 0
			for _, sh := range r.Shards {
				gpus += len(sh.GPUs)
			}
			s.IdleReplicas = append(s.IdleReplicas, IdleReplica{
				ID: r.ID, Model: r.Model, IdleSec: r.IdleSec, GPUs: gpus,
			})
		}
	}
	sort.Slice(s.IdleReplicas, func(i, j int) bool {
		return s.IdleReplicas[i].IdleSec > s.IdleReplicas[j].IdleSec
	})

	for _, m := range a.Models() {
		s.Models.Total++
		if m.Parked {
			s.Parked = append(s.Parked, ParkedModel{Model: m.Name, Failures: m.Failures, Reason: m.Reason})
		}
		// Against what it should have now, not its minimum: an on-demand
		// model's minimum is 0, so one that was asked for and fitted nowhere
		// used to count as satisfied.
		if m.Ready >= m.Desired && !m.Parked {
			s.Models.Satisfied++
			continue
		}
		s.Models.Short++
		// Short of replicas and unable to place more is the actionable case;
		// short and still starting is just a cluster doing its job.
		if m.Unplaceable || (m.Total <= m.Ready && !m.Parked) {
			s.Unplaceable = append(s.Unplaceable, m.Name)
		}
	}
	sort.Strings(s.Unplaceable)
	return s
}

// ParkedModel is a model no longer retried, and why.
type ParkedModel struct {
	Model    string `json:"model"`
	Failures int    `json:"failures"`
	Reason   string `json:"reason"`
}