"""
Modell-Capabilities — EINE Quelle der Wahrheit für Modell-Eigenschaften
(MoE / Tools / Vision / Coder / Reasoning / Embedding / Kontext).
Portiert aus Mission Control v1 (model_caps.py).

Quellen, geschichtet: GGUF-Header (offline, authoritativ) → cmd-Flags
(--jinja/--mmproj) → HF-Block (tags + chat_template) → Familien-Fallback.
Tool-Fähigkeit dreistufig: yes (bestätigt) | likely (Familie) | no.
"""

import re
import struct

from services.fit import extract_active_params_b, extract_params_b

_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 Tokenizer-Array ab → schnell, lädt NICHT das Modell."""
    out: dict = {}
    try:
        with open(path, "rb") as f:
            if f.read(4) != b"GGUF":
                return {}
            struct.unpack("<I", f.read(4))[0]
            f.read(8)
            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:
                    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":
                    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


_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", "glm-5", "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 für 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 "")

    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))
        or bool(re.search(r"a\d+\.?\d*b", low))
    )
    active_b = extract_active_params_b(name)

    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

    # Cap context length at 131072 for Qwen / Hermes models to prevent reporting scaled RoPE context of 256k+ which might OOM or be unstable.
    if ctx and ctx > 131072 and ("qwen" in low or "hermes" in low):
        ctx = 131072

    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, "vision": vision,
        "coder": coder, "reasoning": reasoning, "embedding": embedding,
        "ctx": ctx, "params_b": params_b or None, "arch": arch or None,
    }