#!/usr/bin/env python3
"""
vllmctl — run vLLM (docker, ROCm) on fl WITHOUT freezing the box.

The RX 7600 XT drives the display; VRAM overcommit spills to GTT and freezes the PC. vLLM grabs
`gpu-memory-utilization` × total VRAM at startup and keeps it, so this script gates every start:

  1. llama.service must be STOPPED (or --stop-router): a LAN client hitting :11343 would make the
     router load a model on top of vLLM = the freeze. --unsafe-keep-router skips this gate.
  2. the card must be clean (VRAM used ≤ drain-gb; resident router models are unloaded via its API).
  3. projected free VRAM = total − used − budget must be ≥ reserve-gb (else abort, --force overrides).

Presets live in vllm.ini next to this file (same key = --flag convention as llamacpp/config.ini).

  ./vllmctl.py up <preset> [--stop-router] [--force] [--dry-run]   start, wait for /health
  ./vllmctl.py down                                                stop, wait for VRAM to drain
  ./vllmctl.py status | logs | pull | gpu-check | exec <cmd...>
Stdlib only. Run ON fl.
"""
import argparse, grp, json, os, re, shlex, subprocess, sys, time, urllib.request, urllib.error

HERE = os.path.dirname(os.path.abspath(__file__))
INI = os.path.join(HERE, "vllm.ini")
DRM = "/sys/class/drm/renderD128/device"


# ─── ini ────────────────────────────────────────────────────────────────────────────────────
def parse_ini(path):
    """{section: {key: value}} — inline `# comments` stripped, order kept."""
    sections, sec = {}, None
    for ln in open(path, encoding="utf-8").read().splitlines():
        s = ln.strip()
        if not s or s.startswith("#"):
            continue
        if s.startswith("[") and s.endswith("]"):
            sec = s[1:-1]; sections[sec] = {}
        elif sec and "=" in s:
            k, v = s.split("=", 1)
            sections[sec][k.strip()] = v.split(" #", 1)[0].split("\t#", 1)[0].strip()
    return sections


def build_args(preset):
    """vllm serve args from a preset section (x-* and env.* are not flags)."""
    args = []
    for k, v in preset.items():
        if k.startswith("x-") or k.startswith("env."):
            continue
        if v.lower() == "false":
            continue
        if v.lower() == "true":
            args.append(f"--{k}")
        else:
            args += [f"--{k}", v]
    return args


# ─── gpu / router ───────────────────────────────────────────────────────────────────────────
def sysfs(name):
    try:
        return int(open(os.path.join(DRM, name)).read().strip())
    except (OSError, ValueError):
        return None


def vram():
    """(used, total, gtt_used) bytes from amdgpu sysfs; None where unreadable."""
    return sysfs("mem_info_vram_used"), sysfs("mem_info_vram_total"), sysfs("mem_info_gtt_used")


def gb(b):
    return f"{b / 1e9:.1f} GB" if b is not None else "-"


def http(url, payload=None, timeout=10):
    data = json.dumps(payload).encode() if payload is not None else None
    hdrs = {"Content-Type": "application/json"} if data else {}
    return urllib.request.urlopen(urllib.request.Request(url, data=data, headers=hdrs), timeout=timeout)


def _status(m):
    """Router /models status: newer builds send {'value': 'loaded'|'sleeping'|'unloaded'|…}, older a string."""
    st = m.get("status", "loaded")
    return st.get("value", "loaded") if isinstance(st, dict) else st


def router_loaded(root):
    """ids the llama.cpp router reports as loaded ([] if unreachable)."""
    try:
        d = json.load(http(root.rstrip("/") + "/models", timeout=5))
    except Exception:
        return []
    return [m["id"] for m in d.get("data", []) if _status(m) in ("loaded", "loading")]


def router_unload(root):
    for m in router_loaded(root):
        try:
            http(root.rstrip("/") + "/models/unload", {"model": m}, timeout=30).read()
        except Exception:
            pass


def router_active():
    try:
        r = subprocess.run(["systemctl", "is-active", "llama.service"], capture_output=True, text=True, timeout=10)
    except (OSError, subprocess.TimeoutExpired):
        return False
    return r.stdout.strip() == "active"


def router_hint():
    return "" if router_active() else "\n  (llama.service is STOPPED → sudo systemctl start llama.service to bring the router back)"


def wait_drain(drain_gb, window_s, label):
    used, _, _ = vram()
    end = time.time() + window_s
    said = False
    while used is not None and used > drain_gb * 1e9 and time.time() < end:
        if not said:
            print(f"  … waiting for VRAM to drain {label} ({gb(used)} used)"); said = True
        time.sleep(5)
        used, _, _ = vram()
    return used


# ─── docker ─────────────────────────────────────────────────────────────────────────────────
def docker(*args, check=True, capture=True):
    return subprocess.run(["docker", *args], check=check, capture_output=capture, text=True)


def container_state(name):
    r = docker("inspect", "-f", "{{.State.Status}}", name, check=False)
    return r.stdout.strip() if r.returncode == 0 else None


def gpu_flags():
    flags = ["--device", "/dev/kfd", "--device", "/dev/dri", "--group-add", "video",
             "--cap-add", "SYS_PTRACE", "--security-opt", "seccomp=unconfined", "--ipc=host"]
    try:
        flags += ["--group-add", str(grp.getgrnam("render").gr_gid)]   # Arch: /dev/dri/renderD128 is group render
    except KeyError:
        pass
    return flags


def run_cmd(g, name, preset, envs):
    cmd = ["docker", "run", "-d", "--name", g["container"], *gpu_flags(),
           "-p", f'{g["port"]}:8000',
           "-v", f'{g["models-dir"]}:/models:ro',
           "-v", f'{g["hf-cache"]}:/root/.cache/huggingface']
    for k, v in envs.items():
        cmd += ["-e", f"{k}={v}"]
    cmd.append(g["image"])
    cmd += ["--host", "0.0.0.0", "--port", "8000", "--served-model-name", preset.get("served-model-name", name)]
    cmd += build_args({k: v for k, v in preset.items() if k != "served-model-name"})
    return cmd


ROOT_CAUSE = re.compile(r"hipError|HSA_STATUS|code object|gfx\d+|no kernel image|EngineCore.*(ERROR|Error|error)"
                        r"|^\S*(ERROR|CRITICAL)|RuntimeError|ValueError|NotImplementedError|OutOfMemory|Killed|Segmentation"
                        r"|free memory|not supported|Unsupported|Traceback|memory profiling|kv cache memory|weights take|peak", re.I)


def root_cause(lines):
    """The lines that explain a dead engine core: known ROCm/vLLM failure signatures, the FIRST
    traceback's error line, and the last 8 lines — not the parent's 40-line stack trace."""
    out, seen = [], set()
    for l in lines:
        if ROOT_CAUSE.search(l) and "api_server" not in l and "core_client" not in l and "launch_core_engines" not in l:
            key = l.strip()[:120]
            if key not in seen:
                seen.add(key); out.append(l.rstrip())
        if len(out) >= 30:
            break
    return out + ["…"] + [l.rstrip() for l in lines[-8:]]


# ─── commands ───────────────────────────────────────────────────────────────────────────────
def cmd_up(a, cfg):
    g = cfg["global"]
    if a.preset not in cfg or a.preset == "global":
        sys.exit(f"unknown preset {a.preset!r}; have: {', '.join(k for k in cfg if k != 'global')}")
    preset = dict(cfg[a.preset])
    preset.update(dict(kv.split("=", 1) for kv in a.set))      # --set gpu-memory-utilization=0.73 --set max-model-len=8192
    preset.setdefault("gpu-memory-utilization", g.get("gpu-memory-utilization", "0.75"))
    envs = {k[4:]: v for k, v in g.items() if k.startswith("env.")}
    envs.update({k[4:]: v for k, v in preset.items() if k.startswith("env.")})
    envs.update(dict(e.split("=", 1) for e in a.env))
    util = float(preset["gpu-memory-utilization"])
    reserve = float(a.reserve if a.reserve is not None else g.get("reserve-gb", 2.0))
    drain = float(g.get("drain-gb", 4.0))

    cmd = run_cmd(g, a.preset, preset, envs)
    print("docker command:\n  " + " ".join(shlex.quote(c) for c in cmd) + "\n")
    if a.dry_run:
        return

    st = container_state(g["container"])
    if st == "running":
        sys.exit(f"container {g['container']} is already running — `vllmctl.py down` first")
    if st:   # leftover from a crash: its log is the evidence, show the tail, then clear the name
        r = docker("logs", "--tail", "15", g["container"], check=False)
        print(f"removing old container ({st}); its last lines were:\n" + "\n".join("    | " + l for l in (r.stdout + r.stderr).strip().splitlines()[-15:]))
        docker("rm", g["container"], check=False)

    # 1. clean card — unload resident router models through its API while it is still up, then drain
    if router_active():
        router_unload(g["router"])
    used = wait_drain(drain, 180, "before start")
    used, total, gtt = vram()
    if used is None or total is None:
        if not a.force:
            sys.exit(f"ABORT: cannot read {DRM}/mem_info_vram_* — no VRAM guard, not starting blind (--force overrides)")
        print("⚠ FORCED without a VRAM guard (sysfs unreadable)"); used, total = 0, 1
    if used > drain * 1e9 and not a.force:
        sys.exit(f"ABORT: card not clean — {gb(used)} VRAM used (> drain-gb {drain}). Leftover model / GPU apps?")

    # 2. projection — abort here costs nothing, llama.service is still running
    budget = util * total
    free_after = total - used - budget
    weights = float(preset.get("x-weights-gb", 0) or 0) * 1e9
    print(f"VRAM: total {gb(total)} · used now {gb(used)} (GTT {gb(gtt)}) · vLLM budget {util:.2f} × total = {gb(budget)}"
          f"\n      → projected free {gb(free_after)} (reserve {reserve} GB)")
    if weights:
        room = budget - weights
        print(f"      budget − weights {gb(weights)} = {gb(room)} for activations + KV cache"
              + ("  ⚠ < 1 GB: expect 'no memory for KV cache' → raise util or lower max-model-len" if room < 1e9 else ""))
    if free_after < reserve * 1e9:
        suggest = max(0.0, (total - used - reserve * 1e9) / total)
        msg = (f"projected free {gb(free_after)} < reserve {reserve} GB. Max safe gpu-memory-utilization right now "
               f"= {suggest:.2f} (desktop is using {gb(used)}).")
        if not a.force:
            sys.exit("ABORT: " + msg + f"\n  → free the desktop (vramx) or `--set gpu-memory-utilization={suggest:.2f}` "
                     f"(+ `--set max-model-len=8192` if weights leave < 1 GB); --force to start anyway.")
        print("⚠ FORCED: " + msg)

    # 3. router gate — the only thing that stops a LAN client from loading a model on top of vLLM
    if router_active():
        if a.stop_router:
            print("stopping llama.service (sudo) …")
            subprocess.run(["sudo", "systemctl", "stop", "llama.service"], check=True)
        elif a.unsafe_keep_router:
            print("⚠ llama.service stays ACTIVE: any client request to :11343 while vLLM runs = GTT freeze.")
        else:
            sys.exit("ABORT: llama.service is active. Any client request to :11343 would load a model on top "
                     "of vLLM and freeze the PC.\n  → re-run with --stop-router (sudo systemctl stop llama.service), "
                     "or --unsafe-keep-router if you really know nothing will call the router.")
        used2, _, _ = vram()
        if used2 is not None and used2 > used + 1e9:   # a client loaded something between projection and stop
            print(f"⚠ VRAM grew to {gb(used2)} meanwhile — vLLM's own free-memory check should refuse; watch the start")

    # 4. start + wait
    cid = subprocess.run(cmd, check=True, capture_output=True, text=True).stdout.strip()
    print(f"started {cid[:12]} — waiting for http://127.0.0.1:{g['port']}/health (timeout {g.get('health-timeout', 1800)} s; "
          f"first start compiles kernels)")
    end = time.time() + int(g.get("health-timeout", 1800))
    last_log, shown = 0, set()
    while time.time() < end:
        st = container_state(g["container"])
        if st != "running":
            r = docker("logs", g["container"], check=False)
            print("root cause candidates (full log: docker logs " + g["container"] + "):")
            for l in root_cause((r.stdout + r.stderr).splitlines()):
                print("    | " + l[:200])
            sys.exit(f"ABORT: container is {st or 'gone'}. If a line names gfx1102 / 'code object' / hipErrorNoBinaryForGpu: "
                     f"uncomment env.HSA_OVERRIDE_GFX_VERSION = 11.0.0 in [global] and retry." + router_hint())
        try:
            if http(f"http://127.0.0.1:{g['port']}/health", timeout=3).status == 200:
                break
        except Exception:
            pass
        if time.time() - last_log > 30:
            r = docker("logs", "--tail", "3", g["container"], check=False)
            tail = (r.stdout + r.stderr).strip().splitlines()[-3:]
            for t in tail:
                if t not in shown:                         # only lines not printed before (compile phases are quiet for minutes)
                    print("    | " + t[:160]); shown.add(t)
            last_log = time.time()
        time.sleep(3)
    else:
        sys.exit("ABORT: no /health within the timeout — `vllmctl.py logs`, then `vllmctl.py down`" + router_hint())
    used, total, gtt = vram()
    print(f"UP. VRAM used {gb(used)} / free {gb(total - used)} · GTT {gb(gtt)}")
    try:
        ids = [m["id"] for m in json.load(http(f"http://127.0.0.1:{g['port']}/v1/models")).get("data", [])]
        print(f"models: {', '.join(ids)}  →  http://192.168.0.204:{g['port']}/v1")
    except Exception as e:
        print(f"(/v1/models not readable yet: {e})")
    print(f"next: ./bench-vs.py -t vllm:{preset.get('served-model-name', a.preset)}   ·   when done: ./vllmctl.py down")


def cmd_down(a, cfg):
    g = cfg["global"]
    st = container_state(g["container"])
    if st:
        print(f"stopping {g['container']} ({st}) …")
        docker("stop", "-t", "60", g["container"], check=False)
        docker("rm", g["container"], check=False)
    else:
        print("container not running")
    used = wait_drain(float(g.get("drain-gb", 4.0)), 120, "after stop")
    used, total, gtt = vram()
    print(f"VRAM used {gb(used)} / free {gb(total - used) if total and used is not None else '-'} · GTT {gb(gtt)}")
    if not router_active():
        print("llama.service is stopped → `sudo systemctl start llama.service` to bring the router back")


def cmd_status(a, cfg):
    g = cfg["global"]
    used, total, gtt = vram()
    print(f"container {g['container']}: {container_state(g['container']) or 'not running'}")
    print(f"VRAM used {gb(used)} / total {gb(total)} · GTT {gb(gtt)}")
    print(f"llama.service: {'ACTIVE' if router_active() else 'stopped'} · loaded: {router_loaded(g['router']) or '-'}")
    try:
        ids = [m["id"] for m in json.load(http(f"http://127.0.0.1:{g['port']}/v1/models", timeout=3)).get("data", [])]
        print(f"vLLM :{g['port']} serving: {', '.join(ids)}")
    except Exception:
        print(f"vLLM :{g['port']}: not answering")


def cmd_logs(a, cfg):
    os.execvp("docker", ["docker", "logs", "-f", "--tail", "100", cfg["global"]["container"]])


def cmd_pull(a, cfg):
    os.execvp("docker", ["docker", "pull", cfg["global"]["image"]])


GPU_CHECK = r"""
import os, torch
print('HSA_OVERRIDE_GFX_VERSION =', os.environ.get('HSA_OVERRIDE_GFX_VERSION', '(unset)'))
print('cuda available:', torch.cuda.is_available())
p = torch.cuda.get_device_properties(0)
print('device:', torch.cuda.get_device_name(0) or '(no name)', p.gcnArchName, f'{p.total_memory / 1e9:.1f} GB')
print('torch kernel …', end=' ', flush=True)   # a hard abort here (no code object for this arch) prints nothing more
x = torch.ones(1024, device='cuda'); torch.cuda.synchronize(); print('ok', x.sum().item())
import vllm._C  # noqa
ops = sorted({s.name.split('::', 1)[1] for s in torch._C._jit_get_all_schemas() if s.name.startswith('_C::')})
print(f'vllm _C ops registered: {len(ops)}')
f16 = dict(device='cuda', dtype=torch.float16)
probes = {'silu_and_mul': lambda: torch.ops._C.silu_and_mul(torch.empty(4, 8, **f16), torch.randn(4, 16, **f16)),
          'gelu_fast':    lambda: torch.ops._C.gelu_fast(torch.empty(4, 8, **f16), torch.randn(4, 8, **f16)),
          'rms_norm':     lambda: torch.ops._C.rms_norm(torch.empty(4, 8, **f16), torch.randn(4, 8, **f16), torch.ones(8, **f16), 1e-6)}
for name, fn in probes.items():
    if name in ops:
        print(f'vllm custom op {name} …', end=' ', flush=True); fn(); torch.cuda.synchronize(); print('ok -> the image has kernels for this arch')
        break
else:
    print('none of the probe ops exist in this build; first ops:', ops[:15])
"""


def cmd_gpu_check(a, cfg):
    """Does the image see the card? Prints torch's device name + arch (needs no model)."""
    g = cfg["global"]
    envs = [x for k, v in g.items() if k.startswith("env.") for x in ("-e", f"{k[4:]}={v}")]
    envs += [x for e in a.env for x in ("-e", e)]
    code = GPU_CHECK
    cmd = ["docker", "run", "--rm", *gpu_flags(), *envs, "--entrypoint", "python3", g["image"], "-u", "-c", code]
    print(" ".join(shlex.quote(c) for c in cmd))
    os.execvp("docker", cmd)


def cmd_exec(a, cfg):
    g = cfg["global"]
    if container_state(g["container"]) == "running":
        os.execvp("docker", ["docker", "exec", "-it", g["container"], *a.cmd])
    envs = [x for k, v in g.items() if k.startswith("env.") for x in ("-e", f"{k[4:]}={v}")]
    os.execvp("docker", ["docker", "run", "--rm", "-it", *gpu_flags(), *envs, "--entrypoint", a.cmd[0], g["image"], *a.cmd[1:]])


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--ini", default=INI)
    sub = ap.add_subparsers(dest="cmd", required=True)
    up = sub.add_parser("up"); up.add_argument("preset")
    up.add_argument("--stop-router", action="store_true", help="sudo systemctl stop llama.service first")
    up.add_argument("--unsafe-keep-router", action="store_true", help="leave llama.service running (freeze risk)")
    up.add_argument("--force", action="store_true", help="ignore the VRAM projection / clean-card checks")
    up.add_argument("--reserve", type=float, help="override reserve-gb")
    up.add_argument("--dry-run", action="store_true", help="print the docker command only")
    up.add_argument("--env", action="append", default=[], metavar="KEY=VAL", help="extra container env (e.g. HSA_OVERRIDE_GFX_VERSION=11.0.0)")
    up.add_argument("--set", action="append", default=[], metavar="KEY=VAL", help="override a preset key for this start (e.g. gpu-memory-utilization=0.73)")
    sub.add_parser("down"); sub.add_parser("status"); sub.add_parser("logs"); sub.add_parser("pull")
    gc = sub.add_parser("gpu-check"); gc.add_argument("--env", action="append", default=[], metavar="KEY=VAL")
    ex = sub.add_parser("exec"); ex.add_argument("cmd", nargs=argparse.REMAINDER)
    a = ap.parse_args()
    cfg = parse_ini(a.ini)
    if "global" not in cfg:
        sys.exit(f"{a.ini}: no [global] section")
    {"up": cmd_up, "down": cmd_down, "status": cmd_status, "logs": cmd_logs, "pull": cmd_pull,
     "gpu-check": cmd_gpu_check, "exec": cmd_exec}[a.cmd](a, cfg)


if __name__ == "__main__":
    main()