#!/usr/bin/env python3
"""Resume one checkpointed YOLO run after container recreation, fail closed."""

from __future__ import annotations

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


def container_running(container: str) -> bool:
    result = subprocess.run(
        ["docker", "inspect", "--format", "{{.State.Running}}", container],
        capture_output=True, text=True, check=False,
    )
    return result.returncode == 0 and result.stdout.strip() == "true"


def training_active(container: str, run_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(
        "train" in line and run_marker in line
        for line in result.stdout.splitlines()
    )


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


def load_completion_command(path: Path) -> list[str]:
    command = json.loads(path.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):
        raise ValueError("completion command must be a non-empty JSON list of non-empty strings")
    return command


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--container", required=True)
    parser.add_argument("--host-run-dir", type=Path, required=True)
    parser.add_argument("--container-checkpoint", required=True)
    parser.add_argument("--run-marker", required=True)
    parser.add_argument("--yolo", default="/opt/geointel/venv/bin/yolo")
    parser.add_argument("--poll-seconds", type=int, default=30)
    parser.add_argument("--max-resumes", type=int, default=20)
    parser.add_argument("--completion-command-json", type=Path)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    if args.poll_seconds < 1 or args.max_resumes < 1:
        raise SystemExit("poll-seconds and max-resumes must be positive")

    state_path = args.host_run_dir / "supervisor-state.json"
    state = {"schema_version": 1, "status": "monitoring", "resume_count": 0}
    if state_path.is_file():
        state.update(json.loads(state_path.read_text(encoding="utf-8")))

    while True:
        if (args.host_run_dir / "results.png").is_file():
            if state.get("completion_handoff_started"):
                state["status"] = "training_finished_handoff_already_started"
                write_state(state_path, state)
                return 0
            if args.completion_command_json:
                try:
                    command = load_completion_command(args.completion_command_json)
                except (OSError, json.JSONDecodeError, ValueError) as exc:
                    state["status"] = "invalid_completion_command"
                    state["completion_handoff_error"] = str(exc)
                    write_state(state_path, state)
                    return 4
                result = subprocess.run(command, check=False)
                if result.returncode != 0:
                    state["status"] = "completion_handoff_failed"
                    state["completion_handoff_returncode"] = result.returncode
                    write_state(state_path, state)
                    return 5
                state["completion_handoff_started"] = True
                state["completion_command"] = command
                state["status"] = "training_finished_handoff_started"
            else:
                state["status"] = "training_finished"
            write_state(state_path, state)
            return 0
        checkpoint = args.host_run_dir / "weights" / "last.pt"
        if not checkpoint.is_file() or checkpoint.stat().st_size < 1024 * 1024:
            state["status"] = "checkpoint_missing_or_incomplete"
            write_state(state_path, state)
            return 2
        if container_running(args.container) and not training_active(args.container, args.run_marker):
            if int(state["resume_count"]) >= args.max_resumes:
                state["status"] = "resume_budget_exhausted"
                write_state(state_path, state)
                return 3
            result = subprocess.run(
                ["docker", "exec", "-d", args.container, args.yolo, "train",
                 f"resume={args.container_checkpoint}", "device=0"],
                check=False,
            )
            if result.returncode == 0:
                state["resume_count"] = int(state["resume_count"]) + 1
                state["status"] = "resumed"
            else:
                state["status"] = "resume_command_failed"
            write_state(state_path, state)
        else:
            state["status"] = "monitoring"
            write_state(state_path, state)
        if args.once:
            return 0
        time.sleep(args.poll_seconds)


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