package modelplan

import (
	"fmt"
	"regexp"
	"sort"
	"strings"
)

// Options: the ways the files a repository actually has could run on the
// cluster as it is — each file on one card, on several of a node's cards, on
// several nodes' cards over a fabric, or on a CPU node — with what each card
// would hold, and whether the cards are free now or only exist.

// Artifact is something a repository offers to serve: one GGUF file (or one
// split into parts), or its original safetensors weights.
type Artifact struct {
	// Name is the file, or the set's name; for original weights, "original".
	Name   string `json:"name"`
	Format string `json:"format"` // "gguf" or "safetensors"
	Bytes  int64  `json:"bytes"`
	Shards int    `json:"shards,omitempty"`
	// Quant is the quantisation read from the name, like Q4_K_M; "bf16" for
	// original weights.
	Quant string `json:"quant,omitempty"`
}

// Shape is what sizes a model's KV cache, and constrains how it splits.
type Shape struct {
	// KVBytesPerToken is one token of context, the whole model.
	KVBytesPerToken int64
	// KVHeads must divide by a tensor-parallel width; 0 is unknown.
	KVHeads int
	// Params is the parameter count, for bits per weight; 0 is unknown.
	Params int64
}

// PoolNode is one ready node, as a plan sees it.
type PoolNode struct {
	ID       string
	Class    string // "" for a node with no GPU
	Fabric   string
	GPUMemGB int
	GPUs     int
	FreeGPUs int
	// HostMemGB is its RAM, and HostFreeGB what is unused now.
	HostMemGB  int
	HostFreeGB int
	// Kinds are the node's devices by class, for a node with more than one:
	// a card in MIG mode is one or several classes of instance, and a whole
	// card beside it another. Empty means Class, GPUMemGB, GPUs and FreeGPUs
	// say it all.
	Kinds []PoolKind
}

// PoolKind is one class of a node's devices.
type PoolKind struct {
	Class    string
	GPUMemGB int
	GPUs     int
	FreeGPUs int
}

// kinds is the node's devices by class.
func (n PoolNode) kinds() []PoolKind {
	if len(n.Kinds) > 0 {
		return n.Kinds
	}
	return []PoolKind{{Class: n.Class, GPUMemGB: n.GPUMemGB, GPUs: n.GPUs, FreeGPUs: n.FreeGPUs}}
}

// isMIG reports whether a class names MIG instances: "a100-80g-mig-3g.40gb".
// A replica runs on exactly one, never a group of them, so a plan offers one
// instance and nothing wider. Read from the class label, as the card's memory
// is, which keeps this package free of the agent's.
func isMIG(class string) bool { return strings.Contains(class, "-mig-") }

// Option is one way to run one artifact.
type Option struct {
	Artifact string `json:"artifact"`
	Format   string `json:"format"`
	Engine   string `json:"engine"`
	CPU      bool   `json:"cpu,omitempty"`
	Class    string `json:"gpu_class,omitempty"`
	Fabric   string `json:"fabric,omitempty"`
	// GPUsPerNode cards on each of Nodes nodes.
	GPUsPerNode int `json:"gpus_per_node,omitempty"`
	Nodes       int `json:"nodes,omitempty"`
	// NeedGB is what one card holds, or for CPU the whole replica in RAM;
	// HaveGB what a card, or the node, offers for it.
	NeedGB float64 `json:"need_gb"`
	HaveGB float64 `json:"have_gb"`
	// Hosts is how many nodes could run it at all, and FreeNow whether one
	// (or enough, together) could right now.
	Hosts   int  `json:"hosts"`
	FreeNow bool `json:"free_now"`
	// BitsPerWeight is the artifact's precision, 0 when unknown.
	BitsPerWeight float64 `json:"bits_per_weight,omitempty"`
}

// GPUs is every card one replica takes.
func (o Option) GPUs() int { return o.GPUsPerNode * max(o.Nodes, 1) }

// Label says where it runs, in a few words.
func (o Option) Label() string {
	if o.CPU {
		return "CPU"
	}
	if o.Nodes > 1 {
		return fmt.Sprintf("%d nodes × %d %s on %s", o.Nodes, o.GPUsPerNode, o.Class, o.Fabric)
	}
	return fmt.Sprintf("%d × %s", o.GPUsPerNode, o.Class)
}

// gpuShare is how much of a card a plan may count on, as in Fits.
const gpuShare = usableFraction

// hostShare is how much of a node's RAM one replica may be planned to use:
// the system and the agent keep the rest.
const hostShare = 0.85

// ggufOverhead is llama.cpp's own cost beyond weights and cache, as the GGUF
// header's MemoryBytes counts it: compute buffers and the server.
func ggufOverhead(weights int64) int64 { return 1<<30 + weights/20 }

// Options lists every way each artifact could run, best first.
//
// context is the token pool: llama.cpp's --ctx-size, or vLLM's
// --max-model-len for one sequence.
func Options(arts []Artifact, sh Shape, nodes []PoolNode, context int) []Option {
	var out []Option
	for _, a := range arts {
		bpw := quantBits[strings.ToUpper(a.Quant)]
		if sh.Params > 0 {
			bpw = float64(a.Bytes) * 8 / float64(sh.Params)
		}
		kv := int64(context) * sh.KVBytesPerToken
		switch a.Format {
		case "gguf":
			out = append(out, ggufOptions(a, kv, nodes, bpw)...)
		case "safetensors":
			out = append(out, vllmOptions(a, kv, sh, nodes, bpw)...)
		}
	}
	Rank(out)
	return out
}

func ggufOptions(a Artifact, kv int64, nodes []PoolNode, bpw float64) []Option {
	var out []Option
	// The CPU: the whole replica in a node's RAM.
	total := a.Bytes + kv + ggufOverhead(a.Bytes)
	cpu := Option{Artifact: a.Name, Format: a.Format, Engine: "llamacpp", CPU: true,
		NeedGB: GB(total), BitsPerWeight: bpw}
	for _, n := range nodes {
		have := float64(n.HostMemGB) * hostShare
		if n.HostMemGB <= 0 || GB(total) > have {
			continue
		}
		cpu.Hosts++
		cpu.HaveGB = max(cpu.HaveGB, have)
		if float64(n.HostFreeGB) >= GB(total) {
			cpu.FreeNow = true
		}
	}
	if cpu.Hosts > 0 {
		out = append(out, cpu)
	}
	// GPUs: one node's cards, the model split across them by layer.
	for _, c := range classes(nodes) {
		budget := float64(c.mem) * gpuShare
		for width := 1; width <= c.maxGPUs; width++ {
			per := (a.Bytes+kv)/int64(width) + ggufOverhead(a.Bytes/int64(width))
			if GB(per) > budget {
				continue
			}
			o := Option{Artifact: a.Name, Format: a.Format, Engine: "llamacpp", Class: c.class,
				GPUsPerNode: width, Nodes: 1, NeedGB: GB(per), HaveGB: budget, BitsPerWeight: bpw}
			for _, fab := range c.fabrics {
				for _, n := range fab.nodes {
					if n.GPUs >= width {
						o.Hosts++
						if n.FreeGPUs >= width {
							o.FreeNow = true
						}
					}
				}
			}
			out = append(out, o)
			break // more cards than the fewest that hold it buys nothing here
		}
	}
	return out
}

func vllmOptions(a Artifact, kv int64, sh Shape, nodes []PoolNode, bpw float64) []Option {
	var out []Option
	perCard := func(width int) int64 {
		return (a.Bytes+kv)/int64(width) + overheadBytes(a.Bytes/int64(width))
	}
	// Tensor parallel within a node must divide the KV heads.
	divides := func(w int) bool { return sh.KVHeads == 0 || (w <= sh.KVHeads && sh.KVHeads%w == 0) }
	for _, c := range classes(nodes) {
		budget := float64(c.mem) * gpuShare
		opt := func(perNode, n int, fabric string) Option {
			return Option{Artifact: a.Name, Format: a.Format, Engine: "vllm", Class: c.class, Fabric: fabric,
				GPUsPerNode: perNode, Nodes: n, NeedGB: GB(perCard(perNode * n)), HaveGB: budget, BitsPerWeight: bpw}
		}
		// One node: the narrowest power of two that holds it.
		single := false
		for w := 1; w <= c.maxGPUs && divides(w); w *= 2 {
			if GB(perCard(w)) > budget {
				continue
			}
			o := opt(w, 1, "")
			for _, fab := range c.fabrics {
				for _, nd := range fab.nodes {
					if nd.GPUs >= w {
						o.Hosts++
						o.FreeNow = o.FreeNow || nd.FreeGPUs >= w
					}
				}
			}
			out = append(out, o)
			single = true
			break
		}
		// A MIG instance does not join others, on its node or across a
		// fabric: what does not fit one does not fit the class.
		if single || isMIG(c.class) {
			continue
		}
		// Several nodes, a node's full width on each, pipeline-parallel
		// between them: only within one fabric.
		w := 1
		for w*2 <= c.maxGPUs {
			w *= 2
		}
		if !divides(w) {
			continue
		}
		for _, fab := range c.fabrics {
			if fab.name == "" {
				continue
			}
			for n := 2; n <= len(fab.nodes); n *= 2 {
				if GB(perCard(w*n)) > budget {
					continue
				}
				o := opt(w, n, fab.name)
				free := 0
				for _, nd := range fab.nodes {
					if nd.GPUs >= w {
						o.Hosts++
					}
					if nd.FreeGPUs >= w {
						free++
					}
				}
				if o.Hosts < n {
					break
				}
				o.FreeNow = free >= n
				out = append(out, o)
				break
			}
		}
	}
	return out
}

type classPool struct {
	class   string
	mem     int
	maxGPUs int
	fabrics []fabricPool
}

type fabricPool struct {
	name  string
	nodes []PoolNode
}

// classes groups GPU nodes by card, and each by fabric: a replica spans
// nodes only within one. A node with devices of several classes is in each
// class's pool, as the node it is for that class: its devices of that class
// and no others.
func classes(nodes []PoolNode) []classPool {
	by := map[string]*classPool{}
	var order []string
	for _, node := range nodes {
		for _, k := range node.kinds() {
			if k.Class == "" || k.GPUs == 0 || k.GPUMemGB == 0 {
				continue
			}
			n := node
			n.Class, n.GPUMemGB, n.GPUs, n.FreeGPUs, n.Kinds = k.Class, k.GPUMemGB, k.GPUs, k.FreeGPUs, nil
			if isMIG(n.Class) {
				// One instance a replica, however many the node has.
				n.GPUs, n.FreeGPUs = min(n.GPUs, 1), min(n.FreeGPUs, 1)
			}
			c := by[n.Class]
			if c == nil {
				c = &classPool{class: n.Class, mem: n.GPUMemGB}
				by[n.Class] = c
				order = append(order, n.Class)
			}
			c.maxGPUs = max(c.maxGPUs, n.GPUs)
			i := -1
			for j := range c.fabrics {
				if c.fabrics[j].name == n.Fabric {
					i = j
				}
			}
			if i < 0 {
				c.fabrics = append(c.fabrics, fabricPool{name: n.Fabric})
				i = len(c.fabrics) - 1
			}
			c.fabrics[i].nodes = append(c.fabrics[i].nodes, n)
		}
	}
	sort.Strings(order)
	out := make([]classPool, 0, len(order))
	for _, k := range order {
		out = append(out, *by[k])
	}
	return out
}

// quantBits is roughly what each quantisation spends a weight, for when the
// parameter count is not known to work it out from the file's size.
var quantBits = map[string]float64{
	"F32": 32, "F16": 16, "BF16": 16, "FP16": 16, "Q8_0": 8.5, "Q6_K": 6.6,
	"Q5_K_M": 5.7, "Q5_K_S": 5.5, "Q5_0": 5.5, "Q5_1": 6, "Q4_K_M": 4.8, "Q4_K_S": 4.6, "Q4_0": 4.5, "Q4_1": 5,
	"IQ4_XS": 4.3, "IQ4_NL": 4.5, "Q3_K_L": 4.3, "Q3_K_M": 3.9, "Q3_K_S": 3.5, "IQ3_M": 3.7, "IQ3_XS": 3.3,
	"IQ3_XXS": 3.1, "Q2_K": 3.0, "IQ2_M": 2.7, "IQ2_XS": 2.3, "IQ2_XXS": 2.1, "IQ1_M": 1.8, "IQ1_S": 1.6,
}

// qualityCap is the precision past which a bigger file is not a better
// model worth its memory: eight bits keep what sixteen do, near enough. On a
// CPU the cap is lower: there every bit is read from memory for every token,
// so a file twice the size answers at half the speed, and a Q4_K_M or Q5 is
// the usual balance.
const (
	qualityCap    = 8.6
	cpuQualityCap = 5.8
)

// quality scores an option's precision: rising to the cap, falling past it.
func quality(o Option) float64 {
	if o.BitsPerWeight == 0 {
		return 0
	}
	limit := qualityCap
	if o.CPU {
		limit = cpuQualityCap
	}
	if o.BitsPerWeight <= limit {
		return o.BitsPerWeight
	}
	return limit - (o.BitsPerWeight-limit)*0.25
}

// Rank orders options best first by what the hardware is, not by what it is
// doing: on a GPU before the CPU; the fewest cards, then nodes; then the
// better precision, up to eight bits. Whether the cards are free now only
// breaks a tie, then the most room left.
//
// Busy is not unfit. Ranking what was free first made every file of every
// repository come out on the CPU on a cluster whose one card was in use —
// which turned down the very file the card was there for. A model served onto
// cards that are busy waits for one, or with its fallbacks runs elsewhere
// until then.
func Rank(opts []Option) {
	sort.SliceStable(opts, func(i, j int) bool {
		a, b := opts[i], opts[j]
		if a.CPU != b.CPU {
			return !a.CPU
		}
		if a.GPUs() != b.GPUs() {
			return a.GPUs() < b.GPUs()
		}
		if a.Nodes != b.Nodes {
			return a.Nodes < b.Nodes
		}
		if qa, qb := quality(a), quality(b); qa != qb {
			return qa > qb
		}
		if a.FreeNow != b.FreeNow {
			return a.FreeNow
		}
		return a.HaveGB-a.NeedGB > b.HaveGB-b.NeedGB
	})
}

var quantPattern = regexp.MustCompile(`(?i)(IQ\d_[A-Z]+|Q\d_K_[SML]|Q\d_K|Q\d_\d|Q\d|BF16|F16|F32|FP16)`)

// QuantOf reads a quantisation out of a file name.
func QuantOf(name string) string {
	m := quantPattern.FindAllString(name, -1)
	if len(m) == 0 {
		return ""
	}
	return strings.ToUpper(m[len(m)-1])
}

// MemFromClass reads the card size out of a class label such as "h100-80g",
// and a MIG instance's out of its profile: 40 for "a100-80g-mig-3g.40gb",
// whose card is 80 but whose instance holds 40 of it.
func MemFromClass(class string) int {
	if _, profile, ok := strings.Cut(class, "-mig-"); ok {
		var compute, gb int
		if _, err := fmt.Sscanf(profile, "%dg.%dgb", &compute, &gb); err != nil {
			return 0
		}
		return gb
	}
	i := strings.LastIndex(class, "-")
	if i < 0 || !strings.HasSuffix(class, "g") {
		return 0
	}
	var gb int
	if _, err := fmt.Sscanf(class[i+1:], "%dg", &gb); err != nil {
		return 0
	}
	return gb
}