"""GGUF metadata + tensor-table reader (stdlib only).

Reads the header only (metadata + tensor infos); never touches tensor data, so it is fast enough to
run at picker time on multi-GB files.
"""

from __future__ import annotations

import logging
import re
import struct
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger(__name__)

_GGUF_MAGIC = b"GGUF"

# Split GGUF naming: "<stem>-00001-of-00003.gguf"; the part suffix is not part of the model id.
SPLIT_PART_RE = re.compile(r"-(\d{5})-of-(\d{5})\.gguf$")
_PART_SUFFIX_RE = re.compile(r"-\d{5}-of-\d{5}$")
# Same tensor selection as context_policy's per-block FFN -ot override.
_FFN_WEIGHT = re.compile(r"blk\.(\d+)\.ffn_.*\.weight")
# A block with an SSM conv filter keeps recurrent state (Mamba2 and gated delta-net alike).
_SSM_CONV = re.compile(r"blk\.(\d+)\.ssm_conv1d\.weight")

# Tensors llama.cpp leaves in the file and reads row by row on demand instead of loading them
# (create_tensor's TENSOR_READ_LAZY, marked per architecture in the engine's model code). Under the
# engine's default --lazy-mode auto, only a marked tensor larger than 4 GiB is read this way;
# a smaller one loads like any other weight.
_LAZY_READ_TENSORS = {
    "qwen4exp": frozenset({"per_layer_token_embd.weight"}),
    "gemma4": frozenset({"per_layer_token_embd.weight"}),
}
_LAZY_READ_CANDIDATES = frozenset().union(*_LAZY_READ_TENSORS.values())
_LAZY_READ_MIN_BYTES = 4 << 30


def model_id_from_stem(stem: str) -> str:
    """Model id from a GGUF file stem (split-part suffix stripped)."""
    return _PART_SUFFIX_RE.sub("", stem)


# ggml tensor type sizes: type_id -> (block_bytes, block_elems). IQ-family verified against
# ggml-common.h.
_GGML_TYPE_SIZES = {
    0: (4, 1), 1: (2, 1), 2: (18, 32), 3: (20, 32), 6: (22, 32), 7: (24, 32),
    8: (34, 32), 9: (36, 32), 10: (84, 256), 11: (110, 256), 12: (144, 256),
    13: (176, 256), 14: (210, 256), 15: (292, 256), 16: (66, 256),
    17: (74, 256), 18: (98, 256), 19: (50, 256), 20: (18, 32),
    21: (110, 256), 22: (82, 256), 23: (136, 256), 24: (1, 1), 25: (2, 1),
    26: (4, 1), 27: (8, 1), 28: (8, 1), 29: (56, 256), 30: (2, 1),
    39: (17, 32),  # MXFP4 — 32 elements per 17-byte block (gpt-oss family)
    40: (36, 64),  # NVFP4 — 4 E4M3 sub-block scales + 32 bytes of E2M1 (Nemotron 3.5 Bolt)
    41: (18, 128),  # Q1_0
    42: (18, 64),  # Q2_0
}

# GGUF metadata value types -> struct format; STRING (8) and ARRAY (9) are variable-length.
_V_STRING, _V_ARRAY = 8, 9
_SCALAR_FMT = {
    0: "<B", 1: "<b", 2: "<H", 3: "<h",       # uint8 int8 uint16 int16
    4: "<I", 5: "<i", 6: "<f", 7: "<?",       # uint32 int32 float32 bool
    10: "<Q", 11: "<q", 12: "<d",             # uint64 int64 float64
}

# general.sampling.* metadata key -> preset INI key.
_SAMPLING_INI_KEY = {"temp": "temp", "temperature": "temp", "top_p": "top-p",
                     "top_k": "top-k", "min_p": "min-p",
                     "repeat_penalty": "repeat-penalty",
                     "presence_penalty": "presence-penalty"}


@dataclass
class GGUFHeader:
    path: str
    version: int
    metadata: dict = field(default_factory=dict)
    n_tensors: int = 0
    tensor_bytes: int = 0          # exact sum over the tensor table
    embd_table_bytes: int = 0      # token_embd.weight (duplicated host-side when fully offloaded)
    # block index -> bytes of that block's FFN weights (the tensors a `blk\.N\.ffn_.*\.weight`
    # -ot override moves), so spill placement can move only as many blocks as it needs.
    ffn_block_bytes: dict[int, int] = field(default_factory=dict)
    # Sizes of the tensors some architecture reads lazily, by name. ``lazy_bytes`` applies this
    # model's architecture and the size threshold.
    lazy_candidate_bytes: dict[str, int] = field(default_factory=dict)
    # Blocks that keep recurrent state. Counted from the tensor table: hybrids like Nemotron-H
    # also have MoE-only blocks without KV heads, which hold none.
    recurrent_blocks: frozenset[int] = frozenset()

    # ── typed accessors ──────────────────────────────────────

    @property
    def architecture(self) -> str:
        return str(self.metadata.get("general.architecture", ""))

    @property
    def lazy_bytes(self) -> int:
        """Bytes of weights llama.cpp keeps on disk and reads row by row instead of loading."""
        marked = _LAZY_READ_TENSORS.get(self.architecture, frozenset())
        return sum(nbytes for name, nbytes in self.lazy_candidate_bytes.items()
                   if name in marked and nbytes > _LAZY_READ_MIN_BYTES)

    def _arch_key(self, suffix: str):
        return self.metadata.get(f"{self.architecture}.{suffix}")

    @property
    def recurrent_state_bytes(self) -> int:
        """One copy of one sequence's recurrent state: an f32 conv window plus the SSM state, per
        block that keeps state. llama.cpp holds one copy per slot per speculative draft position.
        0 when the file declares no SSM shape."""
        def ssm(key: str) -> int:
            return int(self._arch_key(f"ssm.{key}") or 0)

        inner, state = ssm("inner_size"), ssm("state_size")
        if not (self.recurrent_blocks and inner and state):
            return 0
        conv_dim = inner + 2 * ssm("group_count") * state
        per_block = 4 * (inner * state + max(0, ssm("conv_kernel") - 1) * conv_dim)
        return len(self.recurrent_blocks) * per_block

    def _arch_int(suffix: str, doc: str = ""):
        return property(lambda self: int(self._arch_key(suffix) or 0), doc=doc)

    n_layer = _arch_int("block_count")
    n_ctx_train = _arch_int("context_length")
    n_embd = _arch_int("embedding_length")
    sliding_window = _arch_int("attention.sliding_window")
    expert_count = _arch_int("expert_count")
    full_attention_interval = _arch_int(
        "full_attention_interval",
        "GDN-hybrid discriminator (qwen35 family): every Nth layer is full attention, the rest "
        "are linear/recurrent. 0 = not present.")
    key_length_swa = _arch_int(
        "attention.key_length_swa",
        "Per-token key size for sliding-window layers when it differs from the global "
        "attention.key_length (e.g. gemma3/gemma4). 0 = not present.")
    value_length_swa = _arch_int(
        "attention.value_length_swa",
        "Per-token value size for sliding-window layers when it differs from the global "
        "attention.value_length. 0 = not present.")
    del _arch_int

    @property
    def sliding_window_pattern(self) -> list[int] | None:
        """Per-layer SWA pattern declared by the file itself: a truthy entry marks a
        sliding-window layer, falsy marks global. None when the file doesn't declare the per-layer
        array form (older GGUFs, architectures without per-layer SWA metadata, or files that use
        the scalar period form instead — see `sliding_window_pattern_period`)."""
        v = self._arch_key("attention.sliding_window_pattern")
        return [int(x) for x in v] if isinstance(v, list) else None

    @property
    def sliding_window_pattern_period(self) -> int:
        """Scalar SWA period declared by the file: llama.cpp also permits
        `attention.sliding_window_pattern` as a single integer N (every Nth layer is full
        attention, e.g. Gemma-family writers) rather than a per-layer array. 0 when absent or when
        the file uses the array form instead — callers should fall back to a coarser signal."""
        v = self._arch_key("attention.sliding_window_pattern")
        return int(v) if isinstance(v, int) and not isinstance(v, bool) else 0

    @property
    def n_vocab(self) -> int:
        """Vocabulary size (prices the GPU logits buffers): vocab_size metadata when present, else
        the tokenizer list length."""
        v = self._arch_key("vocab_size")
        if v:
            return int(v)
        toks = self.metadata.get("tokenizer.ggml.tokens")
        return len(toks) if isinstance(toks, list) else 0

    @property
    def sampling_defaults(self) -> dict:
        """Upstream's recommended sampling as preset INI keys, when the file carries it.

        Publishers bake ``general.sampling.*`` keys into the GGUF (llama-server reads them as that
        model's defaults), so the file is the source of truth — it ships with the download and
        updates with every re-upload, no catalog needed. Empty if absent.
        """
        out = {}
        for key, value in self.metadata.items():
            if not key.startswith("general.sampling."):
                continue
            name = _SAMPLING_INI_KEY.get(key.rsplit(".", 1)[-1])
            if name is not None and isinstance(value, (int, float)):
                num = round(float(value), 4)
                out[name] = str(int(num)) if num == int(num) else str(num)
        return out

    @property
    def n_head(self) -> int:
        v = self._arch_key("attention.head_count")
        if isinstance(v, list):
            return int(max(v))
        return int(v or 0)

    def head_counts_kv(self) -> list[int]:
        """Per-layer KV head counts; 0 marks a recurrent/linear layer (n_head_kv == 0).

        Three GGUF shapes: a per-layer array (nemotron_h_moe) is used as-is; a scalar plus
        ``full_attention_interval`` (qwen35) applies to every N-th layer (1-indexed) and is zero
        elsewhere — pricing all layers as attention was a 4x overestimate; a plain scalar (dense)
        broadcasts to every layer.
        """
        v = self._arch_key("attention.head_count_kv")
        if isinstance(v, list):
            return [int(x) for x in v]
        scalar = int(v or 0)
        interval = self.full_attention_interval
        if interval > 1:
            return [scalar if (i + 1) % interval == 0 else 0
                    for i in range(self.n_layer)]
        return [scalar] * self.n_layer

    @property
    def head_dim_k(self) -> int:
        v = self._arch_key("attention.key_length")
        if v:
            return int(v)
        return self.n_embd // self.n_head if self.n_head else 0

    @property
    def head_dim_v(self) -> int:
        v = self._arch_key("attention.value_length")
        if v:
            return int(v)
        return self.head_dim_k


def split_parts(path: Path) -> list[Path] | None:
    """Every on-disk part of the split ``path`` belongs to, first part first; None when ``path`` is
    not a split member or no other part is present.

    A split is priced as the set, never as one file: publishers lay shards out so the first part
    can hold little more than metadata while the bulk sits in the later ones, so reading one part
    prices the model at whatever fraction of its weights that part happens to hold."""
    m = SPLIT_PART_RE.search(path.name)
    if m is None:
        return None
    stem, total = path.name[: m.start()], int(m.group(2))
    parts = [p for p in (path.with_name(f"{stem}-{i:05d}-of-{total:05d}.gguf")
                         for i in range(1, total + 1)) if p.is_file()]
    return parts if len(parts) > 1 else None


def _read_part(path: Path) -> GGUFHeader:
    """One file's own header: metadata and that file's tensor table."""
    def read(f, fmt: str):
        # A header cut short raises ValueError, which every caller already treats as "skip this
        # model"; struct.error would escape them and take the whole preset pass down.
        size = struct.calcsize(fmt)
        data = f.read(size)
        if len(data) != size:
            raise ValueError(f"truncated GGUF header: {path}")
        return struct.unpack(fmt, data)

    def read_str(f) -> str:
        (n,) = read(f, "<Q")
        return f.read(n).decode("utf-8", errors="replace")

    def read_value(f, vtype: int):
        if vtype == _V_STRING:
            return read_str(f)
        if vtype == _V_ARRAY:
            etype, n = read(f, "<IQ")
            return [read_value(f, etype) for _ in range(n)]
        return read(f, _SCALAR_FMT[vtype])[0]

    with open(path, "rb") as f:
        if f.read(4) != _GGUF_MAGIC:
            raise ValueError(f"not a GGUF file: {path}")
        version, n_tensors, n_kv = read(f, "<IQQ")

        metadata: dict = {}
        for _ in range(n_kv):
            key = read_str(f)
            (vtype,) = read(f, "<I")
            metadata[key] = read_value(f, vtype)

        tensor_bytes = 0
        embd_bytes = 0
        ffn_block_bytes: dict[int, int] = {}
        lazy_candidate_bytes: dict[str, int] = {}
        recurrent_blocks: set[int] = set()
        for _ in range(n_tensors):
            name = read_str(f)
            (n_dims,) = read(f, "<I")
            dims = read(f, f"<{n_dims}Q")
            (ttype,) = read(f, "<I")
            f.read(8)  # offset
            size = _GGML_TYPE_SIZES.get(ttype)
            if size is None:
                raise ValueError(f"unknown ggml tensor type {ttype} in {path}")
            block_bytes, block_elems = size
            elems = 1
            for d in dims:
                elems *= d
            nbytes = (elems // block_elems) * block_bytes
            tensor_bytes += nbytes
            if name == "token_embd.weight":
                embd_bytes = nbytes
            elif m := _FFN_WEIGHT.match(name):
                block = int(m.group(1))
                ffn_block_bytes[block] = ffn_block_bytes.get(block, 0) + nbytes
            elif name in _LAZY_READ_CANDIDATES:
                lazy_candidate_bytes[name] = nbytes
            elif m := _SSM_CONV.match(name):
                recurrent_blocks.add(int(m.group(1)))

    return GGUFHeader(path=str(path), version=version, metadata=metadata,
                      n_tensors=n_tensors, tensor_bytes=tensor_bytes,
                      embd_table_bytes=embd_bytes, ffn_block_bytes=ffn_block_bytes,
                      lazy_candidate_bytes=lazy_candidate_bytes,
                      recurrent_blocks=frozenset(recurrent_blocks))


def read_gguf_header(path: str | Path) -> GGUFHeader:
    """Header for a MODEL, not for a file: a split GGUF is priced as the sum of its parts.

    Weights, the host-side embedding-table duplicate and the per-block FFN map are all summed
    across the shards on disk, because a split is a layout choice, not a smaller model — pricing
    part 1 alone underprices every model whose first shard is a metadata stub (the Hugging Face
    layout), which then makes the physics check and the residency cap admit giants the card cannot
    hold. Architecture metadata (block count, train context, the per-layer SWA pattern, vocab) is
    taken from the first part: gguf-split writes the model's metadata there only, and every later
    part carries just its ``split.*`` keys.

    A part that has gone missing or become unreadable is skipped rather than fatal: a half-arrived
    split prices at what is actually on disk. Refusing it is ``staged_in(require_complete=True)``'s
    job — an incomplete split is never servable, it is only underpriced here."""
    path = Path(path)
    parts = split_parts(path)
    if parts is None:
        return _read_part(path)

    first = _read_part(parts[0])
    readable = [first]
    for part in parts[1:]:
        try:
            readable.append(_read_part(part))
        except (ValueError, OSError) as exc:
            logger.debug("split part unreadable %s: %s", part.name, exc)
    if len(readable) == 1:
        return first

    ffn_block_bytes: dict[int, int] = {}
    for header in readable:
        for block, nbytes in header.ffn_block_bytes.items():
            ffn_block_bytes[block] = ffn_block_bytes.get(block, 0) + nbytes
    return GGUFHeader(
        # The path stays the part the caller named: it is the model id source and the preset's
        # ``model`` key, and llama.cpp resolves the split from any one of its members.
        path=str(path), version=first.version, metadata=first.metadata,
        n_tensors=sum(h.n_tensors for h in readable),
        tensor_bytes=sum(h.tensor_bytes for h in readable),
        embd_table_bytes=sum(h.embd_table_bytes for h in readable),
        ffn_block_bytes=ffn_block_bytes,
        lazy_candidate_bytes={name: nbytes for h in readable
                              for name, nbytes in h.lazy_candidate_bytes.items()},
        recurrent_blocks=frozenset().union(*(h.recurrent_blocks for h in readable)))