#!/usr/bin/env python
"""Single-turn on-policy RL for Lean provers (lane R). Library: align/rl/lean_rl.py, align/rl/lean_rl_run.py.

    # the iteration loop (CPU-only parent; spawns generation, verification and training children in series)
    .venv/bin/python scripts/p2/rl_lean.py run configs/p2_rl_lean.toml \\
        --set run_dir="RUN" --set problems="POOL.jsonl" [--nproc 4] [--until K] [--dry-run]

    # one training iteration (what `run` launches under torchrun; also the Phase 0b probe, docs/RL_LEAN.md)
    torchrun --standalone --nproc_per_node=4 scripts/p2/rl_lean.py train-iter CFG --iteration K [--set ...]
    # K = -1 writes checkpoint 0 (fresh adapter, fresh optimizer)

    # build the loss-mode contrast on the SAME cached rollouts (K1): a second run with loss.mode="ei"
    ... run CFG --set run_dir="RUN_EI" --set replay_from="RUN" --set 'loss.mode="ei"'

Resumable: re-running `run` continues from the last committed checkpoint; every stage commits atomically.
The parent never takes the heavy-job lock itself: each child does (lean_gen --wait-lock, verify_generations when
workers > 2, train-iter on rank 0), so exactly one heavy child exists at any time.

IMPORT ORDER IS LOAD-BEARING (as in scripts/sft.py): the guard runs before anything imports torch.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path

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

from align.utils.guard import apply_resource_guards  # noqa: E402


def parse_args(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name in ("run", "train-iter"):
        p = sub.add_parser(name)
        p.add_argument("config")
        p.add_argument("--set", dest="overrides", action="append", default=[], metavar="KEY=TOML_VALUE")
        if name == "run":
            p.add_argument("--nproc", type=int, default=4, help="training ranks (torchrun --nproc_per_node)")
            p.add_argument("--until", type=int, default=None, help="stop after this many iterations (resume continues)")
            p.add_argument("--dry-run", action="store_true", help="print iteration 0's child command lines, touch nothing")
        else:
            p.add_argument("--iteration", type=int, required=True, help="-1 = write checkpoint 0")
    return ap.parse_args(argv)


def cmd_run(a) -> int:
    from align.rl import lean_rl as L
    from align.rl import lean_rl_run as R

    cfg = L.load_config(a.config, a.overrides)
    if not cfg.run_dir or not cfg.problems:
        raise SystemExit('run_dir and problems must be set (--set \'run_dir="..."\' --set \'problems="..."\')')
    gen, ver = R.LeanGenCliGenerator(cfg), R.VerifyCliVerifier(cfg)
    trainer = R.TorchrunTrainer(cfg, os.path.abspath(a.config), a.nproc, a.overrides)
    if a.dry_run:
        d = R.iter_dir(cfg.run_dir, 0)
        print(json.dumps({
            "fingerprint": cfg.fingerprint(),
            "init": R.torchrun_command(trainer.cfg_path, -1, a.nproc, a.overrides),
            "gen": gen.command(0, d, os.path.join(d, "problems.jsonl"), os.path.join(R.ckpt_path(cfg.run_dir, 0), "adapter"),
                               os.path.join(R.ckpt_path(cfg.run_dir, 0), "merged") if cfg.export_merged else None,
                               R.derive_seed(cfg.seed, 0)),
            "verify": ver.command(d, os.path.join(d, "problems.jsonl")),
            "train": R.torchrun_command(trainer.cfg_path, 0, a.nproc, a.overrides)}, indent=1))
        return 0
    ledger = R.run_loop(cfg, gen, ver, trainer, until=a.until)
    print(json.dumps({"iterations_committed": len(ledger), "cumulative": ledger[-1]["cumulative"] if ledger else {}}))
    return 0


def cmd_train_iter(a) -> int:
    from align.rl import lean_rl as L

    cfg = L.load_config(a.config, a.overrides)
    apply_resource_guards(cpu_threads=8, heavy=True, label=f"rl-lean-train-{a.iteration}")
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    from align.utils.dist import init_distributed

    init_distributed(run_dir=os.path.join(cfg.run_dir, "nccl"))
    from align.rl import lean_rl_run as R

    trainer = R.InProcessTrainer(cfg, lambda: R.build_real_policy(cfg))
    if a.iteration < 0:
        trainer.init()
    else:
        R.clean_partial(cfg.run_dir) if R._rank() == 0 else None
        trainer.train(a.iteration)
    return 0


def main(argv=None) -> int:
    a = parse_args(argv)
    if a.cmd == "run":
        apply_resource_guards(cpu_threads=2, heavy=False, label="rl-lean-orchestrator")
        return cmd_run(a)
    return cmd_train_iter(a)


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