"""
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 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 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}
            for _ in range(kv_count):
                key = r.gstr()
                vtype = r.u32()
                if key == "tokenizer.ggml.tokens" and vtype == _T_ARRAY:
                    etype = r.u32()
                    fp["n_vocab"] = r.u64()
                    # Wir haben alles (model/pre kommen vor tokens) → abbrechen.
                    if etype != _T_STRING:
                        return None
                    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"))


def fingerprint(path: str) -> dict | None:
    """Tokenizer-Fingerprint eines GGUF (gecacht nach Pfad+mtime+size).
    Returns dict(model, pre, n_vocab, arch) 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]}


def vocab_key(path: str) -> tuple | None:
    """Vergleichsschlüssel für Vocab-Kompatibilität: (model, pre, n_vocab).
    Genau diese Identität verlangt llama.cpp für Speculative Decoding."""
    fp = fingerprint(path)
    if not fp or fp["n_vocab"] is None:
        return None
    return (fp["model"], fp["pre"], fp["n_vocab"])


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