package autogen

// estimate.go exposes the per-profile sizer as a one-shot preview so the web
// editor can show live VRAM/RAM use for a candidate tuning before it is saved.
// It mirrors the solo-profile path of emitModel without writing any config.

// EstimateInput is the subset of editor fields that affect placement/memory.
// Zero Ctx means "let the sizer pick". Zero TargetVramGB uses settings.
type EstimateInput struct {
	Ctx int
	// CtxExact marks Ctx as a custom launch pin (-c): the process runs that
	// window whatever the sizer would have picked, so it is not rounded up to a
	// 4096 multiple the way an auto pick is.
	CtxExact bool
	KvK      string
	KvV      string
	KvInRam  bool
	Spec     string
	// RopeScaling ("linear"/"yarn") lets Ctx exceed the model's trained length —
	// the preview must know, or it silently sizes a clamped window and reports a
	// KV reserve the real launch won't have.
	RopeScaling    string
	TargetVramGB   float64
	CpuOffload     int  // >0 pins layers offloaded to CPU, overriding the sizer
	CpuOffloadSet  bool // CpuOffload == 0 is a pin too (every layer on GPU)
	CtxCheckpoints *int // nil => llama default (32); 0 disables; reserves checkpoint VRAM
	// CpuDevice previews a `device: cpu` pin (--device none): the same ngl-0
	// footprint generate prices, all of it charged to RAM.
	CpuDevice bool
	// CheckpointMinStep pins -cms (checkpoint spacing in prompt tokens), which
	// scales each snapshot's global-KV term. 0 => the arch default (wide on SWA).
	CheckpointMinStep int
	DraftGB           float64 // separate MTP/DFlash draft gguf weights (GB); 0 => baked-in or none
	MmprojGB          float64 // "-vision" twin projector footprint (weights + CLIP reserve); 0 => none
	// Ub pins -ub (physical batch). The compute buffer scales with it — ~0.5 GB
	// between 512 and 1024 on a large-vocab model — so a preview that ignored a
	// per-model Override.Ub charged a buffer the real launch never has. 0 => the
	// same auto pick emit makes (effectiveUb).
	Ub int
	// KvKDraft/KvVDraft pin the draft context's cache type (-ctkd/-ctvd). Empty
	// means "whatever the emitter would default to", i.e. the main K/V quant. A
	// caller reconstructing a HAND-WRITTEN argv that carries -ctk but no -ctkd
	// must pass "f16" explicitly, because that is what llama uses there.
	KvKDraft string
	KvVDraft string
	// Parallel is llama-server's slot count (--parallel). Every slot gets its own
	// window over one shared pool, so the KV cost the sizer solves against is
	// Parallel x the per-token slope. 0/1 => a single slot.
	Parallel int
	// DraftKind is the paired draft sidecar's kind ("mtp"/"dflash", "" for none),
	// as DraftSidecarForDir reports it. Only consulted when Spec is empty, to
	// resolve the same auto spec the emitter picks — see EstimatePlan.
	DraftKind string
}

// EstimateResult is the previewed load plan for a candidate tuning.
type EstimateResult struct {
	Ctx          int     `json:"ctx"`
	Ngl          int     `json:"ngl"`
	NCpuMoe      int     `json:"nCpuMoe"`
	EstVramGB    float64 `json:"estVramGB"`
	EstRamGB     float64 `json:"estRamGB"`
	TargetVramGB float64 `json:"targetVramGB"`
	MaxRamGB     float64 `json:"maxRamGB"`
	KvReserveGB  float64 `json:"kvReserveGB"`
	// CheckpointGB is the VRAM reserved for context checkpoints (included in
	// EstVramGB via overhead, broken out so the UI can attribute it separately
	// from model weights).
	CheckpointGB float64 `json:"checkpointGB"`
	// DraftGB is the VRAM charged for the speculative draft / MTP nextn layer: a
	// baked-in nextn head costs its compute pad, its own compute GRAPH and its
	// OWN ctx-scaled KV cache (a second llama_context over the same model at the
	// same window), a separate draft gguf its weights + pad. Folded into
	// EstVramGB via overhead; broken out so the UI can attribute it separately
	// from the main model weights. 0 when no draft-mtp spec is active.
	DraftGB float64 `json:"draftGB"`
	// ComputeBufGB is the GPU compute buffer: the main context's graph (logits +
	// activations) plus the fixed per-process runtime constant for the backend in
	// use (see runtimeCtxGB). A baked-in MTP drafter's own graph is NOT in here,
	// it is charged to DraftGB. MmprojGB is a "-vision" twin's
	// projector weights + CLIP reserve (0 for non-vision). OverheadGB is the global
	// vramOverheadGB safety headroom. All three are folded into EstVramGB via
	// overhead; broken out so the UI can label them separately from the model
	// weights instead of lumping everything non-KV into "Weights".
	ComputeBufGB float64 `json:"computeBufGB"`
	MmprojGB     float64 `json:"mmprojGB"`
	OverheadGB   float64 `json:"overheadGB"`
	// FixedGB is the whole non-splittable share of EstVramGB: compute buffer,
	// CUDA context, projector, draft, safety headroom. On a --tensor-split load
	// these live on --main-gpu alone, so the spawn guard charges FixedGB to the
	// main device before it re-derives the ratio.
	FixedGB     float64 `json:"fixedGB"`
	RamExceeded bool    `json:"ramExceeded"`
	IsMoE       bool    `json:"isMoE"`
}

// EstimatePlan sizes one candidate tuning against the given settings + gguf
// metadata, returning the chosen ctx and the resulting VRAM/RAM estimate. It
// reuses sizeProfile/forceLowActiveMoE so the preview matches what a save would
// actually emit for the solo profile.
func EstimatePlan(s Settings, meta Metadata, in EstimateInput) (EstimateResult, error) {
	// KV quant: the fleet default (f16, stepping down to q8_0 only under VRAM
	// pressure) unless a valid matched override is given — mirrors emitModel, so
	// the preview's KV reserve matches what a save would emit.
	// An empty spec means AUTO, not "no drafter" — the emitter runs every model
	// through effectiveSpec, which hands an MTP-capable gguf draft-mtp+ngram-mod
	// and charges its drafter (the pad plus its ctx-scaled KV, or the sidecar's
	// real weights). The preview used to read "" as no spec and charge 0, so a
	// model left on auto spec previewed a whole drafter lighter than it bakes. On a tight budget that is the whole margin:
	// Qwen3.8-27B at 106k previewed -ngl 99 / 0 GB RAM and emitted -ngl 64 / 0.32.
	spec := in.Spec
	if spec == "" {
		spec = effectiveSpec(meta, nil, in.DraftKind)
	}
	estTarget := s.TargetVramGB
	if in.TargetVramGB > 0 {
		estTarget = in.TargetVramGB
	}
	// Only charge the sidecar when the active spec will actually attach it as -md
	// (see matchedDraftSizeGB): a dir can pair a DFlash drafter to a model that
	// runs on a baked-in MTP head, and that drafter never loads.
	draftGB := matchedDraftSizeGB(spec, in.DraftKind, in.DraftGB)
	kvDef := defaultKvQuant(s, meta, estTarget, s.VramOverheadGB+draftOverheadGB(spec, draftGB))
	// The editor sends kvK/kvV independently, so a form that resolves only one
	// side must not sink the preview back to the fleet default — see
	// resolveKvPair.
	kvK, kvV := resolveKvPair(in.KvK, in.KvV, kvDef, kvDef)

	perTokGB, kvConstGB := 0.0, 0.0
	if m := GetKvCostModel(meta, kvK, kvV); m.OK {
		perTokGB, kvConstGB = m.SlopeGB, m.ConstGB
	}
	// N slots share one -c pool, each with its own KV, so charge N x. buildCmdLines
	// multiplies the per-slot window back out for -c; charging one slot here
	// under-sized the KV reserve and the admission figure by N.
	slots := EffectiveParallel(&Override{Parallel: in.Parallel})
	if slots > 1 {
		perTokGB *= float64(slots)
		kvConstGB *= float64(slots)
	}

	// Rope scaling lifts the trained-length ceiling; without it this is nativeCtx.
	modelMax := ropeCeiling(meta, in.RopeScaling, in.Ctx)

	target := s.TargetVramGB
	if in.TargetVramGB > 0 {
		target = in.TargetVramGB
	}

	// Draft overhead: a baked-in MTP nextn layer costs its compute pad here. A
	// separate draft file (Gemma-4's MTP sidecar, or any DFlash drafter, always
	// separate) instead charges its real on-disk weights + a small KV/compute pad,
	// so big drafts scale up rather than under-counting.
	specOh := draftOverheadGB(spec, draftGB)
	// ...plus, for a baked-in nextn head, its own KV over the whole window. That
	// part is a slope, so it has to reach the sizer as one: charging it as flat
	// overhead would price it at a ctx the sizer had not picked yet.
	dKvK, dKvV := draftKvPair(in.KvKDraft, in.KvVDraft, kvK, kvV)
	draftSlopeGB := mtpDraftSlopeFor(meta, spec, dKvK, dKvV, draftGB)
	if slots > 1 {
		draftSlopeGB *= float64(slots)
	}

	prof := profile{
		Name:     "estimate",
		Target:   target,
		Overhead: s.VramOverheadGB + specOh,
		Ctx:      in.Ctx,
		CtxExact: in.CtxExact,
		Spec:     spec,
		KvK:      kvK,
		KvV:      kvV,
		IsLong:   in.Ctx >= longCtxThreshold,

		CtxCheckpoints:    in.CtxCheckpoints,
		CheckpointMinStep: in.CheckpointMinStep,
		DraftSlopeGB:      draftSlopeGB,
	}
	// Charge the ub the launch will actually run with: effectiveUb reads the
	// override's pinned value, and passing nil here made the preview size a
	// different compute buffer than emit did for the same model.
	computeUb := effectiveUb(meta, prof, &Override{Ub: in.Ub}, prof.Ctx, target)
	computeBufGB := computeBufferGB(meta, computeUb, s.ComputeBufFactor)
	// The baked-in drafter's second context allocates a graph of its own at the
	// same ub. It is reported under Draft rather than folded into the compute
	// buffer, so the UI keeps attributing it to the thing that caused it.
	draftComputeGB := mtpDraftComputeGB(meta, spec, draftGB, computeUb, s.ComputeBufFactor)
	prof.Overhead += computeBufGB + draftComputeGB
	prof.Overhead += in.MmprojGB // "-vision" projector weights + CLIP compute reserve

	// Everything charged so far lands on --main-gpu alone, so this is the figure
	// the spawn guard re-derives the ratio from (FixedGB below). Every device
	// past the main one pays its own runtime context on top; that part is not
	// splittable either, but it belongs to the POOLED budget rather than to the
	// main card, so it is added after FixedGB is taken. Generate does the same,
	// and the two must agree or a spawn-time retune moves the split for no reason
	// other than which code path computed it.
	mainFixedGB := prof.Overhead
	if gpus := s.GpuSetOrEmpty(); gpus.Multi() {
		prof.Overhead += gpus.ExtraDeviceOverheadGB()
	}

	ctx, plan, kvReserve, planCkptGB, err := sizeProfile(meta, s, prof, perTokGB, kvConstGB, modelMax, in.KvInRam)
	if err != nil {
		return EstimateResult{}, err
	}
	// Checkpoints live wherever the KV cache does. When KV is in RAM they cost
	// no VRAM, so report 0 to stay consistent with EstVramGB (which sizeProfile
	// only inflates in the VRAM-KV branch).
	checkpointGB := 0.0
	if !in.KvInRam {
		ckptCtxCeil := modelMax
		if prof.Ctx != 0 {
			ckptCtxCeil = min(ckptCtxCeil, prof.Ctx)
		}
		checkpointGB = checkpointReserveGB(prof, perTokGB, kvConstGB, ckptCtxCeil, meta.FullAttnInterval > 0)
	}

	ngl, ncpuMoe := forceLowActiveMoE(meta, plan, prof, kvReserve)
	if in.CpuOffload > 0 || in.CpuOffloadSet {
		ngl, ncpuMoe = applyForcedOffload(meta, in.CpuOffload)
	}
	if in.CpuDevice {
		ngl, ncpuMoe = 0, 0
	}
	// Any placement that differs from the one GetLoadPlan priced has to be
	// re-priced: forceLowActiveMoE rewrites -ngl/--n-cpu-moe in place and used to
	// leave plan.EstVramGB describing the placement it just discarded.
	if ngl != plan.Ngl || ncpuMoe != plan.NCpuMoe || in.CpuDevice {
		plan.EstVramGB, plan.EstRamGB = estForOffload(meta, prof, kvReserve, planCkptGB, ngl, ncpuMoe)
		if in.CpuDevice {
			plan.EstVramGB, plan.EstRamGB = 0, plan.EstVramGB+plan.EstRamGB
		}
		plan.RamExceeded = s.MaxRamGB > 0 && plan.EstRamGB > s.MaxRamGB
	}

	// Dense checkpoints are split GPU/RAM by the layer placement (see sizeProfile),
	// so the VRAM portion is only the GPU-resident fraction. Scale the reported
	// figure to match EstVramGB. MoE keeps KV (and checkpoints) fully in VRAM.
	if in.CpuDevice {
		checkpointGB = 0
	}
	if checkpointGB > 0 && !meta.IsMoE && meta.BlockCount > 0 {
		gpuFrac := float64(min(ngl, int(meta.BlockCount))) / float64(meta.BlockCount)
		if gpuFrac < 0 {
			gpuFrac = 0
		}
		checkpointGB *= gpuFrac
	}

	// sizeProfile folds the drafter's KV into kvReserve (it sizes against the
	// combined slope). Split it back out so the UI's KV segment stays the main
	// cache and the Draft segment shows what the drafter actually costs.
	// KV-in-RAM is the exception: that branch sizes against the main slope alone
	// (the reserve it reports is a flat RAM-side placeholder), so there is
	// nothing folded in to take back out.
	draftKvGB := 0.0
	if !in.KvInRam {
		draftKvGB = min(draftSlopeGB*float64(ctx), kvReserve)
		kvReserve -= draftKvGB
	}

	res := EstimateResult{
		Ctx:          ctx,
		Ngl:          ngl,
		NCpuMoe:      ncpuMoe,
		EstVramGB:    plan.EstVramGB,
		EstRamGB:     plan.EstRamGB,
		TargetVramGB: target,
		MaxRamGB:     s.MaxRamGB,
		KvReserveGB:  kvReserve,
		CheckpointGB: checkpointGB,
		DraftGB:      specOh + draftComputeGB + draftKvGB,
		ComputeBufGB: computeBufGB,
		MmprojGB:     in.MmprojGB,
		OverheadGB:   s.VramOverheadGB,
		FixedGB:      mainFixedGB,
		RamExceeded:  plan.RamExceeded,
		IsMoE:        meta.IsMoE,
	}
	// A CPU-device pin has no GPU, so every VRAM component is zero, not just the
	// total: the UI draws each one as its own segment of the VRAM bar, and the
	// spawn guard charges FixedGB to the main card.
	if in.CpuDevice {
		res.KvReserveGB, res.DraftGB, res.ComputeBufGB = 0, 0, 0
		res.MmprojGB, res.OverheadGB, res.FixedGB = 0, 0, 0
	}
	return res, nil
}