"""
Modell-Capabilities (Phase 10) — EINE Quelle der Wahrheit fuer Modell-Eigenschaften.

Statt Einzel-Label ("Code"/"Text") ein Satz unabhaengiger Tags, die ein Modell
gleichzeitig tragen kann (MoE + Tools + Reasoning + Coder + Long-Context …).

Quellen, geschichtet (sicher → Fallback):
  1) GGUF-Metadaten der lokalen Datei (architecture, expert_count, context_length,
     parameter_count) — authoritativ & **offline** (kein Netz, passt zu 100%-lokal).
  2) cmd-Flags der llama-swap-Config: `--jinja` = Tool-Template aktiv, `--mmproj` = Vision.
  3) duenner Namens-/Familien-Fallback, wenn Metadaten fehlen.
  4) optional `hf`-Block (HF-API: tags + gguf.chat_template) fuer die Profi-Suche (Schritt 2).

Tool-Faehigkeit kommt NICHT aus dem Dateinamen, sondern aus `--jinja` (bestaetigt) /
dem Chat-Template (HF) / der Familie (vermutet) → drei Stufen: yes | likely | no.
"""

import re
import struct

from hw_math import extract_params_b, extract_active_params_b

# GGUF-Skalar-Typen → Bytebreite (fuer Skip)
_GGUF_FIXED = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}


def _read_gguf_meta(path: str) -> dict:
    """Liest nur den GGUF-Metadaten-Header (architecture/context_length/expert_count/
    parameter_count). Bricht vor dem riesigen Tokenizer-Array ab → schnell, laedt NICHT
    das Modell. Robust: gibt {} bei jedem Fehler."""
    out: dict = {}
    try:
        with open(path, "rb") as f:
            if f.read(4) != b"GGUF":
                return {}
            struct.unpack("<I", f.read(4))[0]      # version
            f.read(8)                              # tensor_count (u64)
            kv = struct.unpack("<Q", f.read(8))[0]

            def ru32() -> int: return struct.unpack("<I", f.read(4))[0]
            def ru64() -> int: return struct.unpack("<Q", f.read(8))[0]
            def rstr() -> str: return f.read(ru64()).decode("utf-8", "replace")

            def rval(t: int):
                if t == 8:  return rstr()
                if t == 0:  return struct.unpack("<B", f.read(1))[0]
                if t == 1:  return struct.unpack("<b", f.read(1))[0]
                if t == 2:  return struct.unpack("<H", f.read(2))[0]
                if t == 3:  return struct.unpack("<h", f.read(2))[0]
                if t == 4:  return struct.unpack("<I", f.read(4))[0]
                if t == 5:  return struct.unpack("<i", f.read(4))[0]
                if t == 6:  return struct.unpack("<f", f.read(4))[0]
                if t == 7:  return f.read(1) != b"\x00"
                if t == 10: return struct.unpack("<Q", f.read(8))[0]
                if t == 11: return struct.unpack("<q", f.read(8))[0]
                if t == 12: return struct.unpack("<d", f.read(8))[0]
                if t == 9:  # Array: nur ueberlesen (brauchen wir nicht)
                    et = ru32(); cnt = ru64()
                    if et == 8:
                        for _ in range(cnt):
                            f.seek(ru64(), 1)
                    elif et == 9:
                        for _ in range(cnt):
                            rval(9)
                    else:
                        f.seek(cnt * _GGUF_FIXED.get(et, 0), 1)
                    return None
                raise ValueError(f"unbekannter GGUF-Typ {t}")

            want = {"architecture", "context_length", "expert_count", "parameter_count"}
            for _ in range(kv):
                key = rstr()
                t = ru32()
                if key == "tokenizer.ggml.tokens":   # ab hier nur noch riesige Arrays
                    break
                v = rval(t)
                short = key.split(".")[-1]
                if short in want and short not in out:
                    out[short] = v
    except Exception:
        return out
    return out


# Familien-Fallback (nur wenn Metadaten schweigen). Bewusst kurz gehalten.
_TOOL_FAMILIES = (
    "qwen2.5", "qwen3", "qwen2", "hermes", "mistral", "mixtral", "devstral",
    "command-r", "command_r", "llama-3.1", "llama3.1", "llama-3.3", "llama-4", "llama4",
    "functionary", "watt", "firefunction", "granite", "glm-4", "ministral",
)
_REASON_KW = (
    "-r1", "deepseek-r1", "qwq", "magistral", "-think", "thinking", "-o1",
    "gpt-oss", "reasoning", "exaone-deep", "phi-4-reasoning", "phi-4-mini-reasoning",
)
_CODE_KW = ("coder", "-code", "code-", "codestral", "starcoder", "deepseek-coder")
_VISION_KW = ("-vl", "vision", "llava", "pixtral", "multimodal", "-mm-", "qwen3vl", "qwen2-vl")
_EMBED_KW = ("bge", "e5-", "gte-", "nomic-embed", "embed")
_MOE_ARCH = ("moe", "mixtral", "deepseek2", "deepseek3", "llama4", "qwen3moe", "grok")


def capabilities(name: str = "", cmd: str = "", gguf_path: str = "", hf: dict | None = None) -> dict:
    """Capability-Tag-Set fuer ein Modell. Alle Quellen optional — nutzt, was da ist."""
    low = (name or "").lower()
    cmdl = (cmd or "").lower()
    hf = hf or {}

    meta = _read_gguf_meta(gguf_path) if gguf_path else {}
    arch = str(meta.get("architecture") or hf.get("architecture") or "").lower()
    tags = [str(t).lower() for t in (hf.get("tags") or [])]
    chat_tpl = str(hf.get("chat_template") or "")

    # --- MoE ---
    expert_count = int(meta.get("expert_count") or 0)
    moe = (
        expert_count > 1
        or any(a in arch for a in _MOE_ARCH)
        or bool(re.search(r"\d+x\d+\.?\d*b", low))     # 8x7B
        or bool(re.search(r"a\d+\.?\d*b", low))        # 30B-A3B
    )
    active_b = extract_active_params_b(name)

    # --- Params / Kontext ---
    pcount = int(meta.get("parameter_count") or 0)
    params_b = round(pcount / 1e9, 1) if pcount else extract_params_b(name)
    ctx = meta.get("context_length")
    if not ctx:
        m = re.search(r"-(?:c|-ctx-size)\s+(\d+)", cmdl)
        ctx = int(m.group(1)) if m else None

    # --- Tools: yes (bestaetigt) | likely (Familie) | no ---
    tool_confirmed = "--jinja" in cmdl or "tool_call" in chat_tpl or "<tools>" in chat_tpl
    tool_family = any(fam in low for fam in _TOOL_FAMILIES) or "function-calling" in tags
    tools = "yes" if tool_confirmed else ("likely" if tool_family else "no")

    vision = "--mmproj" in cmdl or "vl" in arch or "clip" in arch or any(k in low for k in _VISION_KW)
    coder = any(k in low for k in _CODE_KW)
    reasoning = any(k in low for k in _REASON_KW) or "reasoning" in tags
    embedding = "bert" in arch or any(k in low for k in _EMBED_KW)

    return {
        "moe": moe,
        "active_b": active_b,
        "tools": tools,                 # yes | likely | no
        "vision": vision,
        "coder": coder,
        "reasoning": reasoning,
        "embedding": embedding,
        "ctx": ctx,
        "params_b": params_b or None,
        "arch": arch or None,
    }