#!/usr/bin/env python3
"""Keep a checkpoint-aware training-loop parent alive across container recreation."""

from __future__ import annotations

import argparse
import json
import subprocess
import time
from datetime import UTC, datetime
from pathlib import Path

TERMINAL_STATUSES = {"training_complete"}


def read_loop_status(path: Path) -> str | None:
    if not path.is_file():
        return None
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return "invalid"
    return value.get("status") if isinstance(value, dict) else "invalid"


def process_active(container: str, marker: str) -> bool:
    result = subprocess.run(
        ["docker", "top", container, "-eo", "pid,args"],
        capture_output=True, text=True, check=False,
    )
    return result.returncode == 0 and any(marker in line for line in result.stdout.splitlines())


def write_state(path: Path, state: dict) -> None:
    state["updated_at"] = datetime.now(UTC).isoformat()
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(state, indent=2), encoding="utf-8")
    temporary.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--container", required=True)
    parser.add_argument("--process-marker", required=True)
    parser.add_argument("--loop-state", type=Path, required=True)
    parser.add_argument("--command-json", type=Path, required=True)
    parser.add_argument("--host-script", type=Path, required=True)
    parser.add_argument("--container-script", required=True)
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--poll-seconds", type=int, default=30)
    parser.add_argument("--max-launches", type=int, default=20)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    state = {"schema_version": 1, "status": "monitoring", "launch_count": 0}
    if args.state.is_file():
        state.update(json.loads(args.state.read_text(encoding="utf-8")))
    while True:
        loop_status = read_loop_status(args.loop_state)
        state["loop_status"] = loop_status
        if loop_status in TERMINAL_STATUSES:
            state["status"] = "loop_finished"
            write_state(args.state, state)
            return 0
        if loop_status == "invalid":
            state["status"] = "invalid_loop_state"
            write_state(args.state, state)
            return 2
        if not process_active(args.container, args.process_marker):
            if int(state["launch_count"]) >= args.max_launches:
                state["status"] = "launch_budget_exhausted"
                write_state(args.state, state)
                return 3
            command = json.loads(args.command_json.read_text(encoding="utf-8"))
            if not isinstance(command, list) or not command or not all(isinstance(x, str) and x for x in command):
                state["status"] = "invalid_command"
                write_state(args.state, state)
                return 4
            copied = subprocess.run(
                ["docker", "cp", str(args.host_script), f"{args.container}:{args.container_script}"],
                check=False,
            )
            launched = subprocess.run(command, check=False) if copied.returncode == 0 else copied
            if launched.returncode != 0:
                state["status"] = "launch_failed"
                write_state(args.state, state)
                return 5
            state["launch_count"] = int(state["launch_count"]) + 1
            state["status"] = "loop_relaunched"
        else:
            state["status"] = "monitoring"
        write_state(args.state, state)
        if args.once:
            return 0
        time.sleep(args.poll_seconds)


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