// Package modelplan works out whether a model fits some hardware, and how it
// should be split when it does.
//
// It exists because that arithmetic was the real barrier to running a model
// here. Everything else about adding one is typing: the name, the path, the
// engine. Deciding tp, gpus_per_node, gpu_class and the context that will
// actually fit means knowing the model's layer count, its KV head count, what
// its weights weigh at a given quantisation, and how much of a card is left
// after the runtime takes its share — and getting it wrong is not an error
// message, it is an out-of-memory twenty minutes into a load.
package modelplan

import "fmt"

// Model is what a model's own configuration says about it.
//
// Every field here is read from the model rather than guessed: config.json
// carries the architecture, and the repository's index carries the parameter
// count. A guess would defeat the point, because the answer is only worth
// having if it is the one the hardware will agree with.
type Model struct {
	Name string
	// Params is the total parameter count.
	Params int64
	// Layers, KVHeads and HeadDim size the KV cache, which is what scales with
	// context and is usually what actually runs a card out of memory.
	Layers  int
	KVHeads int
	HeadDim int
	// MaxContext is the model's trained context length.
	MaxContext int
}

// Precision is how many bytes one parameter weighs.
//
// Named rather than numeric because the names are what people use, and because
// the difference between them is the difference between a model fitting and
// not.
type Precision struct {
	Name           string
	BytesPerParam  float64
	BytesPerKVElem float64
	// Engine is the one that can serve this precision. Carried here because
	// choosing a precision and choosing an engine are the same decision, and
	// separating them produced a plan recommending a GGUF quantisation with a
	// vLLM profile — which would have loaded full-precision weights that the
	// same plan had just shown do not fit.
	Engine string
	// Quantization is the engine's own flag value, where it needs one.
	Quantization string
	// GGUF marks a precision that lives in a file rather than a repository, so
	// the source is a path and not a model id.
	GGUF bool
}

// Known precisions, from full to the quantisations worth serving.
//
// Engine matters as much as the number of bytes. q8_0 and q4_k_m are GGUF
// quantisations and only llama.cpp serves them, from a .gguf file rather than a
// repository; bf16 and fp8 are what vLLM and SGLang load from the weights as
// published. A plan that picks a precision without saying which engine can
// serve it is a plan that does not work — and worse, one that looks like it
// should.
var (
	BF16   = Precision{Name: "bf16", BytesPerParam: 2, BytesPerKVElem: 2, Engine: "vllm"}
	FP8    = Precision{Name: "fp8", BytesPerParam: 1, BytesPerKVElem: 1, Engine: "vllm", Quantization: "fp8"}
	Q8GGUF = Precision{Name: "q8_0", BytesPerParam: 1.06, BytesPerKVElem: 2, Engine: "llamacpp", GGUF: true}
	Q4GGUF = Precision{Name: "q4_k_m", BytesPerParam: 0.61, BytesPerKVElem: 2, Engine: "llamacpp", GGUF: true}
)

// PrecisionByName resolves a precision, reporting what it knows when it cannot.
func PrecisionByName(s string) (Precision, error) {
	for _, p := range []Precision{BF16, FP8, Q8GGUF, Q4GGUF} {
		if p.Name == s {
			return p, nil
		}
	}
	return Precision{}, fmt.Errorf("unknown precision %q; known: bf16, fp8, q8_0, q4_k_m", s)
}

// Demand is what one replica needs, in bytes.
type Demand struct {
	Weights int64
	KVCache int64
	// Overhead is what the runtime costs on a single GPU. Across a tensor
	// parallel group it is recomputed from each rank's share — see PerGPU.
	Overhead int64
}

// Total is everything one replica needs across all its GPUs.
func (d Demand) Total() int64 { return d.Weights + d.KVCache + d.Overhead }

// PerGPU is what one GPU of a replica must hold.
//
// Tensor parallelism splits the weights and the KV cache, so both divide. The
// overhead does not divide — each process pays its own CUDA context and
// allocator — but neither is it constant: activations scale with the shard the
// rank actually holds, so it is recomputed from the divided weights.
//
// Charging the whole model's overhead to every rank was the first version of
// this, and it was badly wrong in exactly the case that matters. Llama-3.3-70B
// at bf16 across two 80 GB cards came out at 85.9 GB a card and was reported as
// not fitting, which is a configuration people run every day.
func (d Demand) PerGPU(gpus int) int64 {
	if gpus <= 0 {
		gpus = 1
	}
	n := int64(gpus)
	return (d.Weights+d.KVCache)/n + overheadBytes(d.Weights/n)
}

// KVBytesPerToken is what one token of context costs, for the whole model.
//
// Two for the key and the value, times the layers, times the heads the cache
// actually keeps. Grouped-query attention is why KVHeads and not the attention
// head count: a 14B model with 40 attention heads and 8 KV heads pays for
// eight, and the difference is a factor of five.
func (m Model) KVBytesPerToken(p Precision) int64 {
	return int64(2 * float64(m.Layers) * float64(m.KVHeads) * float64(m.HeadDim) * p.BytesPerKVElem)
}

// Needs computes what one replica of this model needs.
//
// context is the total token pool the engine will hold, which is the context
// length times however many requests it serves at once — the number an engine
// is configured with, not the number a single request sees.
func (m Model) Needs(p Precision, context int) Demand {
	weights := int64(float64(m.Params) * p.BytesPerParam)
	kv := int64(context) * m.KVBytesPerToken(p)
	return Demand{
		Weights:  weights,
		KVCache:  kv,
		Overhead: overheadBytes(weights),
	}
}

// overheadBytes is what the runtime costs beyond weights and cache, per GPU.
//
// A flat allowance for the CUDA context, the framework and its allocator —
// which a small model pays as much as a large one — plus a small fraction of
// the shard for activations, which do scale with it.
//
// The fraction is two percent rather than the ten it started at, because the
// ten was counted twice. usableFraction already holds back a tenth of every
// card for exactly this, so charging another tenth on top made a plan reject
// Llama-3.3-70B at bf16 across two 80 GB cards — a configuration people run
// every day. Reserving slack in two places, each unaware of the other, reads as
// caution and behaves as a bug.
func overheadBytes(weights int64) int64 {
	const flat = 2 << 30 // 2 GiB
	return flat + weights/50
}

// GB renders bytes the way people discuss hardware.
func GB(b int64) float64 { return float64(b) / (1 << 30) }