"""
Speicher-Budget & SETUP-BEWUSSTE ctx-Vergabe — EINE Quelle der Wahrheit.

Modelliert die auf der Box VERIFIZIERTE Residenz-Realität (llama-swap, GTT ~124 GB):
  • Die `persistent`-Gruppe (brains = Hirn+embed+vision) bleibt IMMER resident —
    seit dem Key-Fix 03.07.2026 greift der Schutz wirklich (vorher stand `persist`
    in der Config, das llama-swap stillschweigend ignorierte; on-demand-Last
    verdrängte damals die ganze Gruppe).
  • Ein on-demand-Modell (heavy/coder/…) lädt NEBEN die brains-Gruppe und muss
    deren Footprint mit einplanen.
Daraus folgt, wie viel Speicher NEBEN einem Zielmodell reserviert bleiben muss —
und damit der größte Kontext, der wirklich passt (nicht nur für das Modell allein).

Vorher rechnete nur der Hirn-Wechsel (agent.py) setup-bewusst; die allgemeine
ctx-Vergabe nahm den Gesamt-RAM in Isolation. Dieses Modul vereint beides.
"""

import re

import psutil

from services.fit import (
    QUANT_BYTES_PER_PARAM,
    estimate_memory_gb,
    extract_params_b,
    max_ctx_in_budget,
)

HEADROOM_GB = 4.0   # OS/Treiber/Fragmentierung


def gtt_budget_gb() -> float:
    """GPU-adressierbarer Speicher (GTT) in GB — die harte Obergrenze. Liest
    amdgpu.gttsize aus /proc/cmdline, sonst RAM minus OS-Reserve."""
    try:
        with open("/proc/cmdline") as f:
            m = re.search(r"amdgpu\.gttsize=(\d+)", f.read())
            if m:
                return round(int(m.group(1)) / 1024.0, 1)
    except Exception:
        pass
    return round(psutil.virtual_memory().total / (1024 ** 3) - 6.0, 1)


def params_of_model(model: dict) -> float:
    """Robuste Params (Mrd.) eines INSTALLIERTEN Modells: MAXIMUM aus Caps-Schätzung und
    Dateigröße. Deckt 'Coder-Next' ohne Größe im Namen (→ aus Datei) und Split-GGUFs
    (size_bytes = nur erster Teil → ignoriert) ab."""
    caps = model.get("capabilities") or {}
    quant = model.get("quant") or "Q4_K_M"
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.55)
    size_gb = (model.get("size_bytes") or 0) / (1024 ** 3)
    pb_size = (size_gb / bpp) if size_gb > 1.0 else 0.0
    return max(float(caps.get("params_b") or 0), pb_size, 7.0)


_CTK_RE = re.compile(r"(?:--cache-type-k|(?<![\w-])-ctk)\s+(\S+)")
_CTV_RE = re.compile(r"(?:--cache-type-v|(?<![\w-])-ctv)\s+(\S+)")


def _cache_types(cmd: str) -> tuple[str | None, str | None]:
    """K/V-Cache-Quantisierung aus dem llama-server-Cmd (Default f16 → None)."""
    ck = m.group(1) if (m := _CTK_RE.search(cmd or "")) else None
    cv = m.group(1) if (m := _CTV_RE.search(cmd or "")) else None
    return ck, cv


def _real_kv_gb(model: dict, ctx: int) -> float | None:
    """ECHTE KV-Cache-Größe (GiB) aus den GGUF-Architektur-Metadaten (Layer × KV-Heads ×
    Head-Dim) + der cache-type-Quantisierung des Cmds. None, wenn das GGUF nicht lesbar ist
    → Aufrufer fällt auf die params-basierte Heuristik zurück."""
    from services import gguf_meta
    path = model.get("gguf_path")
    if not path:
        return None
    meta = gguf_meta.arch_meta(path)
    if not meta:
        return None
    ck, cv = _cache_types(model.get("cmd") or "")
    return gguf_meta.kv_cache_gb(meta, ctx, ck, cv)


def footprint_gb(model: dict) -> float:
    """Loaded-Footprint eines Modells = Gewichte + KV-Cache (bei seinem aktuellen ctx).
    KV kommt aus den ECHTEN Architektur-Metadaten des GGUF (nicht mehr params-geschätzt) —
    entscheidend bei MoE (A3B): die alte Schätzung hing an den Gesamt-Params und überschätzte
    grob (z.B. „68 GB reserviert" statt real ~25 GB). Heuristik bleibt Fallback."""
    quant = model.get("quant") or "Q4_K_M"
    ctx = int(model.get("ctx") or 32768)
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.55)
    size_gb = (model.get("size_bytes") or 0) / (1024 ** 3)
    pb = params_of_model(model)
    weights = max(pb * bpp, size_gb)
    kv = _real_kv_gb(model, ctx)
    if kv is None:
        kv = estimate_memory_gb(pb, quant, ctx) - pb * bpp
    return weights + max(kv, 0.0)


def params_b_for(name: str) -> float:
    """Parameter (Mrd.) für einen Modell-/Repo-Namen: KATALOG (echte Metadaten) zuerst,
    sonst Namens-Schätzung. Gemeinsam für Fit-Vorschau und ctx-Vergabe."""
    from services import catalog
    meta = catalog.meta_for_name(name) if name else None
    if meta and meta.get("total_params_b"):
        return float(meta["total_params_b"])
    return extract_params_b(name)


def _coresident_members(groups: dict) -> set:
    """Modelle, die GLEICHZEITIG warm sind: Mitglieder aller `swap:false`-Gruppen
    (Ko-Residenz, z.B. brains = Hirn+embed+vision). Seit dem `persistent`-Fix (03.07.2026)
    überlebt die Gruppe auch on-demand-Last: Coder/heavy laden DANEBEN, nicht an ihre
    Stelle (live verifiziert: Coder + Qwen3.6 gleichzeitig `ready`). Die frühere
    Beobachtung „heavy verdrängt die Gruppe" war der ignorierte `persist`-Key."""
    out: set = set()
    for g in (groups or {}).values():
        if isinstance(g, dict) and g.get("swap") is False:
            out.update(g.get("members") or [])
    return out


def reserved_gb(role: str | None) -> dict:
    """Speicher, der NEBEN einem Zielmodell der gegebenen Rolle resident bleibt — gemäß der
    seit dem `persistent`-Fix (03.07.2026) geltenden Semantik: die ko-residente
    `swap:false`-Gruppe (brains) bleibt IMMER geladen, on-demand-Modelle laden daneben.

    - Modell IN der Ko-Residenz-Gruppe (Hirn/embed/vision): koexistiert mit den ÜBRIGEN
      Gruppen-Mitgliedern → reserviert deren Summe.
    - Modell AUSSERHALB (heavy/coder/coder-lite/scout): lädt NEBEN die Gruppe →
      reserviert deren GESAMTE Summe (früher 0.0, weil der kaputte `persist`-Key die
      Gruppe verdrängen ließ — diese Rechnung erlaubte zu große Kontexte).
    """
    from services import llamaswap
    models = llamaswap.list_models()
    groups = llamaswap.list_groups()
    cores = _coresident_members(groups)
    brain = next((m for m in models if (m.get("role") == "hermes")), None)
    brain_gb = footprint_gb(brain) if brain else 0.0
    role = (role or "").strip().lower()

    holder = next((m for m in models if (m.get("role") == role)), None) if role else None
    holder_name = holder["name"] if holder else None
    # Hirn (hermes) ist per Definition Teil der Ko-Residenz-Gruppe; sonst Gruppen-Mitgliedschaft prüfen.
    in_group = role == "hermes" or bool(holder_name and holder_name in cores)

    if in_group:
        others = sum(footprint_gb(m) for m in models
                     if m["name"] in cores and m["name"] != holder_name)
        return {"reserved_gb": others, "mode": "co-resident", "brain_gb": brain_gb}
    # on-demand: lädt neben die (persistente) Ko-Residenz-Gruppe → deren Summe reservieren.
    warm = sum(footprint_gb(m) for m in models if m["name"] in cores)
    return {"reserved_gb": warm, "mode": "ondemand-beside-warmset", "brain_gb": brain_gb}


def setup_aware_ctx(params_b: float, quant: str, role: str | None = None) -> dict:
    """Größter Kontext, der für ein Modell (params_b/quant) der gegebenen Rolle NEBEN dem
    bestehenden Setup passt. Gibt ctx + die Budget-Herleitung zurück (für UI/Transparenz)."""
    gtt = gtt_budget_gb()
    r = reserved_gb(role)
    budget = max(gtt - r["reserved_gb"] - HEADROOM_GB, 0.0)
    ctx = max_ctx_in_budget(params_b, quant, budget)
    return {
        "ctx": ctx,
        "gtt_gb": gtt,
        "reserved_gb": round(r["reserved_gb"], 1),
        "budget_gb": round(budget, 1),
        "mode": r["mode"],
    }


def _snap_ctx(raw_ctx: float, cap: int | None = None) -> int:
    """Größter 'schöner' Kontext ≤ raw_ctx (und ≤ Trainings-Kontext des Modells, falls bekannt)."""
    from services.fit import _NICE_CTX
    if cap:
        raw_ctx = min(raw_ctx, cap)
    best = _NICE_CTX[0]
    for c in _NICE_CTX:
        if c <= raw_ctx:
            best = c
    return best


def setup_aware_ctx_for_model(model: dict) -> dict:
    """Setup-bewusster Optimal-ctx für ein INSTALLIERTES Modell. Für den 'Auto'-Button an der
    Modellkarte. Nutzt die ECHTE KV-Größe des GGUF (gleiche Zahlensprache wie footprint_gb) —
    Fallback auf die params-Heuristik nur, wenn das GGUF nicht lesbar ist."""
    from services import gguf_meta
    quant = model.get("quant") or "Q4_K_M"
    path = model.get("gguf_path")
    meta = gguf_meta.arch_meta(path) if path else None
    if not meta:
        return setup_aware_ctx(params_of_model(model), quant, role=model.get("role"))

    gtt = gtt_budget_gb()
    r = reserved_gb(model.get("role"))
    budget = max(gtt - r["reserved_gb"] - HEADROOM_GB, 0.0)
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.55)
    size_gb = (model.get("size_bytes") or 0) / (1024 ** 3)
    weights = max(params_of_model(model) * bpp, size_gb)
    ck, cv = _cache_types(model.get("cmd") or "")
    per_tok = gguf_meta.kv_gb_per_token(meta, ck, cv)
    ctx = _snap_ctx((budget - weights) / per_tok, cap=meta.get("n_ctx_train")) if per_tok > 0 else 2048
    return {"ctx": ctx, "gtt_gb": gtt, "reserved_gb": round(r["reserved_gb"], 1),
            "budget_gb": round(budget, 1), "mode": r["mode"]}