"""Generation and verification as ONE heavy job, overlapped.

THE IDLE-GPU DEFECT. The P2 queue ran `lean_gen.py` (GPU) and then `verify_generations.py`
(CPU, Lean) back to back. Verification of a 1,952-sample miniF2F run took 1.43 h and the pool
round's would take ~12 h; the four A10Gs sat at 0% for all of it. Running the verifier as a
second job beside the generator is not an option either: /apps/CLAUDE.md rule 2 forbids two heavy
jobs at once, and the repo's flock enforces it.

So this supervisor takes the heavy-job lock ONCE and runs both phases as children that inherit it
(`--inherit-lock`, the mechanism scripts/p2/launch_baseline.py already uses): the generator owns
the GPUs, and every `--every` seconds a verification pass checks the COMPLETE lines written so far.
A pass works on a snapshot that stops at the last newline, because the generator appends while we
read and the strict JSONL reader (p2_common.iter_jsonl) rightly refuses a half-written line.
Verification is keyed and resumable, so each pass only verifies what is new, and the final pass
after the generator exits covers the tail. One lock, one job, both resources busy.

Exit codes: 0 both phases complete; the generator's or verifier's non-zero code otherwise
(143 when stopped by a signal); 4 when the generator exited 0 but its manifest does not say
`complete` (a partial sample set must never read as a finished stage -- U12-2); 5 when GPU memory
did not drain after the generator (an orphan holds it -- U12-3). Re-running the same command
resumes both phases.

After EVERY generator exit (clean, failed or stopped) the generator's process group is reaped by its
exact pgid, orphan `psm_*` /dev/shm segments of this phase are swept, and GPU memory is waited on
until drained, using scripts/p2/launch_baseline.py's helpers (vLLM's EngineCore is documented to
outlive its parent: 19 GB of GPU memory, 0.19-1.4 GB of /dev/shm). The next GPU stage (sft.py,
export, eval generation) therefore starts on a clean host even though this supervisor releases the
lock on exit.
"""

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))

PYR = str(REPO / ".venv-rollout" / "bin" / "python")
PY = str(REPO / ".venv" / "bin" / "python")


_LB = None


def launch_baseline():
    """scripts/p2/launch_baseline.py, loaded by path (the repo idiom): its kill_group / sweep_orphan_shm /
    wait_gpu_drained are the reviewed implementation; a second copy would drift."""
    global _LB
    if _LB is None:
        import importlib.util

        spec = importlib.util.spec_from_file_location("_p2_launch_baseline", REPO / "scripts/p2/launch_baseline.py")
        _LB = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(_LB)
    return _LB


def generation_complete(gen_out: str) -> tuple[bool, str]:
    """Did the LAST lean_gen session on `gen_out` finish every planned sample? (its manifest says so)"""
    mp = gen_out + ".manifest.json"
    if not os.path.exists(mp):
        return False, "no manifest"
    with open(mp, encoding="utf-8") as f:
        sessions = json.load(f).get("sessions", [])
    if not sessions:
        return False, "manifest has no sessions"
    st = sessions[-1].get("status")
    return st == "complete", f"last lean_gen session status {st!r}"


def reap_after_gen(proc, started_at: float, events, lb=None, log=None) -> None:
    """Kill survivors of the generator's process group, sweep its /dev/shm litter, wait for the GPUs to drain."""
    lb = lb or launch_baseline()
    log = log or (lambda **kw: log_event(events, **kw))
    lb.Runner.kill_group(proc, 15.0)
    gone, nbytes = lb.sweep_orphan_shm(started_at)
    if gone:
        log(event="shm_swept", segments=len(gone), bytes=nbytes)
    lb.wait_gpu_drained()
    log(event="gpu_drained")


def snapshot_complete_lines(src: Path, dst: Path) -> int:
    """Copy `src` up to and including its last newline into `dst` (atomically). Returns the
    number of complete lines. A trailing partial line -- the generator mid-append -- is left out."""
    if not src.exists():
        return 0
    data = src.read_bytes()
    cut = data.rfind(b"\n") + 1
    tmp = dst.with_suffix(dst.suffix + ".tmp")
    tmp.write_bytes(data[:cut])
    os.replace(tmp, dst)
    return data[:cut].count(b"\n")


def log_event(path: Path, **kw) -> None:
    kw["t"] = time.strftime("%Y-%m-%d %H:%M:%S")
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(kw) + "\n")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--run-dir", required=True)
    ap.add_argument("--gen-args", required=True, help="arguments for lean_gen.py, as ONE quoted string (must include --out)")
    ap.add_argument("--problems", required=True)
    ap.add_argument("--verify-out", required=True)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--every", type=float, default=300.0, help="seconds between verification passes")
    ap.add_argument("--feedback", action="store_true")
    ap.add_argument("--label", default="p2-gen-verify")
    a = ap.parse_args(argv)

    from align.utils import guard

    run = Path(a.run_dir)
    run.mkdir(parents=True, exist_ok=True)
    events = run / "pipeline.jsonl"
    gen_args = a.gen_args.split()
    gen_out = Path(gen_args[gen_args.index("--out") + 1])
    snap = run / (gen_out.stem + ".snapshot.jsonl")

    guard.acquire_heavy_job_lock(a.label, blocking=True)
    guard.preflight_host_check()  # AFTER the wait: a lock wait can last hours, a stale preflight proves nothing
    me = os.getpid()
    log_event(events, event="lock_acquired", pid=me)

    stop = {"sig": None}

    def on_signal(sig, _frm):
        stop["sig"] = sig

    signal.signal(signal.SIGTERM, on_signal)
    signal.signal(signal.SIGINT, on_signal)

    gen_log = open(run / "gen.log", "ab")
    started_at = time.time() - 1.0
    gen = subprocess.Popen(
        [PYR, str(REPO / "extras/vllm_rollout/lean_gen.py"), *gen_args, "--inherit-lock", str(me)],
        stdout=gen_log, stderr=subprocess.STDOUT, start_new_session=True,
    )
    log_event(events, event="gen_start", pid=gen.pid, args=gen_args)

    def verify_pass(source: Path) -> int:
        cmd = [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", str(source),
               "--problems", a.problems, "--out", a.verify_out, "--workers", str(a.workers),
               "--inherit-lock", str(me)]
        if a.feedback:
            cmd.append("--feedback")
        with open(run / "verify.log", "ab") as vl:
            v = subprocess.Popen(cmd, stdout=vl, stderr=subprocess.STDOUT, start_new_session=True)
            while v.poll() is None:
                if stop["sig"] is not None:
                    os.killpg(v.pid, signal.SIGTERM)
                    v.wait()
                    break
                time.sleep(2)
        return v.returncode

    rc_v = 0
    last_lines = -1
    while gen.poll() is None and stop["sig"] is None:
        n = snapshot_complete_lines(gen_out, snap)
        if n > last_lines and n > 0:
            log_event(events, event="verify_pass", lines=n)
            rc_v = verify_pass(snap)
            log_event(events, event="verify_pass_end", lines=n, rc=rc_v)
            if rc_v != 0 and stop["sig"] is None:
                # a failing verifier must not cost the GPU work: keep generating, stop the
                # overlapped passes, and let the final pass (or a re-run) surface the error
                log_event(events, event="verify_overlap_disabled", rc=rc_v)
                last_lines = float("inf")
                rc_v = 0
                continue
            last_lines = n
        t_end = time.monotonic() + a.every
        while time.monotonic() < t_end and gen.poll() is None and stop["sig"] is None:
            time.sleep(5)

    if stop["sig"] is not None:
        if gen.poll() is None:
            os.killpg(gen.pid, signal.SIGTERM)
            try:
                gen.wait(timeout=120)
            except subprocess.TimeoutExpired:
                os.killpg(gen.pid, signal.SIGKILL)
                gen.wait()
        try:
            reap_after_gen(gen, started_at, events)
        except RuntimeError as e:
            log_event(events, event="gpu_not_drained", error=str(e))
        log_event(events, event="stopped", signal=stop["sig"], verify_rc=rc_v)
        return 143

    rc_g = gen.wait()
    log_event(events, event="gen_end", rc=rc_g)
    try:
        reap_after_gen(gen, started_at, events)
    except RuntimeError as e:
        log_event(events, event="gpu_not_drained", error=str(e))
        return rc_g or 5
    if rc_g != 0:
        return rc_g
    complete, why = generation_complete(gen_out)
    if not complete:
        log_event(events, event="gen_incomplete", detail=why)
        return 4
    # final pass over the complete file (the generator has exited, so every line is whole)
    rc_v = verify_pass(gen_out)
    log_event(events, event="final_verify_end", rc=rc_v)
    if snap.exists():
        snap.unlink()
    return 0 if stop["sig"] is None and rc_v == 0 else (143 if stop["sig"] is not None else rc_v)


if __name__ == "__main__":
    sys.exit(main())