package modelplan

import (
	"fmt"
	"sort"
)

// Hardware is what one kind of node offers, as the fleet reports it.
type Hardware struct {
	// Class is the GPU class, e.g. "h100-80g". Empty means a CPU-only node.
	Class string
	// GPUMemGB is one card's memory.
	GPUMemGB int
	// GPUsPerNode is how many cards a node of this kind has free.
	GPUsPerNode int
	// Nodes is how many such nodes are available.
	Nodes int
	// Fabric is what they are attached to, empty when they are not. A replica
	// may never span two fabrics, so nodes on different ones are separate
	// pools however identical their cards.
	Fabric string
	// HostMemGB is the machine's RAM, which is what a CPU profile needs.
	HostMemGB int
}

// Fit is one way a model could run.
type Fit struct {
	Precision   Precision
	Class       string
	TP          int
	Nodes       int
	GPUsPerNode int
	Fabric      string
	// Context is the token pool this fit was checked against.
	Context int
	// PerGPUGB is what each card must hold, and Headroom what is left of it.
	PerGPUGB float64
	Headroom float64
	// CPUOnly marks a fit that needs no GPU at all.
	CPUOnly bool
}

// usableFraction is how much of a card a plan may assume it can have.
//
// Not all of it. A framework's allocator fragments, the driver reserves some,
// and a card planned to ninety-nine percent full is a card that fails to load
// on a bad day rather than a good one. vLLM's own default GPU utilisation is
// 0.90 and this matches it deliberately: a plan that disagrees with the engine
// it is planning for is worse than no plan.
const usableFraction = 0.90

// Fits returns every way this model could run on the given hardware, best
// first — fewest GPUs, then fewest nodes, then highest precision.
//
// Every candidate, not one answer, because the choice between them is not the
// planner's to make: a team with spare cards wants the full-precision fit, and
// a team with one card free wants to know a quantised one exists. Returning
// only the "best" would hide the option someone actually wants.
func Fits(m Model, hw []Hardware, context int, precisions []Precision) []Fit {
	if context <= 0 {
		context = m.MaxContext
	}
	if len(precisions) == 0 {
		precisions = []Precision{BF16, FP8, Q8GGUF, Q4GGUF}
	}

	var out []Fit
	for _, p := range precisions {
		d := m.Needs(p, context)
		for _, h := range hw {
			out = append(out, fitsOn(m, h, d, p, context)...)
		}
	}

	out = dedupe(out)
	sort.SliceStable(out, func(i, j int) bool {
		a, b := out[i], out[j]
		if ga, gb := a.TP, b.TP; ga != gb {
			return ga < gb
		}
		if a.Nodes != b.Nodes {
			return a.Nodes < b.Nodes
		}
		// Higher precision first among equals: the same hardware serving a
		// better model is strictly the better answer.
		return a.Precision.BytesPerParam > b.Precision.BytesPerParam
	})
	return out
}

// dedupe drops fits that differ only by a fabric that does not matter.
//
// A cluster with two fabrics of identical cards produced every single-node fit
// twice, once per fabric, which is pure noise: a replica on one node does not
// span anything, so which fabric its node is attached to changes nothing about
// whether the model fits. Only a multi-node fit is fabric-specific, because
// that is the case where the fabric is the constraint.
//
// Found by running it against a real fleet: a 32B produced twenty-six rows of
// which thirteen were duplicates.
func dedupe(in []Fit) []Fit {
	type key struct {
		precision string
		class     string
		tp, nodes int
		fabric    string
	}
	seen := map[key]struct{}{}
	out := in[:0]
	for _, f := range in {
		k := key{f.Precision.Name, f.Class, f.TP, f.Nodes, ""}
		if f.Nodes > 1 {
			k.fabric = f.Fabric
		}
		if _, dup := seen[k]; dup {
			continue
		}
		seen[k] = struct{}{}
		out = append(out, f)
	}
	return out
}

func fitsOn(m Model, h Hardware, d Demand, p Precision, context int) []Fit {
	if h.Class == "" {
		return cpuFit(h, d, p, context)
	}

	var out []Fit
	budget := float64(h.GPUMemGB) * usableFraction

	// Powers of two only. Tensor parallelism divides attention heads across
	// ranks, and an engine that cannot divide them evenly refuses the profile —
	// so offering tp=3 would produce a plan that fails at launch.
	//
	// One wide on MIG instances, however many there are: a process opens
	// only one, and they have no NVLink between them.
	widest := h.GPUsPerNode * h.Nodes
	if isMIG(h.Class) {
		widest = min(widest, 1)
	}
	for width := 1; width <= widest; width *= 2 {
		nodes := 1
		perNode := width
		if width > h.GPUsPerNode {
			perNode = h.GPUsPerNode
			nodes = width / h.GPUsPerNode
		}
		if nodes > h.Nodes {
			break
		}
		// A replica spanning nodes needs a fabric. Without one every node is
		// its own island, and a multi-node replica would hang on its first
		// all-reduce rather than fail to start.
		if nodes > 1 && h.Fabric == "" {
			continue
		}
		// The KV heads have to divide too, and they are far fewer than the
		// attention heads under grouped-query attention: a model with 8 KV
		// heads cannot be split sixteen ways.
		if m.KVHeads > 0 && width > m.KVHeads {
			break
		}

		per := GB(d.PerGPU(width))
		if per > budget {
			continue
		}
		out = append(out, Fit{
			Precision: p, Class: h.Class, TP: width, Nodes: nodes,
			GPUsPerNode: perNode, Fabric: h.Fabric, Context: context,
			PerGPUGB: per, Headroom: budget - per,
		})
	}
	return out
}

// cpuFit checks a node with no cards, where the whole working set lives in RAM.
func cpuFit(h Hardware, d Demand, p Precision, context int) []Fit {
	if h.HostMemGB <= 0 {
		return nil
	}
	// Less of the machine than of a card: the operating system, the page cache
	// and everything else on the node want their share, and a CPU node is
	// rarely dedicated the way a GPU node is.
	budget := float64(h.HostMemGB) * 0.75
	need := GB(d.Total())
	if need > budget {
		return nil
	}
	return []Fit{{
		Precision: p, TP: 1, Nodes: 1, Context: context,
		PerGPUGB: need, Headroom: budget - need, CPUOnly: true,
	}}
}

// Describe renders a fit the way someone would say it out loud.
func (f Fit) Describe() string {
	if f.CPUOnly {
		return fmt.Sprintf("CPU, %s, %.1f GB of RAM", f.Precision.Name, f.PerGPUGB)
	}
	where := fmt.Sprintf("%d×%s", f.TP, f.Class)
	if f.Nodes > 1 {
		where = fmt.Sprintf("%d nodes × %d×%s on %s", f.Nodes, f.GPUsPerNode, f.Class, f.Fabric)
	}
	return fmt.Sprintf("%s, %s, %.1f GB a card (%.1f free)",
		where, f.Precision.Name, f.PerGPUGB, f.Headroom)
}