from __future__ import annotations
import argparse
import json
import logging
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from riemann.paths import CHECKPOINT_ROOT

logger = logging.getLogger("riemann.supervise")


def current_step(checkpoint_dir: Path) -> int:
    p = checkpoint_dir / "training_state.json"
    if not p.exists():
        return 0
    try:
        return int(json.loads(p.read_text(encoding="utf-8")).get("global_step", 0))
    except Exception:
        return 0


def main_gpu_present() -> bool:
    try:
        import torch_directml

        if torch_directml.device_count() == 0:
            return False
        name = torch_directml.device_name(0)
        return "880M" not in name
    except Exception:
        try:
            import torch

            return torch.cuda.is_available()
        except Exception:
            return False


def wait_for_gpu(poll_seconds: int) -> None:
    waited = 0
    while not main_gpu_present():
        if waited % (poll_seconds * 6) == 0:
            logger.warning(
                "Main GPU not present; waiting for it to reconnect (%ds so far)", waited
            )
        time.sleep(poll_seconds)
        waited += poll_seconds
    if waited:
        logger.info("Main GPU is back after %ds; resuming", waited)
        time.sleep(10)


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s [supervise] %(message)s"
    )
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", required=True)
    parser.add_argument("--max-steps", type=int, required=True)
    parser.add_argument("--max-restarts", type=int, default=200)
    parser.add_argument("--gpu-poll-seconds", type=int, default=15)
    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
    parser.add_argument("--log-file", default=None)
    parser.add_argument("pretrain_args", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    checkpoint_dir = Path(args.checkpoint_root) / args.stage
    script = Path(__file__).resolve().parent / "pretrain.py"
    extra = [a for a in args.pretrain_args if a != "--"]
    restarts = 0
    while True:
        step = current_step(checkpoint_dir)
        if step >= args.max_steps:
            logger.info("Target reached: %s at step %d", args.stage, step)
            return
        wait_for_gpu(args.gpu_poll_seconds)
        cmd = [
            sys.executable,
            "-u",
            str(script),
            "--stage",
            args.stage,
            "--resume",
            "--max-steps",
            str(args.max_steps),
            "--checkpoint-root",
            args.checkpoint_root,
            *extra,
        ]
        logger.info("Launch %d from step %d: %s", restarts + 1, step, " ".join(cmd[1:]))
        if args.log_file:
            with open(args.log_file, "a", encoding="utf-8") as fh:
                rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT)
        else:
            rc = subprocess.call(cmd)
        new_step = current_step(checkpoint_dir)
        if rc == 0 and new_step >= args.max_steps:
            logger.info("Finished: %s at step %d", args.stage, new_step)
            return
        restarts += 1
        logger.warning(
            "pretrain.py exited with code %d at step %d (restart %d/%d)",
            rc,
            new_step,
            restarts,
            args.max_restarts,
        )
        if restarts >= args.max_restarts:
            logger.error("Too many restarts; giving up at step %d", new_step)
            sys.exit(1)
        time.sleep(20)


if __name__ == "__main__":
    main()