"""llama.cpp VRAM Calculator — Hugging Face Space.

A pre-flight VRAM predictor + llama-server command-builder. Hero use case:
"I want N context on a model trained for less — give me the command and
warn me about coherence." Two-phase UI:

  1. Fetch phase (network, once): repo id + GGUF file -> parse_hf_range ->
     cached ModelArch. Network happens here and never again.
  2. Live phase (microseconds): every input change re-runs estimate() on the
     cached arch; the command, breakdown, and graphs update live.

Two graphs: VRAM-vs-n_ctx (one curve per quant, the context ceiling) and
VRAM-vs-quant at fixed n_ctx (the cost+quality frontier, quality as color).
See PLAN.md for the full design spine.
"""

from __future__ import annotations

import gradio as gr
import matplotlib
matplotlib.use("Agg")  # headless: render to PNG, no display needed
import matplotlib.pyplot as plt
from huggingface_hub import HfApi
import spaces  # noqa: F401  — present so ZeroGPU detects a GPU-aware Space

from vramcalc import (
    QUANT_BPW,
    ModelArch,
    Inputs,
    DraftInputs,
    Mmproj,
    estimate,
    command_preview,
    format_bytes,
    quant_from_filename,
    parse_hf_range,
    max_context,
    best_quant,
    min_gpu_setup,
    mmproj_bytes_from_tensors,
    auto_configure_yarn,
    yarn_coherence_warnings,
    extension_ratio,
)
from vramcalc.draft import SPEC_TYPES as REAL_SPEC_TYPES
from vramcalc.draft import SPEC_DEFAULTS, SPEC_FAMILY
from vramcalc.presets import PRESETS, PRESET_NAMES


QUANT_CHOICES = list(QUANT_BPW.keys())
CACHE_DTYPES = ["f16", "bf16", "f32", "q8_0", "q8_1", "q4_0", "q4_1", "q5_0", "q5_1"]
COMPUTE_DTYPES = ["f16", "bf16", "f32"]
SPLIT_MODES = ["layer", "tensor", "none"]
SPEC_TYPES = list(REAL_SPEC_TYPES)  # the real llama.cpp --spec-type enum
MTP_CACHE_DTYPES = ["(same as target)", "f16", "bf16", "f32", "q8_0"]

# Grouped spec choices for the dropdown: (label, value) tuples. Gradio
# returns the bare type (value), so the existing plumbing (Inputs,
# command_preview, is_weightless) is unchanged; the label carries the family
# grouping for the UI.
SPEC_GROUPED = [
    (f"{family} / {t}", t)
    for family, types in SPEC_FAMILY.items() for t in types
]
SPEC_GROUPED.append(("none / none", "none"))

# The 4-bit-class floor quants (q4_0 / ROCmFP4 / nvfp4). Sub-4-bit quants
# (2-bit/3-bit family) are offered only on explicit override and flagged
# "you asked for it." See PLAN.md "Objective".
FLOOR_QUANTS = {"q4_0", "ROCmFP4", "ROCmFPX", "nvfp4"}  # ~4.0+ bpw

# n_ctx sweep for the VRAM-vs-context graph (log-ish, capped at 1M).
GRAPH_CTX_GRID = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144,
                  524288, 1048576]


def list_gguf_files(repo_id: str, hf_token: str):
    """List *.gguf files in a repo and auto-detect their quant names.

    Returns gr.update objects (not bare lists) so Gradio 6 reliably re-renders
    the dropdown and pre-selects the first file. A plain list return no longer
    repopulates a Dropdown under SSR in Gradio 6.
    """
    if not repo_id or not repo_id.strip():
        return gr.update(choices=[], value=None), "Enter a Hugging Face repo id."
    try:
        api = HfApi(token=hf_token or None)
        files = api.list_repo_files(repo_id=repo_id.strip())
    except Exception as e:  # noqa: BLE001
        return gr.update(choices=[], value=None), f"Error listing repo: {e}"
    ggufs = sorted(f for f in files if f.lower().endswith(".gguf"))
    if not ggufs:
        return gr.update(choices=[], value=None), f"No .gguf files found in {repo_id!r}."
    choices = []
    for f in ggufs:
        q = quant_from_filename(f)
        label = f"{f}" + (f"  [{q}]" if q else "")
        choices.append((label, f))
    return (
        gr.update(choices=choices, value=ggufs[0]),
        f"Found {len(ggufs)} GGUF file(s).",
    )


def fetch_arch(repo_id: str, filename: str, hf_token: str):
    """Range-read a GGUF header from HF and return editable arch fields.

    Also returns a gr.update for the Quantization dropdown that syncs it to the
    quant detected from the filename (e.g. ROCmFP4/ROCmFPX) — a no-op when the
    filename has no recognized quant token, so the user's manual pick is kept.
    The GGUF header's tensor types for ROCm formats use custom IDs stock
    llama.cpp does not map, so the filename is the reliable source here.
    """
    detected = quant_from_filename(filename) if filename else None
    if detected and detected not in QUANT_CHOICES:
        detected = None
    quant_update = gr.update(value=detected) if detected else gr.update()
    if not repo_id or not filename:
        return (*_empty_arch_fields(), "Pick a GGUF file first.", quant_update)
    try:
        meta = parse_hf_range(
            repo_id.strip(), filename, token=hf_token or None
        )
    except Exception as e:  # noqa: BLE001
        return (*_empty_arch_fields(), f"Error reading GGUF header: {e}", quant_update)
    quant_note = f", quant {detected}" if detected else ""
    if not meta.n_layer:
        return (
            *_arch_to_fields(meta),
            "Parsed header but architecture fields look empty; "
            "edit them manually below.",
            quant_update,
        )
    # Resolve head_dim (GGUF may set it explicitly; otherwise derive) and the
    # count of KV-bearing layers for the success message, so hybrid-attention
    # models (Qwen3.6 / Ornith / Gemma4) surface the correction up front.
    resolved_head_dim = meta.head_dim or (
        meta.n_embd // meta.n_head if meta.n_head > 0 else 0
    )
    kv_layers = meta.n_full_attn_layers or meta.n_layer
    hybrid_note = ""
    if meta.n_full_attn_layers and meta.n_layer \
            and meta.n_full_attn_layers != meta.n_layer:
        hybrid_note = (
            f", 🧬 hybrid {kv_layers}/{meta.n_layer} layers carry KV"
        )
    return (
        *_arch_to_fields(meta),
        f"Fetched {meta.architecture or 'model'}: "
        f"{meta.n_layer} layers, {meta.n_embd} embd, "
        f"{meta.n_head}/{meta.n_head_kv} heads (head_dim {resolved_head_dim}), "
        f"ctx {meta.training_ctx}, "
        f"params {meta.params or 'n/a'}{quant_note}{hybrid_note}.",
        quant_update,
    )


def _empty_arch_fields():
    return _arch_to_fields(ModelArch())


def list_mmproj_files(repo_id: str, hf_token: str):
    """List mmproj-*.gguf files in a repo (for the multimodal projector picker)."""
    if not repo_id or not repo_id.strip():
        return gr.update(choices=[], value=None), ""
    try:
        api = HfApi(token=hf_token or None)
        files = api.list_repo_files(repo_id=repo_id.strip())
    except Exception as e:  # noqa: BLE001
        return gr.update(choices=[], value=None), f"Error listing repo: {e}"
    mmps = sorted(
        f for f in files
        if f.lower().endswith(".gguf") and "mmproj" in f.lower()
    )
    if not mmps:
        return gr.update(choices=[], value=None), "No mmproj files in this repo."
    choices = [(f, f) for f in mmps]
    return gr.update(choices=choices, value=mmps[0]), f"Found {len(mmps)} mmproj file(s)."


def fetch_mmproj_bytes(repo_id: str, filename: str, hf_token: str):
    """Range-read an mmproj GGUF header and estimate its weight bytes."""
    if not repo_id or not filename:
        return 0, "Pick an mmproj file first."
    try:
        from vramcalc.gguf import parse_hf_range, parse_header_with_tensors
        # parse_hf_range returns a GGUFMetadata; use the raw dict's dtype hist.
        meta = parse_hf_range(repo_id.strip(), filename, token=hf_token or None)
    except Exception as e:  # noqa: BLE001
        return 0, f"Error reading mmproj header: {e}"
    hist = meta.raw.get(".tensor_dtype_hist") if meta.raw else None
    if isinstance(hist, dict) and hist:
        b = mmproj_bytes_from_tensors({int(k): int(v) for k, v in hist.items()})
    else:
        # fall back to params * f16 bpw
        b = (meta.params or 0) * 2.0
    return int(b), f"mmproj {filename}: ~{format_bytes(b)} ({meta.params or 0} elems)"


def _base_inputs(arch_fields, n_ctx, quant, gpu_vram_text, split_mode, main_gpu,
                 n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype, safety_margin):
    """Build (ModelArch, Inputs) for the auto-fit solvers from UI fields."""
    arch = _fields_to_arch(arch_fields)
    gpus = _parse_gpu_list(gpu_vram_text.replace("u", "").replace("U", ""))
    unified = _parse_unified_flags(gpu_vram_text)
    nb = int(n_batch or 2048)
    ub = int(n_ubatch or nb)
    if ub > nb:
        ub = nb
    inp = Inputs(
        quant=quant, n_ctx=int(n_ctx or 8192),
        cache_dtype=cache_dtype, flash_attn=bool(flash_attn),
        compute_dtype=compute_dtype, n_batch=nb, n_ubatch=ub,
        gpu_vram_gb=gpus, split_mode=split_mode or "layer",
        main_gpu=int(main_gpu or 0), unified_flags=unified,
        safety_margin_pct=float(safety_margin or 5.0),
    )
    return arch, inp


def run_max_context(arch_fields, quant, gpu_vram_text, split_mode, main_gpu,
                    n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype,
                    safety_margin):
    arch, inp = _base_inputs(arch_fields, 8192, quant, gpu_vram_text, split_mode,
                             main_gpu, n_batch, n_ubatch, cache_dtype, flash_attn,
                             compute_dtype, safety_margin)
    if arch.n_layer <= 0 or arch.params <= 0:
        return "⚠️ Architecture incomplete — fetch a GGUF or load a preset first."
    res = max_context(arch, inp, step=1024, max_ctx=max(131072, arch.training_ctx * 4))
    if not res.fits:
        return f"❌ {res.note}\n\n(Total @ n_ctx={res.n_ctx}: {format_bytes(res.total_bytes)})"
    return (
        f"**Max context that fits:** **{res.n_ctx:,}** tokens\n\n"
        f"Total VRAM: {format_bytes(res.total_bytes)}\n\n"
        f"_{res.note}_"
    )


def run_best_quant(arch_fields, n_ctx, gpu_vram_text, split_mode, main_gpu,
                   n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype,
                   safety_margin):
    arch, inp = _base_inputs(arch_fields, n_ctx, "Q4_K_M", gpu_vram_text, split_mode,
                             main_gpu, n_batch, n_ubatch, cache_dtype, flash_attn,
                             compute_dtype, safety_margin)
    if arch.n_layer <= 0 or arch.params <= 0:
        return "⚠️ Architecture incomplete — fetch a GGUF or load a preset first."
    res = best_quant(arch, inp)
    if not res.fits:
        return f"❌ {res.note}\n\nLowest-bpw attempt: {res.fallback_quant} ({res.fallback_bpw} bpw) → {format_bytes(res.fallback_total)}"
    lines = [
        f"**Recommended quant:** **{res.quant}** ({res.bpw} bpw) → {format_bytes(res.total_bytes)}",
    ]
    if res.fallback_quant:
        lines.append(
            f"**Safe fallback:** {res.fallback_quant} ({res.fallback_bpw} bpw) → "
            f"{format_bytes(res.fallback_total)}"
        )
    return "\n\n".join(lines)


def run_min_gpu_setup(arch_fields, n_ctx, quant, gpu_vram_text, split_mode, main_gpu,
                      n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype,
                      safety_margin):
    arch, inp = _base_inputs(arch_fields, n_ctx, quant, gpu_vram_text, split_mode,
                             main_gpu, n_batch, n_ubatch, cache_dtype, flash_attn,
                             compute_dtype, safety_margin)
    if arch.n_layer <= 0 or arch.params <= 0:
        return "⚠️ Architecture incomplete — fetch a GGUF or load a preset first."
    res = min_gpu_setup(arch, inp)
    badge = "✅" if res.fits else "❌"
    subset = ", ".join(f"{g} GB" for g in res.subset)
    return f"{badge} {res.note}\n\nSubset: [{subset}]\n\nTotal: {format_bytes(res.total_bytes)}"


def auto_configure_yarn_handler(arch_fields, n_ctx):
    """Auto-derive YaRN/RoPE params from training_ctx + target n_ctx.

    Writes the five YaRN gr.Number inputs back via gr.update and returns a
    markdown status line as the sixth output.
    """
    arch = _fields_to_arch(arch_fields)
    cfg = auto_configure_yarn(arch.training_ctx, int(n_ctx or 8192))
    badge = "✅" if cfg.scaling else "ℹ️"
    status = f"{badge} {cfg.note}" if cfg.note else ""
    return (
        gr.update(value=cfg.rope_freq_scale),
        gr.update(value=cfg.yarn_ext_factor),
        gr.update(value=cfg.yarn_attn_factor),
        gr.update(value=cfg.yarn_beta_fast),
        gr.update(value=cfg.yarn_beta_slow),
        status,
    )


def _arch_to_fields(m: ModelArch):
    return [
        m.name, m.architecture, m.n_layer, m.n_embd, m.n_head, m.n_head_kv,
        m.training_ctx, m.params, m.rope_freq_base, m.n_expert,
        m.n_expert_used, m.n_mtp, m.head_dim, m.n_full_attn_layers,
    ]


ARCH_FIELD_NAMES = [
    "name", "architecture", "n_layer", "n_embd", "n_head", "n_head_kv",
    "training_ctx", "params", "rope_freq_base", "n_expert",
    "n_expert_used", "n_mtp", "head_dim", "n_full_attn_layers",
]


def _fields_to_arch(fields) -> ModelArch:
    values = {k: v for k, v in zip(ARCH_FIELD_NAMES, fields)}
    return ModelArch(
        name=str(values["name"] or ""),
        architecture=str(values["architecture"] or ""),
        n_layer=int(values["n_layer"] or 0),
        n_embd=int(values["n_embd"] or 0),
        n_head=int(values["n_head"] or 0),
        n_head_kv=int(values["n_head_kv"] or 0),
        training_ctx=int(values["training_ctx"] or 0),
        params=int(values["params"] or 0),
        rope_freq_base=float(values["rope_freq_base"] or 10000.0),
        n_expert=int(values["n_expert"] or 0),
        n_expert_used=int(values["n_expert_used"] or 0),
        n_mtp=int(values["n_mtp"] or 0),
        head_dim=int(values["head_dim"] or 0),
        n_full_attn_layers=int(values["n_full_attn_layers"] or 0),
    )


def load_preset(name: str):
    if name and name in PRESETS:
        return (*_arch_to_fields(PRESETS[name]), f"Loaded preset: {name}")
    return (*_empty_arch_fields(), "")


def _hybrid_md(*fields) -> str:
    """Markdown for the 🧬 Hybrid attention badge + KV-bearing-layers line.

    Takes the arch field values as varargs (wired to all arch_inputs
    components) and round-trips them through _fields_to_arch. Empty for
    pure-attention / unknown archs (n_full_attn_layers == 0 or == n_layer), so
    the badge only appears when the ~4× KV correction is actually in effect —
    the case users need to see explained, not the common case.
    """
    arch = _fields_to_arch(fields)
    if not arch.n_full_attn_layers or not arch.n_layer:
        return ""
    if arch.n_full_attn_layers == arch.n_layer:
        return ""
    full = arch.n_full_attn_layers
    mtp = arch.n_mtp
    total = full + mtp
    resolved_head_dim = arch.head_dim or (
        arch.n_embd // arch.n_head if arch.n_head > 0 else 0
    )
    head_note = (
        f" head_dim {resolved_head_dim} (≠ n_embd/n_head)"
        if arch.head_dim and arch.n_head and arch.head_dim != arch.n_embd // arch.n_head
        else ""
    )
    mtp_part = f" + {mtp} MTP" if mtp else ""
    return (
        f"**🧬 Hybrid attention** — {full}/{arch.n_layer} layers carry KV; "
        f"the other {arch.n_layer - full} are linear/recurrent (Gated DeltaNet) "
        f"with a constant-size state (no O(n_ctx) KV)."
        f"{head_note}  \n"
        f"**KV-bearing layers:** {full}{mtp_part} = **{total}** "
        f"(vs {arch.n_layer} for a pure-attention model)."
    )


def _mtp_hint_md(*fields) -> str:
    """Neutral MTP hint next to the MTP control.

    Replaces the earlier "MTP conflicts with --parallel / --mmproj" warnings —
    those conflicts aren't documented in current llama.cpp. We keep a single
    factual line so users know how MTP is actually enabled (no standalone
    `--mtp N` server flag exists; the only documented path is
    `--spec-type draft-mtp`).
    """
    arch = _fields_to_arch(fields)
    if arch.n_mtp <= 0:
        return ""
    return (
        f"ℹ️ Model has {arch.n_mtp} MTP head(s). VRAM budgeted above. "
        f"Enable at inference with `--spec-type draft-mtp` — there is no "
        f"standalone `--mtp` server flag in current llama.cpp."
    )


def _parse_gpu_list(text: str) -> list[float]:
    out = []
    for tok in (text or "").replace(";", ",").split(","):
        tok = tok.strip()
        if tok:
            try:
                out.append(float(tok))
            except ValueError:
                pass
    return out or [24.0]


def _parse_unified_flags(text: str) -> list[bool]:
    """Parse the 'u' suffix per device: 24,16u -> [False, True]."""
    flags = []
    for tok in (text or "").replace(";", ",").split(","):
        tok = tok.strip().lower()
        if not tok:
            continue
        flags.append(tok.endswith("u"))
    return flags


def _parse_tensor_split(text: str) -> list[float] | None:
    out = []
    for tok in (text or "").replace(";", ",").split(","):
        tok = tok.strip()
        if tok:
            try:
                out.append(float(tok))
            except ValueError:
                pass
    return out or None


# --- Live recompute + graphs -------------------------------------------------
#
# The fetch phase caches a ModelArch; everything below is pure estimate()
# math on that cached arch, so it runs in microseconds and can fire on every
# input change (live sliders). No network in the live phase. See PLAN.md.


def _build_inputs(
    arch_fields, quant, n_ctx, cache_dtype, flash_attn, compute_dtype,
    n_batch, n_ubatch, n_prompt, parallel_slots, parallel_sizing,
    rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
    yarn_beta_fast, yarn_beta_slow,
    gpu_vram_text, split_mode, main_gpu, tensor_split_text, safety_margin,
    mtp_cache_dtype,
    spec_type, draft_quant, draft_params, draft_n_layer, draft_n_max,
    draft_n_min, draft_p_min, draft_p_split,
    mmproj_enabled, mmproj_file, mmproj_offload, mmproj_bytes_text,
) -> tuple[ModelArch, Inputs]:
    arch = _fields_to_arch(arch_fields)
    gpus = _parse_gpu_list(gpu_vram_text.replace("u", "").replace("U", ""))
    unified = _parse_unified_flags(gpu_vram_text)
    tensor_split = _parse_tensor_split(tensor_split_text)

    draft = None
    if spec_type and spec_type != "none":
        draft = DraftInputs(
            spec_type=spec_type,
            quant=draft_quant,
            params=int(draft_params or 0),
            n_layer=int(draft_n_layer or 1),
            n_ctx=int(n_ctx),
            cache_dtype=cache_dtype,
            n_max=int(draft_n_max or 0),
            n_min=int(draft_n_min or 0),
            p_min=float(draft_p_min or 0.0),
            p_split=float(draft_p_split or 0.0),
        )

    mmproj = None
    if mmproj_enabled:
        mmproj = Mmproj(
            filename=mmproj_file or "",
            enabled=True,
            offload=bool(mmproj_offload),
            bytes_=float(mmproj_bytes_text or 0.0),
        )

    # ubatch must be <= batch (llama.cpp constraint); clamp defensively.
    nb = int(n_batch or 512)
    ub = int(n_ubatch or nb)
    if ub > nb:
        ub = nb

    inp = Inputs(
        quant=quant,
        n_ctx=int(n_ctx),
        cache_dtype=cache_dtype,
        flash_attn=bool(flash_attn),
        compute_dtype=compute_dtype,
        n_batch=nb,
        n_ubatch=ub,
        n_prompt=int(n_prompt or 0),
        n_parallel=max(1, int(parallel_slots or 1)),
        parallel_sizing=bool(parallel_sizing),
        rope_freq_scale=float(rope_freq_scale),
        yarn_ext_factor=float(yarn_ext_factor),
        yarn_attn_factor=float(yarn_attn_factor),
        yarn_beta_fast=float(yarn_beta_fast),
        yarn_beta_slow=float(yarn_beta_slow),
        gpu_vram_gb=gpus,
        split_mode=split_mode or "layer",
        main_gpu=int(main_gpu or 0),
        tensor_split=tensor_split,
        unified_flags=unified,
        mtp_cache_dtype=mtp_cache_dtype if mtp_cache_dtype else None,
        draft=draft,
        mmproj=mmproj,
        safety_margin_pct=float(safety_margin),
    )
    return arch, inp


def graph_vram_vs_ctx(arch: ModelArch, inp: Inputs) -> plt.Figure | None:
    """Graph 1: VRAM vs n_ctx, one curve per quant.

    Shows the context ceiling each quant reaches on the configured GPUs. The
    user's current n_ctx + the budget total are overlaid so the slider sits
    on the curve. Returns None when the arch is incomplete.
    """
    if arch.n_layer <= 0 or arch.n_embd <= 0 or arch.params <= 0:
        return None
    budget_gb = sum(inp.gpu_vram_gb)
    # sweep the grid up to max(training_ctx*4, current n_ctx, 1M)
    cap = max(arch.training_ctx * 4, inp.n_ctx, 1 << 20)
    grid = [c for c in GRAPH_CTX_GRID if c <= cap] or [cap]

    # order quants low->high bpw so the legend reads bottom-to-top
    quants = sorted(QUANT_CHOICES, key=lambda q: QUANT_BPW[q])
    fig, ax = plt.subplots(figsize=(7, 4.2))
    palette = plt.cm.viridis
    n_q = max(1, len(quants))
    for i, q in enumerate(quants):
        ys = []
        for c in grid:
            from dataclasses import replace
            bd = estimate(arch, replace(inp, quant=q, n_ctx=c))
            ys.append(bd.total_bytes / (1 << 30))
        color = palette(i / max(1, n_q - 1))
        ax.plot(grid, ys, marker="o", markersize=3, linewidth=1.6,
                label=q, color=color)

    # current config marker + budget line
    cur = estimate(arch, inp)
    ax.axhline(budget_gb, color="crimson", linestyle="--", linewidth=1.2,
               label=f"budget {budget_gb:.0f} GB")
    ax.scatter([inp.n_ctx], [cur.total_bytes / (1 << 30)],
               color="crimson", zorder=5, s=55,
               label=f"current ({inp.quant})")
    ax.set_xscale("log")
    ax.set_xlabel("context length (n_ctx)")
    ax.set_ylabel("total VRAM (GiB)")
    ax.set_title("VRAM vs context, by quant")
    ax.grid(True, which="both", alpha=0.25)
    ax.legend(fontsize=7, ncol=2, loc="upper left")
    fig.tight_layout()
    return fig


def graph_vram_vs_quant(arch: ModelArch, inp: Inputs) -> plt.Figure | None:
    """Graph 2: VRAM vs quant at fixed n_ctx, colored by quality (bpw).

    The cost+quality frontier: each bar's height is VRAM, its color is the
    quant's bpw (proxy for quality). The budget is a horizontal line; bars
    below it fit. Returns None when the arch is incomplete.
    """
    if arch.n_layer <= 0 or arch.n_embd <= 0 or arch.params <= 0:
        return None
    budget_gb = sum(inp.gpu_vram_gb)
    from dataclasses import replace
    quants = sorted(QUANT_CHOICES, key=lambda q: QUANT_BPW[q])
    ys, bpws, labels = [], [], []
    for q in quants:
        bd = estimate(arch, replace(inp, quant=q))
        ys.append(bd.total_bytes / (1 << 30))
        bpws.append(QUANT_BPW[q])
        labels.append(q)

    fig, ax = plt.subplots(figsize=(7, 4.2))
    cmap = plt.cm.plasma
    norm = plt.Normalize(vmin=min(bpws), vmax=max(bpws))
    colors = [cmap(norm(b)) for b in bpws]
    bars = ax.bar(range(len(quants)), ys, color=colors, edgecolor="black",
                  linewidth=0.4)
    ax.axhline(budget_gb, color="crimson", linestyle="--", linewidth=1.2,
               label=f"budget {budget_gb:.0f} GB")
    # highlight the currently-selected quant
    if inp.quant in quants:
        idx = quants.index(inp.quant)
        bars[idx].set_edgecolor("crimson")
        bars[idx].set_linewidth(2.0)
    ax.set_xticks(range(len(quants)))
    ax.set_xticklabels(labels, rotation=55, ha="right", fontsize=7)
    ax.set_ylabel("total VRAM (GiB)")
    ax.set_title(f"VRAM vs quant @ ctx {inp.n_ctx}  (color = quality/bpw)")
    ax.grid(True, axis="y", alpha=0.25)
    ax.legend(fontsize=8, loc="upper left")
    # colorbar for quality
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    fig.colorbar(sm, ax=ax, label="bits per weight (quality)")
    fig.tight_layout()
    return fig


def graph_scratch_vs_batch(arch: ModelArch, inp: Inputs) -> plt.Figure | None:
    """Graph 3: VRAM scratch vs n_batch, showing the batch/VRAM/prefill tradeoff.

    Sweeps n_batch over a grid and plots compute_scratch (activation/logits
    buffers) plus total VRAM. The current n_batch is overlaid so the slider
    sits on the curve. Returns None when the arch is incomplete.
    """
    if arch.n_layer <= 0 or arch.n_embd <= 0 or arch.params <= 0:
        return None
    from dataclasses import replace
    grid = [256, 512, 1024, 2048, 4096, 8192]
    # keep ubatch = batch for a clean apples-to-apples sweep
    scratch_ys, total_ys = [], []
    for b in grid:
        bd = estimate(arch, replace(inp, n_batch=b, n_ubatch=b))
        scratch_ys.append(bd.compute_scratch_bytes / (1 << 20))  # MiB
        total_ys.append(bd.total_bytes / (1 << 30))             # GiB
    fig, ax1 = plt.subplots(figsize=(7, 4.2))
    color_s = "tab:blue"
    ax1.plot(grid, scratch_ys, marker="o", color=color_s, linewidth=1.6,
             label="compute scratch (MiB)")
    ax1.set_xlabel("n_batch (-b, = -ub here)")
    ax1.set_ylabel("compute scratch (MiB)", color=color_s)
    ax1.tick_params(axis="y", labelcolor=color_s)
    ax1.set_xscale("log")
    ax1.set_xticks(grid)
    ax1.set_xticklabels([str(g) for g in grid], rotation=0)
    ax2 = ax1.twinx()
    color_t = "tab:orange"
    ax2.plot(grid, total_ys, marker="s", linestyle="--", color=color_t,
             linewidth=1.2, label="total VRAM (GiB)")
    ax2.set_ylabel("total VRAM (GiB)", color=color_t)
    ax2.tick_params(axis="y", labelcolor=color_t)
    cur = estimate(arch, inp)
    ax1.axvline(inp.n_batch, color="crimson", linestyle=":", linewidth=1.2)
    ax1.scatter([inp.n_batch], [cur.compute_scratch_bytes / (1 << 20)],
                color="crimson", zorder=5, s=55,
                label=f"current (b={inp.n_batch}, ub={inp.n_ubatch})")
    ax1.set_title("VRAM scratch vs n_batch")
    ax1.grid(True, which="both", alpha=0.25)
    fig.legend(loc="upper left", bbox_to_anchor=(0.12, 0.98), fontsize=8)
    fig.tight_layout()
    return fig


def live_compute(
    arch_fields,
    quant, n_ctx, cache_dtype, flash_attn, compute_dtype,
    n_batch, n_ubatch, n_prompt,
    parallel_slots, parallel_sizing,
    rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
    yarn_beta_fast, yarn_beta_slow,
    gpu_vram_text, split_mode, main_gpu, tensor_split_text, safety_margin,
    mtp_cache_dtype,
    spec_type, draft_quant, draft_params, draft_n_layer, draft_n_max,
    draft_n_min, draft_p_min, draft_p_split,
    mmproj_enabled, mmproj_file, mmproj_offload, mmproj_bytes_text,
):
    """Live recompute: estimate() on the cached arch + graphs + command.

    Returns (summary, command, graph1, graph2, graph3). No network — the arch
    was fetched once and stored in arch_state. Fires on every input change so
    sliders are live.
    """
    arch, inp = _build_inputs(
        arch_fields, quant, n_ctx, cache_dtype, flash_attn, compute_dtype,
        n_batch, n_ubatch, n_prompt, parallel_slots, parallel_sizing,
        rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
        yarn_beta_fast, yarn_beta_slow,
        gpu_vram_text, split_mode, main_gpu, tensor_split_text, safety_margin,
        mtp_cache_dtype,
        spec_type, draft_quant, draft_params, draft_n_layer, draft_n_max,
        draft_n_min, draft_p_min, draft_p_split,
        mmproj_enabled, mmproj_file, mmproj_offload, mmproj_bytes_text,
    )
    if arch.n_layer <= 0 or arch.n_embd <= 0 or arch.params <= 0:
        return (
            "⚠️ Fetch a GGUF (or load a preset) to populate the architecture, "
            "then adjust the sliders.",
            "", None, None, None,
        )
    bd = estimate(arch, inp)

    # breakdown table
    rows = [
        ["Weights (GGUF, " + quant + ")", format_bytes(bd.weights_bytes)],
        ["KV cache (" + cache_dtype + ")", format_bytes(bd.kv_cache_bytes)],
        ["Compute / scratch", format_bytes(bd.compute_scratch_bytes)],
        ["MTP overhead", format_bytes(bd.mtp_overhead_bytes)],
    ]
    if bd.draft_bytes_:
        rows.append(["Draft model (spec. decoding)", format_bytes(bd.draft_bytes_)])
    if bd.mmproj_bytes_:
        rows.append(["Multimodal projector", format_bytes(bd.mmproj_bytes_)])
    rows += [
        ["GGUF header / overhead", format_bytes(bd.gguf_overhead_bytes)],
        ["Safety margin (" + str(inp.safety_margin_pct) + "%)",
         format_bytes(bd.safety_margin_bytes)],
        ["**Total**", f"**{format_bytes(bd.total_bytes)}**"],
    ]
    breakdown_md = "| Component | Size |\n|---|---|\n" + "\n".join(
        f"| {a} | {b} |" for a, b in rows
    )

    # warnings (incl. the YaRN coherence ladder) + the build-support caveat
    warns = list(bd.warnings)
    if quant not in FLOOR_QUANTS and QUANT_BPW.get(quant, 99) < 4.0:
        warns.append(
            f"{quant} ({QUANT_BPW[quant]} bpw) is below the 4-bit quality "
            f"floor — you asked for it."
        )
    warns.append(
        "Estimate assumes your llama.cpp build supports the chosen quant "
        "and KV dtype. ROCmFP4/ROCmFPX need the ciru-ai/ROCmFPX runner."
    )
    warn_md = "\n\n**⚠️ Notes:**\n" + "\n".join(f"- {w}" for w in warns) if warns else ""

    eff_md = (
        f"\n\nEffective context (training_ctx / rope_freq_scale): "
        f"**{bd.effective_context}**"
    )

    # per-GPU table
    gpu_md = ""
    if bd.gpu and bd.gpu.assignments:
        a = bd.gpu.assignments
        header = "| GPU | VRAM | Weights | KV | Used | Free | Fits? |"
        sep = "|---|---|---|---|---|---|---|"
        body = []
        for g in a:
            kv = format_bytes(g.kv_compute_bytes) if g.kv_compute_bytes else "—"
            badge = "✅" if g.fits else "❌"
            tag = f" ({g.name})" if g.name else ""
            unote = " 🔁" if g.is_unified else ""
            body.append(
                f"| {g.index}{tag}{unote} | {format_bytes(g.vram_bytes)} | "
                f"{format_bytes(g.weight_bytes)} | {kv} | "
                f"{format_bytes(g.used_bytes)} | {format_bytes(g.free_bytes)} "
                f"| {badge} |"
            )
        total_badge = "✅ all fit" if bd.gpu.all_fit else "❌ over budget"
        gpu_md = (
            f"**Per-GPU split ({bd.gpu.split_mode} mode):**\n\n"
            + header + "\n" + sep + "\n" + "\n".join(body)
            + f"\n\nTotal VRAM: {format_bytes(bd.gpu.total_vram_bytes)} · "
            f"Total used: {format_bytes(bd.gpu.total_used_bytes)} · "
            f"{total_badge}"
        )

    cmd = command_preview(arch, inp)

    # Prefill info chip (n_prompt is informational — does NOT affect VRAM).
    prefill_chip = ""
    try:
        np_ = int(n_prompt or 0)
        nb = max(1, int(n_batch or 1))
        if np_ > 0:
            import math
            passes = math.ceil(np_ / nb)
            prefill_chip = (
                f"\n\nℹ️ Prefill: ~{passes} pass(es) for a {np_:,}-token prompt "
                f"(ceil(prompt / n_batch)). Prompt length does **not** affect "
                f"VRAM — llama.cpp pre-allocates KV for the full -c window."
            )
            if np_ > 8192:
                prefill_chip += (
                    " System prompts this long benefit from `--cache-prompt` "
                    "+ `--cache-reuse` (host RAM, ~0 VRAM)."
                )
    except (TypeError, ValueError):
        pass

    # Parallel / effective-KV note.
    kv_note = ""
    if inp.n_parallel > 1:
        if inp.parallel_sizing:
            kv_note = (
                f"\n\n🔁 KV sized for {inp.n_parallel} slots × {inp.n_ctx} = "
                f"{bd.effective_kv_ctx:,} tokens (-c {bd.effective_kv_ctx})."
            )
        else:
            per_slot = inp.n_ctx // inp.n_parallel
            kv_note = (
                f"\n\n🔁 --parallel {inp.n_parallel}: -c {inp.n_ctx} is the total "
                f"KV budget → per-slot ctx ≈ {per_slot:,}. (Enable “size KV for "
                f"N slots” to size VRAM for N×C.)"
            )

    summary = (
        f"**{arch.name or arch.architecture or 'Model'}** @ {quant}, "
        f"ctx {inp.n_ctx} ({cache_dtype} KV"
        + (", FA" if inp.flash_attn else ", no FA")
        + f"), {len(inp.gpu_vram_gb)} GPU(s) → "
        f"**{format_bytes(bd.total_bytes)}** total"
    )
    md = (summary + prefill_chip + kv_note + "\n\n" + breakdown_md + eff_md
          + warn_md + "\n\n" + gpu_md)

    g1 = graph_vram_vs_ctx(arch, inp)
    g2 = graph_vram_vs_quant(arch, inp)
    g3 = graph_scratch_vs_batch(arch, inp)
    return md, cmd, g1, g2, g3


# ZeroGPU requires at least one @spaces.GPU-decorated function in the Space at
# startup, or it refuses to start with "No @spaces.GPU function detected". The
# live recompute path is pure CPU math (no GPU), so it can't carry the
# decorator — that would queue every slider drag behind the GPU pool. This
# probe satisfies the startup check: it runs trivial estimate() work under the
# decorator and is never wired to the UI.
@spaces.GPU
def _zero_gpu_probe() -> float:
    """Satisfy the ZeroGPU startup detector. Returns a trivial estimate."""
    arch = PRESETS["Llama-3 8B"]
    bd = estimate(arch, Inputs(quant="Q4_K_M", n_ctx=8192, gpu_vram_gb=[24.0]))
    return float(bd.total_bytes)


def build_ui():
    with gr.Blocks(title="llama.cpp VRAM Calculator") as demo:
        gr.Markdown(
            "# 🦀 llama.cpp VRAM Calculator\n"
            "Estimate VRAM for a Hugging Face GGUF model: quant size, "
            "context, KV cache options, YaRN context extension, MTP heads, "
            "and multi-GPU split. Architecture is auto-fetched from the GGUF "
            "header (range-read — no full model download)."
        )

        arch_state = gr.State(_empty_arch_fields())

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 1. Model source")
                with gr.Tab("Auto-fetch from HF"):
                    repo_id = gr.Textbox(
                        label="HF repo id",
                        placeholder="e.g. bartowski/Llama-3-8B-Instruct-GGUF",
                    )
                    hf_token = gr.Textbox(
                        label="HF token (optional, for gated/private repos)",
                        type="password",
                    )
                    list_btn = gr.Button("List GGUF files")
                    file_picker = gr.Dropdown(
                        label="GGUF file", choices=[], interactive=True
                    )
                    fetch_btn = gr.Button("Fetch architecture from GGUF header")
                    fetch_status = gr.Markdown("")
                    mmproj_picker = gr.Dropdown(
                        label="mmproj file (optional, multimodal)",
                        choices=[], interactive=True,
                    )
                    mmproj_bytes_box = gr.Number(
                        label="mmproj weight bytes (auto from header; edit if needed)",
                        value=0, precision=0,
                    )
                    fetch_mmproj_btn = gr.Button("Fetch mmproj size from header")

                with gr.Tab("Presets / manual"):
                    preset_dd = gr.Dropdown(
                        label="Quick preset", choices=PRESET_NAMES, interactive=True
                    )
                    load_preset_btn = gr.Button("Load preset")

                gr.Markdown("### Architecture (editable)")
                arch_inputs = [
                    gr.Textbox(label="name", value=""),
                    gr.Textbox(label="architecture", value=""),
                    gr.Number(label="n_layer", value=0, precision=0),
                    gr.Number(label="n_embd", value=0, precision=0),
                    gr.Number(label="n_head", value=0, precision=0),
                    gr.Number(label="n_head_kv", value=0, precision=0),
                    gr.Number(label="training_ctx", value=0, precision=0),
                    gr.Number(label="params", value=0, precision=0),
                    gr.Number(label="rope_freq_base", value=10000.0),
                    gr.Number(label="n_expert (MoE)", value=0, precision=0),
                    gr.Number(label="n_expert_used", value=0, precision=0),
                    gr.Number(label="n_mtp", value=0, precision=0),
                    gr.Number(
                        label="head_dim (0 = n_embd/n_head)", value=0, precision=0,
                    ),
                    gr.Number(
                        label="n_full_attn_layers (0 = all carry KV)",
                        value=0, precision=0,
                    ),
                ]
                # Hybrid-attention badge + KV-bearing-layers line. Populated by
                # the recompute handler whenever the arch fields change, so the
                # ~4× KV correction is visible at the arch panel, not just in the
                # breakdown.
                hybrid_md = gr.Markdown("")
                # Inline MTP-conflict hint (MTP+--parallel / MTP+--mmproj), shown
                # next to the n_mtp field so the conflict is visible at the
                # control, not just in global Notes.
                mtp_hint_md = gr.Markdown("")

            with gr.Column(scale=1):
                gr.Markdown("### 2. Inference options")
                quant = gr.Dropdown(
                    label="Quantization", choices=QUANT_CHOICES, value="Q4_K_M"
                )
                n_ctx = gr.Number(label="Target context (n_ctx)", value=8192, precision=0)
                with gr.Row():
                    cache_dtype = gr.Dropdown(
                        label="KV cache dtype", choices=CACHE_DTYPES, value="f16"
                    )
                    compute_dtype = gr.Dropdown(
                        label="Compute dtype", choices=COMPUTE_DTYPES, value="f16"
                    )
                flash_attn = gr.Checkbox(label="Flash attention", value=True)
                with gr.Row():
                    n_batch = gr.Number(label="n_batch (-b, logits buffer)", value=2048, precision=0)
                    n_ubatch = gr.Number(label="n_ubatch (-ub, compute tile)", value=512, precision=0)
                n_prompt = gr.Number(
                    label="Expected prompt length (prefill only — does NOT affect VRAM)",
                    value=0, precision=0,
                )

                # MTP draft KV dtype (only meaningful when n_mtp > 0)
                with gr.Row():
                    mtp_cache_dtype = gr.Dropdown(
                        label="MTP draft KV dtype",
                        choices=MTP_CACHE_DTYPES, value="(same as target)",
                    )

                with gr.Accordion("YaRN / RoPE context extension", open=False):
                    with gr.Row():
                        rope_freq_scale = gr.Number(label="rope_freq_scale", value=1.0)
                        yarn_ext_factor = gr.Number(label="yarn_ext_factor", value=-1.0)
                        yarn_attn_factor = gr.Number(label="yarn_attn_factor", value=1.0)
                    with gr.Row():
                        yarn_beta_fast = gr.Number(label="yarn_beta_fast", value=32.0)
                        yarn_beta_slow = gr.Number(label="yarn_beta_slow", value=1.0)
                    with gr.Row():
                        yarn_auto_btn = gr.Button(
                            "Auto-configure from training_ctx + n_ctx"
                        )
                        yarn_auto_status = gr.Markdown("")

                with gr.Accordion("Speculative decoding (speed add-on)", open=False):
                    gr.Markdown(
                        "_Three families. **n-gram (weightless)** (`ngram-*`): "
                        "~zero extra VRAM, no draft model — **default-on when "
                        "headroom exists**; `ngram-mod` is the safe default for "
                        "reasoning/summarization/code. **Draft model (weighted)** "
                        "(`draft-simple/eagle3/dflash`): a draft model with its "
                        "own weights + KV; enable only with headroom and accept "
                        "the tradeoff. **MTP-from-target** (`draft-mtp`): uses the "
                        "target's MTP heads as a draft source (needs `--mtp N`)._"
                        "\n\n_Note: `--spec-type draft-mtp` (spec source) ≠ "
                        "`--mtp N` (target MTP head count) — different flags._"
                    )
                    with gr.Row():
                        spec_type = gr.Dropdown(
                            label="Spec type (grouped)", choices=SPEC_GROUPED,
                            value="none",
                        )
                        spec_default_btn = gr.Button(
                            "Use spec-default (ngram-mod)", size="sm"
                        )
                    with gr.Row():
                        draft_quant = gr.Dropdown(
                            label="Draft quant (weighted only)",
                            choices=QUANT_CHOICES, value="Q4_K_M"
                        )
                        draft_params = gr.Number(
                            label="Draft params (weighted only)",
                            value=0, precision=0
                        )
                        draft_n_layer = gr.Number(
                            label="Draft n_layer (1 for EAGLE-3)",
                            value=1, precision=0,
                        )
                    with gr.Row():
                        draft_n_max = gr.Number(
                            label="--spec-draft-n-max", value=0, precision=0
                        )
                        draft_n_min = gr.Number(
                            label="--spec-draft-n-min (ngram-mod >0 recommended)",
                            value=0, precision=0,
                        )
                    with gr.Row():
                        draft_p_min = gr.Number(label="p_min", value=0.0)
                        draft_p_split = gr.Number(label="p_split", value=0.0)

                with gr.Accordion("Multimodal (mmproj)", open=False):
                    mmproj_enabled = gr.Checkbox(
                        label="Include multimodal projector", value=False
                    )
                    mmproj_file = gr.Textbox(
                        label="mmproj filename (e.g. mmproj-F16.gguf)",
                        value="",
                    )
                    mmproj_offload = gr.Checkbox(
                        label="Offload mmproj to GPU", value=True
                    )
                    # mmproj_bytes_box + fetch_mmproj_btn declared above

                gr.Markdown("### 3. Multi-GPU budget")
                gpu_vram_text = gr.Textbox(
                    label="Per-GPU VRAM (GB, comma-separated; suffix 'u' = unified)",
                    value="24",
                    placeholder="e.g. 24,16u  (24GB discrete + 16GB unified)",
                )
                with gr.Row():
                    split_mode = gr.Dropdown(
                        label="split-mode", choices=SPLIT_MODES, value="layer"
                    )
                    main_gpu = gr.Number(label="main-gpu (none mode)", value=0, precision=0)
                tensor_split_text = gr.Textbox(
                    label="--tensor-split ratios (optional, comma-sep)",
                    value="",
                    placeholder="e.g. 3,1 (blank = auto by VRAM)",
                )
                with gr.Row():
                    safety_margin = gr.Number(label="Safety margin %", value=5.0)
                with gr.Row():
                    parallel_slots = gr.Number(
                        label="--parallel (server slots)", value=1, precision=0
                    )
                    parallel_sizing = gr.Checkbox(
                        label="Size KV for N slots × C ctx (scales VRAM)",
                        value=False,
                    )

                compute_btn = gr.Button("Compute VRAM", variant="primary")
                gr.Markdown(
                    "_Live: every slider/dropdown above updates the results "
                    "and graphs on change. The Compute button is a manual "
                    "trigger for the same recompute._"
                )

        gr.Markdown("### Results")
        result_md = gr.Markdown("")
        with gr.Row():
            graph_ctx = gr.Plot(label="VRAM vs context (by quant)")
            graph_quant = gr.Plot(label="VRAM vs quant @ ctx (color = quality)")
        with gr.Row():
            graph_batch = gr.Plot(label="VRAM scratch vs n_batch")
        with gr.Accordion("llama.cpp launch command preview", open=True):
            cmd_text = gr.Textbox(
                label="Command (copyable)", lines=10, interactive=False
            )
            gr.Markdown(
                "_Assumed on by llama-server: continuous batching, jinja chat "
                "templates, prompt caching (`--cache-prompt`), flash-attention "
                "auto. The command notes these rather than emitting flags._"
            )

        with gr.Tab("Auto-fit"):
            gr.Markdown(
                "### Auto-fit solvers\n"
                "Uses the current architecture + GPU budget. Run the calculator "
                "first (or load a preset/fetch a GGUF) so the architecture is "
                "populated, then pick a solver."
            )
            with gr.Row():
                af_quant = gr.Dropdown(
                    label="Quant for max-context solver",
                    choices=QUANT_CHOICES, value="Q4_K_M",
                )
                af_maxctx_btn = gr.Button("Max context that fits")
                af_maxctx_out = gr.Markdown("")
            with gr.Row():
                af_bestq_ctx = gr.Number(
                    label="Context for best-quant solver", value=8192, precision=0
                )
                af_bestq_btn = gr.Button("Best quant that fits")
                af_bestq_out = gr.Markdown("")
            with gr.Row():
                af_mingu_ctx = gr.Number(
                    label="Context for min-GPU solver", value=8192, precision=0
                )
                af_mingu_quant = gr.Dropdown(
                    label="Quant for min-GPU solver",
                    choices=QUANT_CHOICES, value="Q4_K_M",
                )
                af_mingu_btn = gr.Button("Min GPUs needed")
                af_mingu_out = gr.Markdown("")

        # --- wiring ---
        def _store_arch(*fields):
            return list(fields)

        arch_inputs_and_state = [*arch_inputs, arch_state]
        # keep arch_state synced whenever arch fields change; also refresh the
        # hybrid-attention badge and the inline MTP-conflict hint so they track
        # manual edits to n_full_attn_layers / head_dim / n_mtp.
        for comp in arch_inputs:
            comp.change(
                fn=_store_arch, inputs=arch_inputs, outputs=arch_state
            )
            comp.change(
                fn=_hybrid_md, inputs=arch_inputs, outputs=hybrid_md,
            )
            comp.change(
                fn=_mtp_hint_md, inputs=arch_inputs,
                outputs=mtp_hint_md,
            )
        # mtp_hint_md only depends on n_mtp, which is in arch_inputs — already
        # covered by the per-component change handlers above (the loop on
        # arch_inputs calls _mtp_hint_md for every arch field edit).

        list_btn.click(
            fn=list_gguf_files, inputs=[repo_id, hf_token],
            outputs=[file_picker, fetch_status],
        )
        fetch_btn.click(
            fn=fetch_arch, inputs=[repo_id, file_picker, hf_token],
            outputs=[*arch_inputs, fetch_status, quant],
        ).then(
            fn=_store_arch, inputs=arch_inputs, outputs=arch_state
        ).then(
            fn=_hybrid_md, inputs=arch_inputs, outputs=hybrid_md
        ).then(
            fn=_mtp_hint_md, inputs=arch_inputs,
            outputs=mtp_hint_md,
        )
        load_preset_btn.click(
            fn=load_preset, inputs=[preset_dd],
            outputs=[*arch_inputs, fetch_status],
        ).then(
            fn=_store_arch, inputs=arch_inputs, outputs=arch_state
        ).then(
            fn=_hybrid_md, inputs=arch_inputs, outputs=hybrid_md
        ).then(
            fn=_mtp_hint_md, inputs=arch_inputs,
            outputs=mtp_hint_md,
        )

        # mmproj picker: populate from the same repo's *.gguf (filter mmproj)
        list_btn.click(
            fn=list_mmproj_files, inputs=[repo_id, hf_token],
            outputs=[mmproj_picker, fetch_status],
        )
        fetch_mmproj_btn.click(
            fn=fetch_mmproj_bytes, inputs=[repo_id, mmproj_picker, hf_token],
            outputs=[mmproj_bytes_box, fetch_status],
        )

        # Live recompute inputs — every input that affects the estimate.
        live_inputs = [
            arch_state, quant, n_ctx, cache_dtype, flash_attn, compute_dtype,
            n_batch, n_ubatch, n_prompt, parallel_slots, parallel_sizing,
            rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
            yarn_beta_fast, yarn_beta_slow,
            gpu_vram_text, split_mode, main_gpu, tensor_split_text,
            safety_margin,
            mtp_cache_dtype,
            spec_type, draft_quant, draft_params, draft_n_layer,
            draft_n_max, draft_n_min, draft_p_min, draft_p_split,
            mmproj_enabled, mmproj_file, mmproj_offload, mmproj_bytes_box,
        ]
        live_outputs = [result_md, cmd_text, graph_ctx, graph_quant, graph_batch]

        compute_btn.click(
            fn=live_compute,
            inputs=live_inputs,
            outputs=live_outputs,
        )

        # Live: fire on every input change (estimate() is microseconds on the
        # cached arch — no network). See PLAN.md "Live phase".
        live_comps = [
            quant, n_ctx, cache_dtype, flash_attn, compute_dtype,
            n_batch, n_ubatch, n_prompt, parallel_slots, parallel_sizing,
            rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
            yarn_beta_fast, yarn_beta_slow,
            gpu_vram_text, split_mode, main_gpu, tensor_split_text,
            safety_margin,
            mtp_cache_dtype,
            spec_type, draft_quant, draft_params, draft_n_layer,
            draft_n_max, draft_n_min, draft_p_min, draft_p_split,
            mmproj_enabled, mmproj_file, mmproj_offload, mmproj_bytes_box,
        ]
        for comp in live_comps:
            comp.change(
                fn=live_compute, inputs=live_inputs, outputs=live_outputs,
            )
        # arch_state changes (fetch/preset) also trigger a recompute.
        arch_state.change(
            fn=live_compute, inputs=live_inputs, outputs=live_outputs,
        )

        # Spec type change → prefill draft_n_max with the per-type default,
        # and set a sane n_min for ngram-mod. Returns gr.update for the two
        # Number components so they update without a full live recompute round.
        def _spec_change_handler(spec_value):
            # spec_value is the bare type (dropdown value), or "none".
            t = spec_value or "none"
            n_max_default = SPEC_DEFAULTS.get(t, 0)
            n_min_default = 8 if t == "ngram-mod" else 0
            return gr.update(value=n_max_default), gr.update(value=n_min_default)

        spec_type.change(
            fn=_spec_change_handler,
            inputs=[spec_type],
            outputs=[draft_n_max, draft_n_min],
        )

        # One-click spec-default (ngram-mod): set the dropdown + n_max + n_min.
        def _spec_default_handler():
            return gr.update(value="ngram-mod"), gr.update(value=64), gr.update(value=8)

        spec_default_btn.click(
            fn=_spec_default_handler,
            inputs=[],
            outputs=[spec_type, draft_n_max, draft_n_min],
        )

        # Auto-fit wiring
        af_maxctx_btn.click(
            fn=run_max_context,
            inputs=[arch_state, af_quant, gpu_vram_text, split_mode, main_gpu,
                    n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype,
                    safety_margin],
            outputs=[af_maxctx_out],
        )
        af_bestq_btn.click(
            fn=run_best_quant,
            inputs=[arch_state, af_bestq_ctx, gpu_vram_text, split_mode, main_gpu,
                    n_batch, n_ubatch, cache_dtype, flash_attn, compute_dtype,
                    safety_margin],
            outputs=[af_bestq_out],
        )
        af_mingu_btn.click(
            fn=run_min_gpu_setup,
            inputs=[arch_state, af_mingu_ctx, af_mingu_quant, gpu_vram_text,
                    split_mode, main_gpu, n_batch, n_ubatch, cache_dtype,
                    flash_attn, compute_dtype, safety_margin],
            outputs=[af_mingu_out],
        )

        # YaRN auto-configure: derive rope_freq_scale from training_ctx + n_ctx
        yarn_auto_btn.click(
            fn=auto_configure_yarn_handler,
            inputs=[arch_state, n_ctx],
            outputs=[rope_freq_scale, yarn_ext_factor, yarn_attn_factor,
                     yarn_beta_fast, yarn_beta_slow, yarn_auto_status],
        )

    return demo


demo = build_ui()


if __name__ == "__main__":
    demo.launch()