"""Speculative-decoding draft model memory.

A draft model (llama.cpp `-md`) runs alongside the target and contributes its
own weights and KV cache. Supported spec types (see docs/speculative.md):
  draft        — a standalone small draft model
  draft-eagle3  — EAGLE-3 single-layer draft (shares the target tokenizer)
  draft-dflash  — DFlash block-diffusion draft
  none         — no draft model

For VRAM purposes the draft model's cost is its weights + its own KV cache.
EAGLE-3 drafts are typically a single layer; standalone drafts have their own
n_layer. The draft usually rides on the main GPU alongside the target, which
the caller handles by adding draft_bytes to that GPU's used.
"""

from __future__ import annotations

from dataclasses import dataclass

from .quant import weight_bytes, QUANT_BPW
from .kv import kv_cache_bytes, cache_dtype_bytes


SPEC_TYPES = (
    "none",
    # weighted (draft model: weights + own KV)
    "draft-simple",
    "draft-eagle3",
    "draft-mtp",
    "draft-dflash",
    # weightless (statistical / n-gram; no draft model, ~zero extra VRAM)
    "ngram-simple",
    "ngram-map-k",
    "ngram-map-k4v",
    "ngram-mod",
    "ngram-cache",
)

# Spec types that need no draft model — pure statistical speculation. These
# have ~zero extra VRAM cost (no draft weights, no draft KV), so they are
# auto-offered when headroom exists without an "acknowledge the tradeoff"
# gate. See PLAN.md "Speculative decoding".
WEIGHTLESS_SPEC_TYPES = (
    "ngram-simple",
    "ngram-map-k",
    "ngram-map-k4v",
    "ngram-mod",
    "ngram-cache",
)

# Spec types backed by a draft model (-md) with its own weights + KV.
WEIGHTED_SPEC_TYPES = ("draft-simple", "draft-eagle3", "draft-mtp", "draft-dflash")


# Per-type default --spec-draft-n-max (from docs/speculative.md). The UI uses
# these to prefill draft_n_max when the user picks a spec type.
SPEC_DEFAULTS: dict[str, int] = {
    "draft-simple": 3,
    "draft-eagle3": 3,
    "draft-dflash": 15,   # clamped to the draft's trained block size
    "draft-mtp": 3,
    "ngram-simple": 64,
    "ngram-map-k": 64,
    "ngram-map-k4v": 64,
    "ngram-mod": 64,
    "ngram-cache": 0,
}

# Family grouping for the UI (3 buckets). Weightless = no draft model, ~0 VRAM.
SPEC_FAMILY: dict[str, tuple[str, ...]] = {
    "Draft model (weighted, needs -md)": (
        "draft-simple", "draft-eagle3", "draft-dflash",
    ),
    "n-gram (weightless, ~0 VRAM)": (
        "ngram-simple", "ngram-map-k", "ngram-map-k4v", "ngram-mod", "ngram-cache",
    ),
    "MTP-from-target (weighted)": ("draft-mtp",),
}


def is_weightless(spec_type: str) -> bool:
    """True if this spec type needs no draft model (n-gram family)."""
    return spec_type in WEIGHTLESS_SPEC_TYPES


def is_weighted(spec_type: str) -> bool:
    """True if this spec type needs a draft model (-md)."""
    return spec_type in WEIGHTED_SPEC_TYPES


@dataclass
class DraftInputs:
    spec_type: str = "none"
    quant: str = "Q4_K_M"
    params: int = 0          # draft parameter count
    n_layer: int = 1        # draft layer count (1 for EAGLE-3 by default)
    n_ctx: int | None = None  # defaults to target n_ctx
    cache_dtype: str = "f16"
    n_max: int = 0          # --spec-draft-n-max (informational)
    n_min: int = 0          # --spec-draft-n-min (ngram-mod wants > 0, PR #19164)
    p_min: float = 0.0
    p_split: float = 0.0

    @property
    def enabled(self) -> bool:
        # Weightless (n-gram) spec types need no draft model: enabled as long
        # as a type is chosen. Weighted types need params > 0.
        if self.spec_type == "none":
            return False
        if is_weightless(self.spec_type):
            return True
        return self.params > 0


@dataclass
class DraftBreakdown:
    weights_bytes: float = 0.0
    kv_bytes: float = 0.0
    total_bytes: float = 0.0
    enabled: bool = False


def draft_bytes(
    *,
    draft: DraftInputs,
    target_n_ctx: int,
    target_n_embd: int = 0,
    target_n_head: int = 0,
    target_n_head_kv: int = 0,
) -> DraftBreakdown:
    """Return (weights, kv) bytes for a draft model.

    Reuses weight_bytes and kv_cache_bytes. The draft KV uses the draft's own
    n_head_kv; if unset we fall back to the target's (EAGLE-3 shares the
    target's attention geometry). n_embd likewise defaults to the target's.
    """
    if not draft.enabled:
        return DraftBreakdown(enabled=False)

    # Weightless spec types (n-gram family) have no draft model: zero VRAM.
    if is_weightless(draft.spec_type):
        return DraftBreakdown(enabled=True)

    w = weight_bytes(draft.params, draft.quant)

    n_embd = draft.n_embd if getattr(draft, "n_embd", 0) else target_n_embd
    n_head = draft.n_head if getattr(draft, "n_head", 0) else target_n_head
    n_head_kv = (
        draft.n_head_kv if getattr(draft, "n_head_kv", 0) else target_n_head_kv
    )
    if n_head_kv == 0:
        n_head_kv = n_head
    n_ctx = draft.n_ctx if draft.n_ctx is not None else target_n_ctx

    # kv_cache_bytes raises on an unknown cache dtype; guard.
    try:
        kv = kv_cache_bytes(
            n_layer=max(1, draft.n_layer),
            n_embd=n_embd,
            n_head=max(1, n_head),
            n_head_kv=max(1, n_head_kv),
            n_ctx=n_ctx,
            cache_dtype=draft.cache_dtype,
        )
    except ValueError:
        kv = 0.0

    return DraftBreakdown(
        weights_bytes=w, kv_bytes=kv, total_bytes=w + kv, enabled=True,
    )