"""
GGUF-Tokenizer-Fingerprint — liest die Tokenizer-Identität direkt aus dem
GGUF-Header (ohne das Modell zu laden), um zu entscheiden, ob ein Draft-Modell
**vocab-kompatibel** mit einem Ziel-Modell ist (Voraussetzung für Speculative
Decoding in llama.cpp — sonst: "draft model vocab type must match target").

Wir lesen nur die Metadaten-KV-Sektion am Dateianfang und brechen ab, sobald
`tokenizer.ggml.tokens` erreicht ist (dessen Länge = n_vocab). model+pre+n_vocab
identifizieren den Tokenizer eindeutig genug, um die in der Praxis relevanten
Fälle zu unterscheiden (Qwen2.5 vs Qwen3 vs Qwen3.6 etc.). Die llama.cpp-Prüfung
beim Laden bleibt der letzte Schiedsrichter.
"""

import hashlib
import struct
from functools import lru_cache

# GGUF value types (https://github.com/ggml-org/ggml/blob/master/docs/gguf.md)
_T_UINT8, _T_INT8, _T_UINT16, _T_INT16, _T_UINT32, _T_INT32, _T_FLOAT32, \
    _T_BOOL, _T_STRING, _T_ARRAY, _T_UINT64, _T_INT64, _T_FLOAT64 = range(13)

_SCALAR_FMT = {
    _T_UINT8: "<B", _T_INT8: "<b", _T_UINT16: "<H", _T_INT16: "<h",
    _T_UINT32: "<I", _T_INT32: "<i", _T_FLOAT32: "<f", _T_BOOL: "<?",
    _T_UINT64: "<Q", _T_INT64: "<q", _T_FLOAT64: "<d",
}
_SCALAR_SIZE = {t: struct.calcsize(f) for t, f in _SCALAR_FMT.items()}

_WANT_STRINGS = {"tokenizer.ggml.model", "tokenizer.ggml.pre", "general.architecture"}


class _Reader:
    def __init__(self, f):
        self.f = f

    def read(self, n: int) -> bytes:
        b = self.f.read(n)
        if len(b) != n:
            raise EOFError("unerwartetes Dateiende beim GGUF-Parsen")
        return b

    def u32(self) -> int:
        return struct.unpack("<I", self.read(4))[0]

    def u64(self) -> int:
        return struct.unpack("<Q", self.read(8))[0]

    def gstr(self) -> str:
        n = self.u64()
        return self.read(n).decode("utf-8", "replace")

    def scalar(self, vtype: int):
        """Liest einen Skalar-Wert (für die Architektur-Metadaten). None bei Nicht-Skalar."""
        fmt = _SCALAR_FMT.get(vtype)
        if not fmt:
            self.skip_value(vtype)
            return None
        return struct.unpack(fmt, self.read(_SCALAR_SIZE[vtype]))[0]

    def skip_value(self, vtype: int) -> None:
        """Liest einen Wert und verwirft ihn (um den Datei-Pointer korrekt
        weiterzuschieben). Arrays werden elementweise konsumiert."""
        if vtype == _T_STRING:
            self.f.seek(self.u64(), 1)
        elif vtype in _SCALAR_SIZE:
            self.f.seek(_SCALAR_SIZE[vtype], 1)
        elif vtype == _T_ARRAY:
            etype = self.u32()
            count = self.u64()
            if etype == _T_STRING:
                for _ in range(count):
                    self.f.seek(self.u64(), 1)
            elif etype in _SCALAR_SIZE:
                self.f.seek(_SCALAR_SIZE[etype] * count, 1)
            else:
                raise ValueError(f"unbekannter Array-Elementtyp {etype}")
        else:
            raise ValueError(f"unbekannter GGUF-Wertetyp {vtype}")


def _read_fingerprint(path: str) -> dict | None:
    """Liest model/pre/n_vocab aus dem GGUF-Header. None bei Fehler/kein GGUF."""
    try:
        with open(path, "rb") as fh:
            r = _Reader(fh)
            if r.read(4) != b"GGUF":
                return None
            r.u32()                 # version
            r.u64()                 # tensor_count
            kv_count = r.u64()
            fp: dict = {"model": None, "pre": None, "arch": None, "n_vocab": None,
                        "tokens_sha": None}
            for _ in range(kv_count):
                key = r.gstr()
                vtype = r.u32()
                if key == "tokenizer.ggml.tokens" and vtype == _T_ARRAY:
                    etype = r.u32()
                    count = r.u64()
                    fp["n_vocab"] = count
                    if etype != _T_STRING:
                        return None
                    # ECHTE Vocab-Identität: sha256 über die tatsächliche Token-Liste
                    # (familienunabhängig — funktioniert für Qwen, Llama, Mistral, …).
                    h = hashlib.sha256()
                    h.update(count.to_bytes(8, "little"))
                    for _ in range(count):
                        n = r.u64()
                        h.update(r.read(n))
                    fp["tokens_sha"] = h.hexdigest()
                    # model/pre kommen vor tokens → wir haben alles. Abbrechen.
                    break
                if key in _WANT_STRINGS and vtype == _T_STRING:
                    val = r.gstr()
                    if key == "tokenizer.ggml.model":
                        fp["model"] = val
                    elif key == "tokenizer.ggml.pre":
                        fp["pre"] = val
                    else:
                        fp["arch"] = val
                else:
                    r.skip_value(vtype)
            if fp["model"] is None and fp["n_vocab"] is None:
                return None
            return fp
    except (OSError, EOFError, ValueError, struct.error):
        return None


@lru_cache(maxsize=256)
def _cached(path: str, mtime: float, size: int) -> tuple | None:
    fp = _read_fingerprint(path)
    if fp is None:
        return None
    return (fp.get("model"), fp.get("pre"), fp.get("n_vocab"), fp.get("arch"), fp.get("tokens_sha"))


def fingerprint(path: str) -> dict | None:
    """Tokenizer-Fingerprint eines GGUF (gecacht nach Pfad+mtime+size).
    Returns dict(model, pre, n_vocab, arch, tokens_sha) oder None wenn nicht lesbar."""
    import os
    try:
        st = os.stat(path)
    except OSError:
        return None
    t = _cached(path, st.st_mtime, st.st_size)
    if t is None:
        return None
    return {"model": t[0], "pre": t[1], "n_vocab": t[2], "arch": t[3], "tokens_sha": t[4]}


def vocab_key(path: str) -> tuple | None:
    """ECHTER Vergleichsschlüssel für Vocab-Kompatibilität: (model, pre, n_vocab, sha256
    der vollständigen Token-Liste). Vergleicht den TATSÄCHLICHEN Vokabular-Inhalt, nicht
    nur Metadaten — familienunabhängig (Qwen, Llama, Mistral, …). Genau diese Identität
    verlangt llama.cpp für Speculative Decoding."""
    fp = fingerprint(path)
    if not fp or fp["n_vocab"] is None or not fp.get("tokens_sha"):
        return None
    return (fp["model"], fp["pre"], fp["n_vocab"], fp["tokens_sha"])


def compatible(target_path: str, draft_path: str) -> bool | None:
    """True/False ob draft vocab-kompatibel zum target ist. None = unbestimmbar
    (eine Datei nicht lesbar) → UI behandelt das als 'nicht bestätigt'."""
    a = vocab_key(target_path)
    b = vocab_key(draft_path)
    if a is None or b is None:
        return None
    return a == b


# ── Architektur-Metadaten für EHRLICHE KV-Cache-Größen ──────────────────────────────
# Der KV-Cache hängt an (Layer × KV-Heads × Head-Dim), NICHT an den Gesamt-Parametern.
# Bei MoE (z.B. Qwen3.6-35B-A3B) ist das entscheidend: die alte params-basierte Schätzung
# überschätzte grob (aktive vs. gesamte Params + GQA), reale KV liest man direkt hier.
# Schlüssel sind arch-präfixiert ('qwen3moe.block_count', 'llama.attention.head_count_kv' …),
# gegen echte GGUFs verifiziert. Wir sammeln die gewünschten Skalar-Schlüssel per Suffix.
_ARCH_WANT = (
    ".block_count", ".attention.head_count_kv", ".attention.head_count",
    ".attention.key_length", ".attention.value_length", ".embedding_length",
    ".context_length",
)


def _read_arch_meta(path: str) -> dict | None:
    try:
        with open(path, "rb") as fh:
            r = _Reader(fh)
            if r.read(4) != b"GGUF":
                return None
            r.u32()                 # version
            r.u64()                 # tensor_count
            kv_count = r.u64()
            raw: dict = {}
            arch = None
            for _ in range(kv_count):
                key = r.gstr()
                vtype = r.u32()
                if key == "general.architecture" and vtype == _T_STRING:
                    arch = r.gstr()
                    continue
                if key == "tokenizer.ggml.tokens":
                    break            # Arch-Metadaten stehen davor → fertig, Rest überspringen
                suf = next((s for s in _ARCH_WANT if key.endswith(s)), None)
                if suf is not None and vtype in _SCALAR_SIZE:
                    raw[suf] = r.scalar(vtype)
                else:
                    r.skip_value(vtype)
            n_layers = raw.get(".block_count")
            n_head = raw.get(".attention.head_count")
            n_head_kv = raw.get(".attention.head_count_kv") or n_head   # GQA fehlt → MHA
            n_embd = raw.get(".embedding_length")
            hd_k = raw.get(".attention.key_length") \
                or (int(n_embd / n_head) if (n_embd and n_head) else None)
            hd_v = raw.get(".attention.value_length") or hd_k
            if not (n_layers and n_head_kv and hd_k and hd_v):
                return None          # unvollständig → Aufrufer nutzt Heuristik-Fallback
            return {"arch": arch, "n_layers": int(n_layers), "n_head_kv": int(n_head_kv),
                    "head_dim_k": int(hd_k), "head_dim_v": int(hd_v),
                    "n_ctx_train": int(raw[".context_length"]) if raw.get(".context_length") else None}
    except (OSError, EOFError, ValueError, struct.error):
        return None


@lru_cache(maxsize=128)
def _arch_cached(path: str, mtime: float, size: int) -> dict | None:
    return _read_arch_meta(path)


def arch_meta(path: str) -> dict | None:
    """Architektur-Metadaten eines GGUF (gecacht nach Pfad+mtime+size):
    {arch, n_layers, n_head_kv, head_dim_k, head_dim_v, n_ctx_train}. None wenn nicht lesbar
    oder unvollständig."""
    import os
    try:
        st = os.stat(path)
    except OSError:
        return None
    return _arch_cached(path, st.st_mtime, st.st_size)


# Bytes pro KV-Cache-Element je cache-type (inkl. Block-Overhead der k-Quants).
_KV_BPE = {
    "f32": 4.0, "f16": 2.0, "bf16": 2.0,
    "q8_0": 1.0625, "q5_1": 0.75, "q5_0": 0.6875,
    "q4_1": 0.625, "q4_0": 0.5625, "iq4_nl": 0.5625,
}
_GIB = 1024 ** 3


def _bpe(cache_type: str | None) -> float:
    return _KV_BPE.get((cache_type or "f16").lower(), 2.0)


def kv_cache_gb(meta: dict, ctx: int, ck: str | None = None, cv: str | None = None) -> float:
    """Echte KV-Cache-Größe (GiB) für ctx Tokens, K/V ggf. quantisiert. Formel wie llama.cpp:
    je Layer & Token hält der Cache n_head_kv × head_dim Elemente für K und für V."""
    per_tok = meta["n_layers"] * meta["n_head_kv"] * ctx
    k = per_tok * meta["head_dim_k"] * _bpe(ck)
    v = per_tok * meta["head_dim_v"] * _bpe(cv)
    return (k + v) / _GIB


def kv_gb_per_token(meta: dict, ck: str | None = None, cv: str | None = None) -> float:
    """KV-GiB pro Kontext-Token — für den analytischen ctx-Solver (linear in ctx)."""
    return kv_cache_gb(meta, 1, ck, cv)