"""Pre-flight and serialization guard for runs that allocate a Metal expert arena.

Every `strata_window_cli` prompt run mmaps `STRATA_METAL_ARENA_GB` of expert slots and, with
`STRATA_METAL_WIRE_ARENA=1`, `mlock`s them.  Wired memory cannot be evicted, so two of those on a
machine that is already serving from a resident engine does not degrade gracefully: the kernel swaps
everything else out and Metal stalls, which reads as "the GPU is dying".  This module exists because
that happened: three arena runs were started before any was awaited, on top of a live 9.27 GB
`--serve` engine, and it took the machine down.

The fix is structural, not a promise to be careful:

    import arena_preflight as ap
    job = await ap.arena_run(*argv, env=env)      # one at a time, env-gated, reported

`arena_run` holds a process-wide asyncio lock, so a concurrent call blocks instead of oversubscribing,
and refuses to start at all when headroom is short unless `force=True`.
"""
import asyncio
import os
import re
import sys

PAGE = 16384  # Apple Silicon
_DEFAULT_ARENA_GB = 6.5
# an arena run also wants its dense weights, the PLE/MTP files, KV, and the framework itself
_OVERHEAD_GB = 4.0
_lock = asyncio.Lock()


def _sh(*argv):
    import subprocess
    return subprocess.run(argv, capture_output=True, text=True).stdout


def processes():
    """Live arena-allocating or arena-holding processes, newest state from ps."""
    out = _sh("ps", "-ax", "-o", "pid,rss,etime,command")
    found = []
    for line in out.splitlines()[1:]:
        parts = line.split(None, 3)
        if len(parts) < 4:
            continue
        pid, rss, etime, cmd = parts
        if "ps -ax" in cmd or "ps -p" in cmd:
            continue
        # arena holders: the engine itself and the serve supervisor that spawns it
        if "strata_window_cli" in cmd or "server.py" in cmd:
            found.append({"pid": int(pid), "rss_gb": int(rss) / 1e6, "etime": etime,
                          "serve": "--serve" in cmd or "server.py" in cmd, "cmd": cmd[:110]})
    return found


def memory():
    out = _sh("vm_stat")
    pages = {}
    for line in out.splitlines():
        m = re.match(r"(.+?):\s+(\d+)\.?$", line.strip())
        if m:
            pages[m.group(1)] = int(m.group(2))
    free_gb = (pages.get("Pages free", 0) + pages.get("Pages speculative", 0)) * PAGE / 1e9
    wired_gb = pages.get("Pages wired down", 0) * PAGE / 1e9
    swapped = _sh("sysctl", "-n", "vm.swapusage").strip()
    used = re.search(r"used = ([\d.]+)M", swapped)
    total = re.search(r"total = ([\d.]+)M", swapped)
    return {
        "free_gb": free_gb,
        "wired_gb": wired_gb,
        "swap_used_gb": float(used.group(1)) / 1024 if used else -1.0,
        "swap_total_gb": float(total.group(1)) / 1024 if total else -1.0,
        "swap_encrypted_unusable": "encrypted" in swapped,
    }


def snapshot(arena_gb=None):
    if arena_gb is None:
        arena_gb = float(os.environ.get("STRATA_METAL_ARENA_GB", _DEFAULT_ARENA_GB))
    procs, mem = processes(), memory()
    held = sum(p["rss_gb"] for p in procs)
    need = arena_gb + _OVERHEAD_GB
    return {"arena_gb": arena_gb, "need_gb": need, "held_gb": held,
            "headroom_gb": mem["free_gb"], "free_after_gb": mem["free_gb"] - need,
            "ok": mem["free_gb"] >= need and mem["swap_used_gb"] < 1.0,
            "processes": procs, "memory": mem}


def report(snap=None):
    s = snap or snapshot()
    print(f"arena preflight: need {s['need_gb']:.1f} GB "
          f"(arena {s['arena_gb']:.1f} + overhead {_OVERHEAD_GB:.1f}), "
          f"free {s['headroom_gb']:.1f} GB, after {s['free_after_gb']:+.1f} GB, "
          f"swap {s['memory']['swap_used_gb']:.2f}/{s['memory']['swap_total_gb']:.2f} GB, "
          f"wired {s['memory']['wired_gb']:.1f} GB")
    if s["processes"]:
        for p in s["processes"]:
            tag = "SERVE" if p["serve"] else "cli"
            print(f"  [{tag}] pid {p['pid']} {p['rss_gb']:.2f} GB up {p['etime']}")
    else:
        print("  no live arena process")
    print(f"  verdict: {'OK to start one arena run' if s['ok'] else 'REFUSE - not enough headroom'}")
    return s["ok"]


async def arena_run(*argv, env=None, arena_gb=None, force=False, quiet=False, tail=40):
    """Run one arena-allocating program to completion, serialized and pre-flighted."""
    from albedo import run  # the session's process launcher
    async with _lock:
        snap = snapshot(arena_gb if arena_gb is not None else
                        (float(env["STRATA_METAL_ARENA_GB"]) if env and "STRATA_METAL_ARENA_GB" in env else None))
        if not quiet:
            report(snap)
        if not snap["ok"] and not force:
            raise RuntimeError(
                "arena_preflight: refusing to start; headroom "
                f"{snap['headroom_gb']:.1f} GB < need {snap['need_gb']:.1f} GB or swap "
                f"{snap['memory']['swap_used_gb']:.2f} GB in use. Free memory, shrink "
                "STRATA_METAL_ARENA_GB, or pass force=True deliberately.")
        job = run(*argv, env=env)
        await job
        if not quiet:
            out = job.read() if job.exit_code == 0 else job.read()
            for line in out.splitlines():
                if "STRATA_RESULT" in line or "expert stream" in line or "error" in line.lower():
                    print("  " + line[:200])
        return job