#!/usr/bin/env python
"""Bounded-retry launch supervisor CLI (stdlib only; no torch import).

    python scripts/supervise.py --run-dir RUN [--max-retries 2] [--backoff 30] \
        -- torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml

Everything after `--` is the child command. See align/utils/supervisor.py for
the retry rules (progress + preflight + non-HALT exit code) and why each exists.
`--run-dir` must be the directory the Trainer writes history.jsonl and
checkpoints/ into; attempts are logged to RUN/supervisor.jsonl.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

# The package is not installed in the venv; without this a bare `python scripts/supervise.py` (the documented
# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py).
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.supervisor import Supervisor


def _checkpoint_validator():
    """`checkpoint_problem` from the Trainer, imported lazily (it needs torch).

    The supervisor module itself stays stdlib-only; THIS entry point is where
    the cost is paid, once, so that "progress" means a checkpoint `latest_checkpoint`
    would actually resume from. If torch is unavailable the supervisor still
    runs, falling back to name-only progress, and says so on stderr.
    """
    try:
        from align.train.trainer import checkpoint_problem
    except Exception as e:  # noqa: BLE001
        print(f"supervisor: checkpoint validation unavailable ({type(e).__name__}: {e}); "
              "progress is judged by directory name only", file=sys.stderr)
        return None
    return checkpoint_problem


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--run-dir", required=True)
    ap.add_argument("--max-retries", type=int, default=2)
    ap.add_argument("--backoff", type=float, default=30.0, help="base seconds; doubles per retry")
    ap.add_argument("--grace", type=float, default=20.0, help="SIGTERM->SIGKILL grace seconds")
    ap.add_argument(
        "--no-validate-checkpoints",
        action="store_true",
        help="count a checkpoint dir as progress by NAME only (skips importing torch for the "
        "completeness check); a truncated newest dir then counts as progress",
    )
    ap.add_argument(
        "--retry-on-stop",
        action="store_true",
        help="also relaunch a run that was stopped by SIGTERM/SIGINT/SIGHUP (exit 143/130/129). "
        "Default: an operator's stop is a decision, not a fault, and ends the supervisor",
    )
    ap.add_argument("cmd", nargs=argparse.REMAINDER, help="-- child command")
    a = ap.parse_args(argv)
    cmd = a.cmd[1:] if a.cmd and a.cmd[0] == "--" else a.cmd
    if not cmd:
        ap.error("no child command given (put it after --)")
    res = Supervisor(cmd, a.run_dir, max_retries=a.max_retries,
                     backoff_s=a.backoff, grace_s=a.grace, retry_on_stop=a.retry_on_stop,
                     checkpoint_validator=None if a.no_validate_checkpoints
                     else _checkpoint_validator()).run()
    print(f"supervisor: {res.reason} after {res.attempts} attempt(s), "
          f"child exits {res.child_codes}", file=sys.stderr)
    return res.exit_code


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