package autogen

import "math"

// kvBytesPerElem is the approximate bytes-per-element of each KV cache quant
// (includes block-scale overhead).
var kvBytesPerElem = map[string]float64{
	"f32":    4.0,
	"f16":    2.0,
	"bf16":   2.0,
	"q8_0":   1.0625,
	"q5_1":   0.75,
	"q5_0":   0.6875,
	"q4_1":   0.625,
	"q4_0":   0.5625,
	"iq4_nl": 0.5625,
}

// kvQuantAllowed is the set of KV cache types quartermaster will emit for
// -ctk/-ctv. It is llama.cpp's own list (common/arg.cpp `kv_cache_types`) minus
// iq4_nl: every emitted command uses flash-attention, which has no iq4_nl KV
// kernel.
var kvQuantAllowed = map[string]bool{
	"f32": true, "f16": true, "bf16": true,
	"q8_0": true, "q5_1": true, "q5_0": true, "q4_1": true, "q4_0": true,
}

// ValidKvPair reports whether a K/V quant pair is emittable: both sides must be
// a known type and must match (flash-attention requires matched K/V).
func ValidKvPair(kvK, kvV string) bool {
	return kvK == kvV && kvQuantAllowed[kvK]
}

// resolveKvPair layers a possibly HALF-specified K/V pair over the defaults.
// Fast flash attention needs matched K and V, so ValidKvPair demands K==V — but
// treating a one-sided pin as invalid dropped BOTH sides back to the default,
// which is the opposite of what the caller asked for. Mirror the given side
// instead, and fall back only when neither side is set or the pair is a genuine
// mismatch (two different quants, or an unsupported type).
//
// This is not cosmetic: the fleet default is f16 and a pin is usually the
// cheaper q8_0, so a dropped half-pin silently doubles the KV reserve. On
// Qwen3.8-27B at 64k that is the difference between -ngl 99 / 0 GB RAM and
// -ngl 64 / 0.33 GB — the editor's estimate bar reported a layer spilled to RAM
// for a config that emits and loads fully on the GPU.
func resolveKvPair(kvK, kvV, defK, defV string) (string, string) {
	if kvK == "" {
		kvK = kvV
	}
	if kvV == "" {
		kvV = kvK
	}
	if !ValidKvPair(kvK, kvV) {
		return defK, defV
	}
	return kvK, kvV
}

// defaultKvQuant picks the fleet-wide default KV cache type for an LLM when no
// per-model override is set. Quality-first: f16 is the baseline for EVERY arch,
// not just MoE — a quantized KV's damage shows up in long-context recall and
// multi-turn tool use long before it shows in perplexity, and quartermaster's
// whole point is that placement is computed, so KV precision shouldn't be the
// thing silently traded away. q8_0 is used only as a step-down, when f16 cannot
// reach DenseMinCtx inside the VRAM budget (there, a shrunken window is the worse
// loss). settings.kvQuant pins one type fleet-wide and skips the decision.
//
// bf16 is deliberately NOT the auto pick despite the same 2 bytes: it trades
// mantissa bits for exponent range that K/V activations don't need, and f16 is
// the native flash-attention path.
func defaultKvQuant(s Settings, meta Metadata, targetVramGB, overheadGB float64) string {
	if kvQuantAllowed[s.KvQuant] {
		return s.KvQuant
	}
	m := GetKvCostModel(meta, "f16", "f16")
	if !m.OK {
		return "f16"
	}
	minCtx := s.DenseMinCtx
	if minCtx <= 0 {
		minCtx = 32768
	}
	if meta.ContextLength > 0 && int(meta.ContextLength) < minCtx {
		minCtx = int(meta.ContextLength)
	}
	if MaxCtxForBudget(targetVramGB-meta.FileSizeGB-overheadGB, m.SlopeGB, m.ConstGB) >= minCtx {
		return "f16"
	}
	return "q8_0"
}

func kvByteWidth(quant string) float64 {
	if b, ok := kvBytesPerElem[quant]; ok {
		return b
	}
	return 1.0625
}

// KvCostModel expresses KV cache size as Slope*ctx + Const (GB). For non-SWA,
// non-hybrid archs Const is 0 (pure linear). GlobalLayers/LocalLayers/SsmLayers
// are diagnostics.
type KvCostModel struct {
	SlopeGB      float64 // GB per token
	ConstGB      float64 // GB, ctx-independent (SWA windows + SSM recurrent state)
	GlobalLayers int
	LocalLayers  int
	SsmLayers    int
	OK           bool // false when core dims are missing
}

// GetKvCostModel derives the KV cost from gguf attention dims and the chosen K/V
// quant. Port of Get-KvCostModel. Returns OK=false (caller falls back to a flat
// reserve) when core dims are missing.
func GetKvCostModel(meta Metadata, kvK, kvV string) KvCostModel {
	blocks := int(meta.BlockCount)
	kvHeads := int(meta.HeadCountKv)
	kLen := int(meta.KeyLength)
	vLen := int(meta.ValueLength)
	if blocks <= 0 || kvHeads <= 0 || kLen <= 0 || vLen <= 0 {
		return KvCostModel{OK: false}
	}
	bK := kvByteWidth(kvK)
	bV := kvByteWidth(kvV)

	win := int(meta.SlidingWindow)
	pattern := 6
	if meta.SlidingWinPattern > 0 {
		pattern = int(meta.SlidingWinPattern)
	}
	kLenL := kLen
	if meta.KeyLengthSwa > 0 {
		kLenL = int(meta.KeyLengthSwa)
	}
	vLenL := vLen
	if meta.ValueLengthSwa > 0 {
		vLenL = int(meta.ValueLengthSwa)
	}
	arr := meta.HeadCountKvArr // nil -> uniform kvHeads on every layer

	perTokGlobal := float64(kLen)*bK + float64(vLen)*bV
	perWinLocal := (float64(kLenL)*bK + float64(vLenL)*bV) * float64(win)

	// Hybrid SSM: full-attention (KV) layer only every Nth; the rest are linear
	// with a fixed recurrent state (ctx-independent constant).
	interval := 0
	if meta.FullAttnInterval > 0 {
		interval = int(meta.FullAttnInterval)
	}
	ssmInner := int(meta.SsmInnerSize)
	ssmConv := int(meta.SsmConvKernel)
	ssmState := int(meta.SsmStateSize)

	var slopeBytes, constBytes float64
	var globalLayers, localLayers, ssmLayers int
	for i := 0; i < blocks; i++ {
		if interval > 0 && ((i+1)%interval) != 0 {
			ssmLayers++
			continue
		}
		kvh := kvHeads
		if arr != nil && i < len(arr) {
			kvh = int(arr[i])
		}
		if kvh <= 0 {
			continue // conv / no-KV layer
		}
		isGlobal := win <= 0 || ((i+1)%pattern) == 0
		if isGlobal {
			slopeBytes += float64(kvh) * perTokGlobal
			globalLayers++
		} else {
			constBytes += float64(kvh) * perWinLocal
			localLayers++
		}
	}
	// Fixed recurrent state for SSM layers (stored f32, 4 bytes), ctx-independent.
	if ssmLayers > 0 && ssmInner > 0 {
		recElems := (ssmInner * ssmState) + (ssmInner * int(math.Max(0, float64(ssmConv-1))))
		constBytes += float64(ssmLayers) * float64(recElems) * 4.0
	}

	return KvCostModel{
		SlopeGB:      slopeBytes / gib,
		ConstGB:      constBytes / gib,
		GlobalLayers: globalLayers,
		LocalLayers:  localLayers,
		SsmLayers:    ssmLayers,
		OK:           true,
	}
}

// mtpDraftSlopeGB is the per-token VRAM cost (GB) of the BAKED-IN MTP drafter's
// own KV cache — i.e. what --spec-type draft-mtp reserves on top of the main
// model's when the gguf carries its own nextn layers and no separate -md file is
// attached.
//
// llama-server does not run the nextn head inside the main context. It creates a
// second llama_context over the SAME model
// (server-context.cpp, "creating MTP draft context against the target model"),
// inheriting the target's cparams — n_ctx included — and filtering its KV cache
// to the nextn layers only (llama-model.cpp: `il >= hparams.n_layer()`; on
// hybrid qwen35 that cache is plain attention, not the hybrid wrapper, so there
// is no second copy of the recurrent state). So the drafter costs
// nextn_layers worth of full-attention KV over the WHOLE context window, which
// is why the old flat 0.34 GB charge was wrong in both directions: ~3x too fat at
// a 32k tier and ~2x too thin at 160k.
//
// kvK/kvV are the DRAFT cache types (-ctkd/-ctvd), not the main ones: the draft
// context takes params.speculative.draft.cache_type_k, which llama defaults to
// f16 regardless of -ctk. See draftKvPair.
func mtpDraftSlopeGB(meta Metadata, kvK, kvV string) float64 {
	nextn := meta.NextnLayers
	kvHeads := meta.HeadCountKv
	if nextn <= 0 || kvHeads <= 0 || meta.KeyLength <= 0 || meta.ValueLength <= 0 {
		return 0
	}
	perTok := float64(kvHeads) * (float64(meta.KeyLength)*kvByteWidth(kvK) + float64(meta.ValueLength)*kvByteWidth(kvV))
	return float64(nextn) * perTok / gib
}

// MaxCtxForBudget returns the largest ctx that fits a VRAM/RAM budget given the
// slope+const KV model. 0 when nothing fits.
func MaxCtxForBudget(budgetGB, slopeGB, constGB float64) int {
	avail := budgetGB - constGB
	if avail <= 0 || slopeGB <= 0 {
		return 0
	}
	return int(math.Floor(avail / slopeGB))
}

// KvReserveGB is the KV cache size (GB) for a chosen ctx.
func KvReserveGB(ctx int, slopeGB, constGB float64) float64 {
	return slopeGB*float64(ctx) + constGB
}

// RoundedCtx rounds ctx down to a multiple of 4096 (clean rope/cache
// boundaries), floored at 4096.
func RoundedCtx(ctx float64) int {
	r := int(math.Floor(ctx/4096)) * 4096
	if r < 4096 {
		return 4096
	}
	return r
}

// DenseCtxResult is the chosen dense context plus a human note.
type DenseCtxResult struct {
	Ctx  int
	Note string
}

// DenseCtxParams are inputs to GetDenseCtx.
type DenseCtxParams struct {
	ModelMax     int
	PerTokGB     float64
	KvConstGB    float64
	FileSizeGB   float64
	TargetVramGB float64
	Overhead     float64
	Ladder       []int
	MinCtx       int
	AllowOffload bool
}

// GetDenseCtx picks a dense model's context, speed-first: never offload layers
// just to grow context unless the model can't reach MinCtx on GPU (or the caller
// opts in via AllowOffload). Port of Get-DenseCtx.
func GetDenseCtx(p DenseCtxParams) DenseCtxResult {
	kvAt99 := p.TargetVramGB - p.FileSizeGB - p.Overhead
	if kvAt99 < 0.1 {
		kvAt99 = 0.1
	}
	ctxAt99 := MaxCtxForBudget(kvAt99, p.PerTokGB, p.KvConstGB)
	// The ladder's top rung is a ceiling, not a target: with no ladder at all the
	// only ceiling is the model's own trained window.
	top := p.ModelMax
	if len(p.Ladder) > 0 {
		top = min(p.ModelMax, p.Ladder[0])
	}

	// 1. Reaches the floor fully on GPU -> max ctx up to the ladder top. Fast.
	if ctxAt99 >= p.MinCtx {
		c := RoundedCtx(float64(min(min(p.ModelMax, top), ctxAt99)))
		note := "max-ctx (full gpu)"
		switch {
		case c >= p.ModelMax:
			note = "model-max"
		case c >= top:
			note = "ladder-top (full gpu)"
		}
		return DenseCtxResult{Ctx: c, Note: note}
	}

	weightsFitGpu := p.FileSizeGB <= (p.TargetVramGB - p.Overhead - 0.1)

	// 3 / manual opt-in: offload to reach the target window.
	if p.AllowOffload || !weightsFitGpu {
		c := RoundedCtx(float64(min(p.ModelMax, p.MinCtx)))
		return DenseCtxResult{Ctx: c, Note: "vram-limited (offload for ctx)"}
	}

	// 2. Weights fit but KV can't reach MinCtx on GPU -> stay full GPU, smaller ctx.
	c := RoundedCtx(float64(min(p.ModelMax, ctxAt99)))
	return DenseCtxResult{Ctx: c, Note: "max-ctx (full gpu, under min)"}
}