#!/usr/bin/env python3
"""model_registry — reads model-registry.yaml (kb#133, A2A-1).

Per DESIGN-a2a-agents.md v2.1 §2-3b: the model registry is data, not logic.
a(t) probe *mechanics* (QuotaProbe, GPUResidencyProbe, ...) already live in
kanboard/bin/kb_worker.py — this module does not reimplement them. It gives
callers two things instead:

  * resolve(tier)        — "an available model for tier X" without the
                            caller naming a model. Structural availability
                            (lifecycle, metered opt-in) is decided here from
                            registry data; live a(t) truth (is the quota
                            window open right now, is the GPU actually free)
                            is decided by an optional `probe_check` callback
                            the caller supplies (e.g. wired to kb_worker's
                            Probe classes via to_probe_config()).
  * preload_check(...)   — the §3b GPU pre-load check, expressed purely from
                            registry numbers (never-evict reservations +
                            candidate footprint) plus a headroom figure the
                            caller supplies. It does not shell nvidia-smi
                            itself — kb_worker.GPUResidencyProbe (or
                            `nvidia-smi` directly) is the live-read path;
                            this stays pure/testable.

Usage (library):
    from model_registry import load_registry, resolve, to_probe_config, preload_check
    reg = load_registry()
    model = resolve(reg, tier="large")               # -> the "kimi" entry
    cfg = to_probe_config(reg, "kimi")                # -> kb_worker probe config dict
    ok, reason = preload_check(reg, "local-small", headroom_mb=1900)

Usage (CLI, for manual verification):
    ./model_registry.py resolve --tier large
    ./model_registry.py resolve --tier large --allow-metered --opted-in agent:torgash
    ./model_registry.py probe-config --id kimi
    ./model_registry.py preload-check --id local-small --headroom-mb 1900
    ./model_registry.py preload-check --id local-small --live   # shells nvidia-smi
    ./model_registry.py list
"""

import argparse
import json
import os
import subprocess
import sys

import yaml

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_REGISTRY_PATH = os.path.join(HERE, "model-registry.yaml")


class RegistryError(Exception):
    pass


def load_registry(path=None):
    """Load and lightly validate model-registry.yaml."""
    path = path or DEFAULT_REGISTRY_PATH
    with open(path) as f:
        reg = yaml.safe_load(f)
    if not reg or "models" not in reg:
        raise RegistryError(f"{path}: missing top-level 'models' list")
    ids = [m["id"] for m in reg["models"]]
    if len(ids) != len(set(ids)):
        raise RegistryError(f"{path}: duplicate model ids in {ids}")
    return reg


def get_model(registry, model_id):
    for m in registry["models"]:
        if m["id"] == model_id:
            return m
    raise RegistryError(f"unknown model id: {model_id!r}")


# ---------------------------------------------------------------------------
# resolve — "an available model for tier X" without the caller naming one.
# ---------------------------------------------------------------------------

def resolve(registry, tier, allow_metered=False, opted_in_key=None, probe_check=None):
    """Return the first model in `tier`'s pool that is structurally usable,
    and (if probe_check is given) currently available.

    Structural filter (from registry data alone):
      - candidate must be listed under routing.tiers[tier]
      - a metered model is only a candidate at all when the CALLER passes
        allow_metered=True AND opted_in_key appears in routing.metered_opt_in
        (§3a: "no metered API by default" — an empty metered_opt_in list, the
        shipped default, makes every metered model structurally unreachable
        regardless of allow_metered).

    Live filter (optional): probe_check(model_dict) -> bool. Wire this to
    kb_worker's Probe.available() (via to_probe_config below) when the
    caller wants real a(t) truth instead of just structural eligibility.
    """
    pools = registry.get("routing", {}).get("tiers", {})
    if tier not in pools:
        raise RegistryError(f"unknown tier: {tier!r} (have: {sorted(pools)})")
    opted_in = set(registry.get("routing", {}).get("metered_opt_in", []) or [])

    candidates = []
    for model_id in pools[tier]:
        m = get_model(registry, model_id)
        if m.get("metered"):
            if not allow_metered:
                continue
            if not m.get("opt_in_required", True):
                # Registry says this metered model doesn't need opt-in — treat
                # as a data error rather than silently routing to it.
                raise RegistryError(
                    f"model {model_id!r} is metered but opt_in_required=false; "
                    "fix the registry entry, this helper will not assume implicit access"
                )
            if opted_in_key is None or opted_in_key not in opted_in:
                continue
        candidates.append(m)

    for m in candidates:
        if probe_check is None or probe_check(m):
            return m

    raise RegistryError(
        f"no available model for tier={tier!r} "
        f"(allow_metered={allow_metered}, opted_in_key={opted_in_key!r}); "
        f"checked candidates: {[m['id'] for m in candidates] or pools[tier]}"
    )


# ---------------------------------------------------------------------------
# to_probe_config — bridges registry quota data into kb_worker's probe cfg
# shape (kanboard/bin/kb_worker.py PROBE_BUILDERS), so probes read registry
# numbers instead of the registry re-implementing probe logic.
# ---------------------------------------------------------------------------

def to_probe_config(registry, model_id):
    """Return a dict matching kb_worker.py's `build_probe(cfg)` input shape
    for `model_id`'s lifecycle. Raises if the model has no probe-relevant
    lifecycle (e.g. cost-gated with no probe_command wired yet)."""
    m = get_model(registry, model_id)
    lifecycle = m["lifecycle"]

    if lifecycle == "always-on":
        return {"type": "always_on"}

    if lifecycle == "quota-gated":
        q = m.get("quota") or {}
        windows = q.get("windows") or []
        if not windows:
            raise RegistryError(f"{model_id}: quota-gated but no quota.windows configured")
        # kb_worker's QuotaProbe checks one field; the tightest (first-to-hit)
        # window in practice is the short one — default to the first entry,
        # callers needing multi-window gating build one probe per window.
        window = windows[0]
        return {
            "type": "quota",
            "command": q["probe_command"],
            "field": window["field"],
            "threshold_pct": q.get("threshold_pct", 95),
        }

    if lifecycle == "cost-gated":
        q = m.get("quota") or {}
        if not q.get("probe_command"):
            raise RegistryError(
                f"{model_id}: cost-gated but no budget probe wired yet "
                "(opt-in path incomplete — see registry comment)"
            )
        return {
            "type": "budget",
            "command": q["probe_command"],
            "field": q["field"],
            "limit": q["limit"],
        }

    if lifecycle == "on-demand":
        ep = (m.get("endpoints") or [{}])[0]
        if not ep.get("health_url"):
            raise RegistryError(f"{model_id}: on-demand but no endpoint.health_url configured")
        return {"type": "on_demand", "url": ep["health_url"]}

    raise RegistryError(f"{model_id}: unknown lifecycle {lifecycle!r}")


# ---------------------------------------------------------------------------
# preload_check — §3b GPU pre-load check, pure registry-data math. The live
# VRAM headroom READ is the caller's job (kb_worker.GPUResidencyProbe or
# nvidia-smi directly) — see live_headroom_mb() below for a thin convenience
# wrapper used only by this module's own CLI, not by the check itself.
# ---------------------------------------------------------------------------

def preload_check(registry, candidate_id, headroom_mb, resident_ids=None):
    """Would loading `candidate_id` fit, given `headroom_mb` free VRAM right
    now (as reported by a live probe) and `resident_ids` already loaded?

    Never-evict models are never subtracted from headroom by a caller in the
    first place (their footprint is a standing reservation baked into any
    correct live headroom read) — this function just re-asserts that policy
    from registry data: if `candidate_id` itself is never-evict, it always
    passes (it's not something a worker "loads speculatively" and might be
    told to skip); anything else must fit in the reported headroom.
    """
    policy = registry.get("gpu_residency_policy") or {}
    m = get_model(registry, candidate_id)
    gr = m.get("gpu_residency")
    if not gr:
        return True, f"{candidate_id} has no gpu_residency entry (not a GPU-resident model)"

    if gr.get("never_evict"):
        return True, f"{candidate_id} is in the never-evict set — always resident by policy"

    required_mb = gr["vram_mb"]
    never_evict_ids = set(policy.get("never_evict_ids", []))
    resident_ids = set(resident_ids or [])
    # Sanity: if the live headroom read already accounts for never-evict
    # reservations (the expected contract — see kb_worker.GPUResidencyProbe's
    # never_evict_reserved_mb param), this is just a straight comparison.
    # If a caller passes raw total-minus-used instead, warn via the reason
    # string rather than silently under/over-reserving.
    reserved_hint = sum(
        get_model(registry, mid)["gpu_residency"]["vram_mb"]
        for mid in never_evict_ids
        if mid not in resident_ids  # already counted as "used" if resident_ids says so
    )
    ok = headroom_mb >= required_mb
    reason = (
        f"headroom={headroom_mb}MB required={required_mb}MB "
        f"(never-evict reserve expected already netted out by the caller's probe: "
        f"~{reserved_hint}MB across {sorted(never_evict_ids)})"
    )
    return ok, reason


def live_headroom_mb(never_evict_reserved_mb=0):
    """Convenience for manual CLI checks only — NOT used by preload_check()
    itself. Shells nvidia-smi the same way kb_worker.GPUResidencyProbe does."""
    out = subprocess.run(
        ["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits"],
        capture_output=True, text=True, timeout=5, check=True,
    ).stdout.strip().splitlines()[0]
    used, total = (int(x) for x in out.split(","))
    return total - used - never_evict_reserved_mb


# ---------------------------------------------------------------------------
# CLI — manual verification only, not part of the library contract.
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--registry", default=None, help="path to model-registry.yaml (default: sibling file)")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("resolve")
    p.add_argument("--tier", required=True)
    p.add_argument("--allow-metered", action="store_true")
    p.add_argument("--opted-in", default=None, help="virtual key claimed to be opted in")

    p = sub.add_parser("probe-config")
    p.add_argument("--id", required=True)

    p = sub.add_parser("preload-check")
    p.add_argument("--id", required=True)
    p.add_argument("--headroom-mb", type=float, default=None)
    p.add_argument("--live", action="store_true", help="read live headroom via nvidia-smi instead of --headroom-mb")
    p.add_argument("--resident", action="append", default=[], help="repeatable: id already resident")

    sub.add_parser("list")

    args = ap.parse_args()
    reg = load_registry(args.registry)

    try:
        if args.cmd == "resolve":
            m = resolve(reg, args.tier, allow_metered=args.allow_metered, opted_in_key=args.opted_in)
            print(json.dumps(m, indent=2))
        elif args.cmd == "probe-config":
            print(json.dumps(to_probe_config(reg, args.id), indent=2))
        elif args.cmd == "preload-check":
            headroom = args.headroom_mb
            if args.live:
                policy = reg.get("gpu_residency_policy") or {}
                reserved = sum(get_model(reg, mid)["gpu_residency"]["vram_mb"]
                               for mid in policy.get("never_evict_ids", []))
                headroom = live_headroom_mb(never_evict_reserved_mb=reserved)
            if headroom is None:
                raise RegistryError("preload-check needs --headroom-mb or --live")
            ok, reason = preload_check(reg, args.id, headroom, resident_ids=args.resident)
            print(json.dumps({"ok": ok, "reason": reason}))
            sys.exit(0 if ok else 1)
        elif args.cmd == "list":
            for m in reg["models"]:
                print(f"{m['id']:16} tier={m['tier']:6} lifecycle={m['lifecycle']:13} "
                      f"metered={m['metered']} role={m['role']}")
    except RegistryError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(2)


if __name__ == "__main__":
    main()