"""Minimal GGUF v3 header parser.

Reads only the metadata key/value section (no tensor data) from a GGUF file,
either from a local path or via a ranged HTTP GET against the Hugging Face
resolve URL so we never download the whole model.

GGUF layout (little-endian):
    magic: u32 = 'GGUF'
    version: u32
    tensor_count: u64
    metadata_kv_count: u64
    metadata_kv[]: { name_len: u64, name: bytes, value_type: u32, value: ... }
    tensor_info[]: (not parsed here)

We parse just enough metadata to recover architecture parameters.
"""

from __future__ import annotations

import io
import struct
from dataclasses import dataclass, field

GGUF_MAGIC = 0x46554747  # "GGUF"
GGUF_DEFAULT_VERSION = 3

GGUF_TYPE_UINT8 = 0
GGUF_TYPE_INT8 = 1
GGUF_TYPE_UINT16 = 2
GGUF_TYPE_INT16 = 3
GGUF_TYPE_UINT32 = 4
GGUF_TYPE_INT32 = 5
GGUF_TYPE_FLOAT32 = 6
GGUF_TYPE_BOOL = 7
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9
GGUF_TYPE_UINT64 = 10
GGUF_TYPE_INT64 = 11
GGUF_TYPE_FLOAT64 = 12

_SCALAR_FMT = {
    GGUF_TYPE_UINT8: ("<B", 1),
    GGUF_TYPE_INT8: ("<b", 1),
    GGUF_TYPE_UINT16: ("<H", 2),
    GGUF_TYPE_INT16: ("<h", 2),
    GGUF_TYPE_UINT32: ("<I", 4),
    GGUF_TYPE_INT32: ("<i", 4),
    GGUF_TYPE_FLOAT32: ("<f", 4),
    GGUF_TYPE_BOOL: ("<?", 1),
    GGUF_TYPE_UINT64: ("<Q", 8),
    GGUF_TYPE_INT64: ("<q", 8),
    GGUF_TYPE_FLOAT64: ("<d", 8),
}


@dataclass
class GGUFMetadata:
    raw: dict[str, object] = field(default_factory=dict)
    architecture: str = ""
    name: str = ""
    n_layer: int = 0
    n_embd: int = 0
    n_head: int = 0
    n_head_kv: int = 0
    training_ctx: int = 0
    rope_freq_base: float = 10000.0
    n_expert: int = 0
    n_expert_used: int = 0
    n_mtp: int = 0
    params: int = 0


class _Reader:
    def __init__(self, buf: bytes):
        self.buf = buf
        self.pos = 0

    def eof(self) -> bool:
        return self.pos >= len(self.buf)

    def need(self, n: int):
        if self.pos + n > len(self.buf):
            raise EOFError("buffer exhausted while parsing GGUF header")

    def u32(self) -> int:
        self.need(4)
        v = struct.unpack_from("<I", self.buf, self.pos)[0]
        self.pos += 4
        return v

    def u64(self) -> int:
        self.need(8)
        v = struct.unpack_from("<Q", self.buf, self.pos)[0]
        self.pos += 8
        return v

    def string(self) -> str:
        n = self.u64()
        self.need(n)
        s = self.buf[self.pos:self.pos + n].decode("utf-8", errors="replace")
        self.pos += n
        return s

    def scalar(self, t: int):
        fmt, size = _SCALAR_FMT[t]
        self.need(size)
        v = struct.unpack_from(fmt, self.buf, self.pos)[0]
        self.pos += size
        return v

    def value(self, t: int):
        if t == GGUF_TYPE_STRING:
            return self.string()
        if t in _SCALAR_FMT:
            return self.scalar(t)
        if t == GGUF_TYPE_ARRAY:
            inner = self.u32()
            count = self.u64()
            return [self.value(inner) for _ in range(count)]
        raise ValueError(f"unsupported GGUF value type {t}")


def parse_header_bytes(buf: bytes) -> dict[str, object]:
    """Parse the metadata KV map from a buffer containing the GGUF header."""
    r = _Reader(buf)
    magic = r.u32()
    if magic != GGUF_MAGIC:
        raise ValueError(f"not a GGUF file (magic={magic:#x})")
    version = r.u32()
    if version < 1 or version > 3:
        raise ValueError(f"unsupported GGUF version {version}")
    # tensor_count (u64) - skipped
    r.u64()
    kv_count = r.u64()
    meta: dict[str, object] = {}
    for _ in range(kv_count):
        if r.eof():
            break
        key = r.string()
        t = r.u32()
        try:
            meta[key] = r.value(t)
        except (EOFError, ValueError):
            # ran out of buffered bytes (range read too small) - stop
            break
    return meta


def parse_header_with_tensors(buf: bytes):
    """Parse metadata KV and sum tensor element counts from the header.

    Returns (meta_dict, total_tensor_elems). The tensor-info section follows
    the metadata: each tensor is name(str) + n_dims(u32) + dims[n](u64) +
    dtype(u32) + offset(u64). We sum product(dims) across all tensors. The
    buffer must cover the metadata *and* the full tensor-info section for the
    sum to be correct — tokenizer-heavy GGUFs push the tensor section past
    10 MiB, so callers should read >= ~32 MiB. If the read is too small to
    reach the tensor section we return -1 (unreachable) so callers can tell
    "incomplete" from "zero tensors".

    Used as a fallback for `params` when general.parameter_count is absent.
    """
    r = _Reader(buf)
    magic = r.u32()
    if magic != GGUF_MAGIC:
        raise ValueError(f"not a GGUF file (magic={magic:#x})")
    version = r.u32()
    if version < 1 or version > 3:
        raise ValueError(f"unsupported GGUF version {version}")
    tensor_count = r.u64()
    kv_count = r.u64()
    meta: dict[str, object] = {}
    for _ in range(kv_count):
        if r.eof():
            return meta, 0
        key = r.string()
        t = r.u32()
        try:
            meta[key] = r.value(t)
        except (EOFError, ValueError):
            # ran out of buffer in metadata; tensor section unreachable
            return meta, 0
    total = 0
    parsed = 0
    for _ in range(tensor_count):
        if r.eof():
            break
        try:
            r.string()                 # tensor name
            n_dims = r.u32()
            elems = 1
            for _ in range(n_dims):
                elems *= r.u64()       # dims multiply, not add
            total += elems
            r.u32()                    # dtype
            r.u64()                    # data offset
            parsed += 1
        except (EOFError, ValueError):
            break
    # If we couldn't reach any tensor infos, signal failure with -1 so callers
    # can distinguish "unreachable" from "0 tensors".
    if parsed == 0 and tensor_count > 0:
        return meta, -1
    return meta, total


def _coerce_int(v: object) -> int:
    if isinstance(v, bool):
        return int(v)
    if isinstance(v, (int, float)):
        return int(v)
    return 0


def _coerce_float(v: object) -> float:
    if isinstance(v, (int, float)):
        return float(v)
    return 0.0


def metadata_to_arch(
    meta: dict[str, object], total_tensor_elems: int | None = None
) -> GGUFMetadata:
    """Project raw GGUF metadata into the architecture fields we need.

    `total_tensor_elems` (from parse_header_with_tensors) is used as a fallback
    for `params` when general.parameter_count is absent. A negative value means
    the tensor section was unreachable (range read too small); treat as no
    fallback.
    """
    arch_name = str(meta.get("general.architecture", "") or "")
    p = f"{arch_name}." if arch_name else ""

    def g(key: str, default=0):
        return meta.get(f"{p}{key}", meta.get(key, default))

    m = GGUFMetadata(raw=meta, architecture=arch_name)
    m.name = str(meta.get("general.name", "") or "")
    m.n_layer = _coerce_int(g("block_count"))
    m.n_embd = _coerce_int(g("embedding_length"))
    m.n_head = _coerce_int(g("attention.head_count"))
    n_head_kv = g("attention.head_count_kv")
    m.n_head_kv = _coerce_int(n_head_kv) if n_head_kv is not None else m.n_head
    if m.n_head_kv == 0:
        m.n_head_kv = m.n_head
    m.training_ctx = _coerce_int(g("context_length"))
    m.rope_freq_base = _coerce_float(g("rope.freq_base", 10000.0))
    m.n_expert = _coerce_int(g("expert_count"))
    m.n_expert_used = _coerce_int(g("expert_used_count"))
    # MTP: not always in GGUF metadata; leave 0 unless present
    m.n_mtp = _coerce_int(g("mtp.count") or g("n_mtp"))
    # parameter count: prefer general.parameter_count; fall back to summed
    # tensor element counts when the GGUF omits it (many GGUFs do).
    m.params = _coerce_int(meta.get("general.parameter_count", 0))
    if not m.params and total_tensor_elems and total_tensor_elems > 0:
        m.params = int(total_tensor_elems)
    return m


def parse_local_file(path: str, max_bytes: int = 1 << 25) -> GGUFMetadata:
    with open(path, "rb") as f:
        buf = f.read(max_bytes)
    meta, total = parse_header_with_tensors(buf)
    return metadata_to_arch(meta, total)


def parse_hf_range(
    repo_id: str,
    filename: str,
    *,
    max_bytes: int = 1 << 25,
    token: str | None = None,
    revision: str = "main",
    timeout: float = 30.0,
) -> GGUFMetadata:
    """Range-read the first `max_bytes` of a GGUF file from the HF Hub.

    Uses a HTTP Range request against the resolve URL so we never download the
    full model. Requires the `requests` package. Reads up to 32 MiB of the
    header by default (enough to cover tokenizer-heavy metadata plus the full
    tensor-info section, so we can sum tensor element counts as a params
    fallback when general.parameter_count is absent).
    """
    import requests

    url = f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename}"
    headers = {"Range": f"bytes=0-{max_bytes - 1}"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    resp = requests.get(url, headers=headers, timeout=timeout, stream=True)
    resp.raise_for_status()
    buf = resp.content
    meta, total = parse_header_with_tensors(buf)
    # If the read was too small to reach the tensor section, retry once with a
    # larger read so the params fallback can work for tokenizer-heavy GGUFs.
    if not meta.get("general.parameter_count") and total == -1:
        big = 1 << 26  # 64 MiB
        headers["Range"] = f"bytes=0-{big - 1}"
        resp = requests.get(url, headers=headers, timeout=timeout, stream=True)
        resp.raise_for_status()
        meta, total = parse_header_with_tensors(resp.content)
    return metadata_to_arch(meta, total)