#!/usr/bin/env python
"""P2: the baseline miniF2F pass@8 run (cap 16384, then the 40960 reference) as ONE supervised, resumable job.

    .venv/bin/python scripts/p2/launch_baseline.py --run-dir .tmp/p2/<run>/baseline \\
        --problems .tmp/p2/problems/minif2f.jsonl            # add --print to show the commands only

What this process IS: the heavy job. It takes the repo heavy-job flock ONCE (waiting if another job
holds it), keeps it for the whole sequence, and runs each phase as a child that *inherits* it
(`--inherit-lock <this pid>`; `align.utils.guard.inherit_heavy_job_lock` verifies the claim).
Phases, in this order so the pre-registered primary result exists first:

    gen_c16384  ->  ver_c16384  ->  gen_c40960  ->  ver_c40960

Generation and Lean verification are both heavy; they never overlap, and no other agent's job can slip
in between phases, because the lock is not released between them.

Properties (each one is a failure this job would otherwise have):
  * RESUMABLE at every level. Re-running the same command skips phases with a `.done` marker; a phase
    killed half way resumes from its append-only output (both children skip finished keys).
  * COMPLETENESS IS CHECKED, not assumed: a generation phase is done only when it holds exactly
    problems x n records at that cap; a verification phase only when its key set equals the generation's.
    A child that exits 0 with fewer records is retried, not accepted.
  * BOUNDED RETRIES that need PROGRESS: a failed phase is retried (max `--max-retries`) only if its
    output grew since the previous attempt; a crash loop with no progress stops with `stalled`.
  * NO ORPHANS: every child runs in its own session; after it exits (cleanly or not) its process group
    is signalled by its exact pgid (vLLM's EngineCore is known to outlive its parent and keep 19 GB of
    GPU memory), and the next GPU phase waits for GPU memory to drain. Nothing is ever killed by a
    name pattern.
  * SIGTERM/SIGINT to this process stops the current child (SIGTERM, grace, SIGKILL of its group),
    writes `stopped` to the status log and exits 143; the lock goes with the process.
  * Host rules: preflight (strict) before every phase; scratch and outputs under /apps; vLLM TP=4,
    `--max-num-seqs 128` (B's measured knee); <= 8 Lean workers.

Status: `<run-dir>/supervisor.jsonl` (one event per line) and `<run-dir>/logs/<phase>.log`.
"""

from __future__ import annotations

import argparse
import json
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))

PY = str(REPO / ".venv" / "bin" / "python")
PYR = str(REPO / ".venv-rollout" / "bin" / "python")
MODEL = ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
         "dfd02e6271a58375dfbf3ece0175277cf6b6a89a")
KEY = ("problem_id", "sample_idx", "round", "cap_tokens")


class Stalled(RuntimeError):
    pass


def phase_plan(a, parent_pid: int) -> list[dict]:
    """The ordered phases and their exact argv. Pure (tests pin it)."""
    d = Path(a.run_dir)
    plan = []
    for cap in a.caps:
        gen, ver = d / f"gen_minif2f_c{cap}.jsonl", d / f"ver_minif2f_c{cap}.jsonl"
        plan.append({
            "name": f"gen_c{cap}", "kind": "gen", "out": str(gen), "cap": cap,
            "cmd": [PYR, str(REPO / "extras/vllm_rollout/lean_gen.py"), "--model", a.model, "--problems", a.problems,
                    "--out", str(gen), "--n", str(a.n), "--cap-tokens", str(cap), "--round", "0", "--tp", str(a.tp),
                    "--seed", "0", "--max-num-seqs", str(a.max_num_seqs), "--inherit-lock", str(parent_pid),
                    "--label", f"p2-baseline-gen-c{cap}"],
        })
        plan.append({
            "name": f"ver_c{cap}", "kind": "ver", "out": str(ver), "gen": str(gen), "cap": cap,
            "cmd": [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", str(gen), "--problems", a.problems,
                    "--out", str(ver), "--workers", str(a.workers), "--inherit-lock", str(parent_pid)],
        })
    return plan


def count_keys(path: str, cap: int | None = None) -> set[tuple]:
    keys = set()
    p = Path(path)
    if not p.is_file():
        return keys
    with open(p, "rb") as f:
        for line in f:
            if not line.endswith(b"\n"):
                break  # torn tail of a killed writer: the child repairs it; do not count it
            r = json.loads(line)
            if cap is None or r.get("cap_tokens") == cap:
                keys.add(tuple(r[k] for k in KEY))
    return keys


def phase_complete(ph: dict, n_problems: int, n: int) -> tuple[bool, str]:
    """Is the phase's OUTPUT complete? gen: problems x n records at this cap. ver: same key set as gen."""
    keys = count_keys(ph["out"], ph["cap"])
    if ph["kind"] == "gen":
        want = n_problems * n
        # prompts too long for the window are skipped by lean_gen and listed in <out>.skipped.jsonl
        skipped = len(list(open(ph["out"] + ".skipped.jsonl"))) if os.path.isfile(ph["out"] + ".skipped.jsonl") else 0
        return len(keys) + skipped >= want and len(keys) > 0, f"{len(keys)} records + {skipped} skipped of {want}"
    gkeys = count_keys(ph["gen"], ph["cap"])
    return bool(gkeys) and keys == gkeys, f"{len(keys)} verdicts for {len(gkeys)} generations"


def n_problems_of(path: str) -> int:
    return sum(1 for l in open(path) if l.strip())


def gpu_memory_used_mib() -> list[int]:
    out = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
                         capture_output=True, text=True, timeout=30).stdout.split()
    return [int(x) for x in out]


def wait_gpu_drained(limit_mib: int = 1024, timeout_s: float = 180.0) -> None:
    t0 = time.time()
    while True:
        used = gpu_memory_used_mib()
        if max(used) <= limit_mib:
            return
        if time.time() - t0 > timeout_s:
            raise RuntimeError(f"GPU memory did not drain after the previous phase: {used} MiB (an orphan holds it; "
                               "find it by pid with nvidia-smi and stop it by exact pid)")
        time.sleep(5)


SHM_DIR = "/dev/shm"
SHM_NAME_RE = r"psm_[0-9a-f]{8}"


def _paths_in_use(shm_dir: str) -> set[str]:
    """Every file under `shm_dir` that some live process has open (fd) or mapped (maps)."""
    held: set[str] = set()
    prefix = shm_dir.rstrip("/") + "/"
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        try:
            for fd in os.listdir(f"/proc/{entry}/fd"):
                try:
                    t = os.readlink(f"/proc/{entry}/fd/{fd}")
                except OSError:
                    continue
                if t.startswith(prefix):
                    held.add(t.split(" (deleted)")[0])
            with open(f"/proc/{entry}/maps") as f:
                for line in f:
                    i = line.find(prefix)
                    if i >= 0:
                        held.add(line[i:].split(" (deleted)")[0].strip())
        except (FileNotFoundError, ProcessLookupError, PermissionError):
            continue
    return held


def sweep_orphan_shm(since: float, shm_dir: str = SHM_DIR, only: "set[str] | None" = None) -> tuple[list[str], int]:
    """Unlink the `psm_*` shared-memory segments a vLLM phase left behind. Returns (names, bytes).

    THE /dev/shm LEAK (R2 review, measured): every vLLM launch leaves its broadcast ring buffers in /dev/shm
    (a clean exit: 2 segments, ~0.19 GB; a SIGKILLed one: ~1.4 GB), and nothing removes them -- 46 segments /
    6.3 GB were sitting there from P2 alone. tmpfs is RAM on a box with NO swap, and /dev/shm exhaustion is one
    of the two 2026-08-12 incidents. Only segments that are provably this phase's litter are removed: named
    exactly like Python's SharedMemory (`psm_` + 8 hex), owned by this uid, modified after the phase began, and
    held (open or mapped) by NO process. Nothing is matched by pattern against processes. `only` restricts the
    candidates to the given names (tests use it so that a mutated copy cannot touch other jobs' segments).
    """
    import re

    removed, total = [], 0
    held = _paths_in_use(shm_dir)
    for name in sorted(os.listdir(shm_dir)):
        if only is not None and name not in only:
            continue
        if not re.fullmatch(SHM_NAME_RE, name):
            continue
        path = os.path.join(shm_dir, name)
        try:
            st = os.stat(path)
        except FileNotFoundError:
            continue
        if st.st_uid != os.getuid() or st.st_mtime < since or path in held:
            continue
        try:
            os.unlink(path)
        except FileNotFoundError:
            continue
        removed.append(name)
        total += st.st_size
    return removed, total


class Runner:
    def __init__(self, a):
        self.a = a
        self.run_dir = Path(a.run_dir)
        self.child: subprocess.Popen | None = None
        self.stop = False

    def log(self, **ev) -> None:
        ev["t"] = time.strftime("%F %T")
        with open(self.run_dir / "supervisor.jsonl", "a") as f:
            f.write(json.dumps(ev) + "\n")
        print(json.dumps(ev), flush=True)

    def _signal(self, signum, _frame) -> None:
        self.stop = True
        self.log(event="signal", signum=signum)
        self.kill_group(self.child, grace=60.0)

    @staticmethod
    def group_alive(pgid: int) -> bool:
        """Does any NON-ZOMBIE process still belong to process group `pgid`?

        `os.killpg(pgid, 0)` is the wrong probe: a killed leader stays a zombie until its parent reaps it,
        and the SIGTERM handler runs in the very thread that is blocked in `Popen.wait()` (which owns the
        reaping), so the probe saw the dead leader as "alive" for the whole grace period -- every stop took
        grace + 10 s (70 s) with the heavy-job lock held, whatever the child did."""
        for entry in os.listdir("/proc"):
            if not entry.isdigit():
                continue
            try:
                with open(f"/proc/{entry}/stat") as f:
                    rest = f.read().rsplit(")", 1)[1].split()
            except (FileNotFoundError, ProcessLookupError, IndexError):
                continue
            if int(rest[2]) == pgid and rest[0] != "Z":  # rest = [state, ppid, pgrp, ...]
                return True
        return False

    @staticmethod
    def kill_group(proc, grace: float) -> None:
        """SIGTERM then SIGKILL the child's process group, by the exact pgid it was given (its own session)."""
        if proc is None:
            return
        for sig, wait in ((signal.SIGTERM, grace), (signal.SIGKILL, 10.0)):
            if not Runner.group_alive(proc.pid):
                return
            try:
                os.killpg(proc.pid, sig)
            except ProcessLookupError:
                return
            t0 = time.time()
            while time.time() - t0 < wait:
                if not Runner.group_alive(proc.pid):
                    return
                time.sleep(0.2)

    def run_child(self, ph: dict, attempt: int) -> int:
        log_path = self.run_dir / "logs" / f"{ph['name']}.log"
        env = dict(os.environ)
        env.update({"TMPDIR": str(self.run_dir / "scratch"), "HF_HOME": "/apps/development/rlhf/hf_cache",
                    "HF_HUB_DISABLE_XET": "1", "PYTHONUNBUFFERED": "1"})
        for k in [k for k in env if k.startswith(("RANK", "LOCAL_RANK", "WORLD_SIZE"))]:
            del env[k]
        started_at = time.time() - 1.0
        with open(log_path, "a") as lf:
            lf.write(f"\n=== attempt {attempt} {time.strftime('%F %T')}: {' '.join(ph['cmd'])}\n")
            lf.flush()
            self.child = subprocess.Popen(ph["cmd"], stdout=lf, stderr=subprocess.STDOUT, env=env, cwd=str(REPO),
                                          start_new_session=True)
            if self.stop:
                # A stop signal that landed between Popen() starting the child and the assignment above found
                # `self.child` still None, so the handler had nothing to kill: without this the supervisor
                # would sit in wait() for the whole phase (hours) while holding the heavy-job lock.
                self.kill_group(self.child, grace=60.0)
            rc = self.child.wait()
        # the leader is gone; anything still in its group is an orphan of THIS child
        self.kill_group(self.child, grace=15.0)
        self.child = None
        gone, nbytes = sweep_orphan_shm(started_at)
        if gone:
            self.log(event="shm_swept", phase=ph["name"], segments=len(gone), bytes=nbytes)
        return rc

    def run_phase(self, ph: dict, n_problems: int) -> None:
        marker = self.run_dir / f"{ph['name']}.done"
        if marker.is_file():
            self.log(event="skip_done", phase=ph["name"])
            return
        last = -1
        for attempt in range(self.a.max_retries + 1):
            if self.stop:
                return
            from align.utils import guard

            guard.preflight_host_check(strict=True)
            if ph["kind"] == "gen":
                wait_gpu_drained()
            t0 = time.time()
            self.log(event="start", phase=ph["name"], attempt=attempt)
            rc = self.run_child(ph, attempt)
            ok, detail = phase_complete(ph, n_problems, self.a.n)
            size = len(count_keys(ph["out"], ph["cap"]))
            self.log(event="end", phase=ph["name"], attempt=attempt, rc=rc, complete=ok, detail=detail,
                     wall_s=round(time.time() - t0, 1))
            if rc == 0 and ok:  # a phase that finished while a stop was pending is still finished
                marker.write_text(json.dumps({"detail": detail, "t": time.strftime("%F %T")}))
                return
            if self.stop:
                return
            if size <= last:
                raise Stalled(f"{ph['name']}: attempt {attempt} made no progress ({detail}); stopping instead of looping")
            last = size
            time.sleep(self.a.backoff)
        raise Stalled(f"{ph['name']}: still incomplete after {self.a.max_retries + 1} attempts ({detail})")

    def summary(self) -> dict:
        out = {}
        for cap in self.a.caps:
            vp = self.run_dir / f"ver_minif2f_c{cap}.jsonl"
            if not vp.is_file():
                continue
            solved, seen, acc = set(), set(), 0
            for l in open(vp):
                r = json.loads(l)
                seen.add(r["problem_id"])
                if r["accepted"] and r["sample_idx"] < self.a.n:
                    solved.add(r["problem_id"])
                acc += int(r["accepted"])
            out[f"solved@{self.a.n}_cap{cap}"] = f"{len(solved)}/{len(seen)}"
            out[f"accepted_samples_cap{cap}"] = acc
        return out

    def run(self) -> int:
        from align.utils import guard

        (self.run_dir / "logs").mkdir(parents=True, exist_ok=True)
        (self.run_dir / "scratch").mkdir(parents=True, exist_ok=True)
        guard.assert_paths_off_root([str(self.run_dir), str(self.run_dir / "scratch")])
        signal.signal(signal.SIGTERM, self._signal)
        signal.signal(signal.SIGINT, self._signal)
        t0 = time.time()
        guard.acquire_heavy_job_lock("p2-baseline-supervisor", blocking=True, lock_path=getattr(self.a, "lock_path", None))  # waits; held until this process exits
        self.log(event="lock_acquired", waited_s=round(time.time() - t0, 1), pid=os.getpid())
        n_problems = n_problems_of(self.a.problems)
        try:
            for ph in phase_plan(self.a, os.getpid()):
                self.run_phase(ph, n_problems)
                if self.stop:
                    self.log(event="stopped")
                    return 143
        except Stalled as e:
            self.log(event="stalled", error=str(e))
            return 3
        self.log(event="complete", **self.summary())
        return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--run-dir", required=True)
    ap.add_argument("--problems", required=True, help="miniF2F Problem JSONL (all 244)")
    ap.add_argument("--model", default=MODEL)
    ap.add_argument("--caps", type=int, nargs="+", default=[16384, 40960])
    ap.add_argument("--n", type=int, default=8)
    ap.add_argument("--tp", type=int, default=4)
    ap.add_argument("--max-num-seqs", type=int, default=128)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--max-retries", type=int, default=2)
    ap.add_argument("--backoff", type=float, default=60.0)
    ap.add_argument("--print", action="store_true", dest="print_only")
    a = ap.parse_args(argv)
    if not 1 <= a.workers <= 8:
        raise SystemExit("--workers must be in [1, 8] (host rule)")
    if a.n != 8 or a.caps[0] != 16384:
        print("note: the registered primary endpoint is n=8 at cap 16384 (docs/P2.md section 1)", file=sys.stderr)
    for p in (a.run_dir,):
        if not os.path.abspath(p).startswith("/apps/"):
            raise SystemExit(f"--run-dir {p} must be on /apps (never / or /tmp)")
    if a.print_only:
        for ph in phase_plan(a, os.getpid()):
            print(" ".join(ph["cmd"]))
        return 0
    return Runner(a).run()


if __name__ == "__main__":
    raise SystemExit(main())