"""
hegm.ctrl -- External control CLI and REST shim.

This module provides:

- ``hegm-ctrl`` CLI with:
  - ``checkpoint`` subcommand to send SIGUSR1 to target PIDs or all
    GPU-using PIDs.
  - ``resume`` subcommand to remove lock files under ``/checkpoint/<PPID>/lock``
    for specific PPIDs or for all pending checkpoints.
- A lightweight HTTP server that exposes the same actions for external
  controllers to call via REST. The server is intended to run *inside* the
  training container; the controller talks to it over the Pod IP.
"""

import argparse
import glob
import json
import os
import signal
import subprocess
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Dict, List, Optional, Tuple

from hegm._config import EXIT_CODE, log, warn

# Base directory for checkpoints (mirrors launcher.py behaviour).
_BASE_DIR: str = os.environ.get("HEGM_CHECKPOINT_DIR", "/checkpoint")

# Default control port for the in-container REST API.
_DEFAULT_CTRL_PORT: int = int(os.environ.get("HEGM_CTRL_PORT", "8298"))

# How long the controller should wait (in seconds) for checkpoint locks to
# appear after signalling workers.
_CHECKPOINT_WAIT_TIMEOUT: float = float(
    os.environ.get("HEGM_CTRL_CHECKPOINT_TIMEOUT", "300")
)

# Poll interval used while waiting for locks.
_CHECKPOINT_WAIT_INTERVAL: float = float(
    os.environ.get("HEGM_POLL_INTERVAL", "1")
)

_api_server: Optional[ThreadingHTTPServer] = None
_api_thread: Optional[threading.Thread] = None
_api_lock = threading.Lock()


# ---------------------------------------------------------------------------
# GPU PID discovery
# ---------------------------------------------------------------------------


def _gpu_pids_via_nvidia_smi() -> List[int]:
    """Best-effort discovery of GPU-using PIDs via ``nvidia-smi``.

    ``nvidia-smi`` reports **host** PIDs, which do not match container PIDs
    inside PID namespaces. We translate host PIDs to the current namespace
    using the ``NSpid`` field in ``/proc/<pid>/status`` when available.

    Returns an empty list if ``nvidia-smi`` is unavailable or fails.
    """
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-compute-apps=pid",
                "--format=csv,noheader,nounits",
            ],
            stderr=subprocess.DEVNULL,
            text=True,
        )
    except FileNotFoundError:
        warn("nvidia-smi not found in PATH; hegm-ctrl checkpoint --all requires it")
        return []
    except Exception:
        return []

    host_pids: List[int] = []
    for line in out.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            host_pids.append(int(line))
        except ValueError:
            continue

    # Build a mapping from host PID -> PID in this namespace using NSpid.
    host_to_local: Dict[int, int] = {}
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        status_path = os.path.join("/proc", entry, "status")
        try:
            with open(status_path, "r", encoding="utf-8") as fh:
                for line in fh:
                    if line.startswith("NSpid:"):
                        parts = line.split()
                        # Format: "NSpid:\t<host> <...> <local>"
                        try:
                            ids = [int(p) for p in parts[1:]]
                        except ValueError:
                            break
                        if len(ids) >= 2:
                            host_pid = ids[0]
                            local_pid = ids[-1]
                            host_to_local[host_pid] = local_pid
                        break
        except FileNotFoundError:
            # Process exited between listing and open.
            continue
        except Exception:
            continue

    local_pids: List[int] = []
    for hpid in host_pids:
        lpid = host_to_local.get(hpid, hpid)
        # Only keep if it looks alive in this namespace.
        if os.path.exists(os.path.join("/proc", str(lpid))):
            local_pids.append(lpid)

    return sorted(set(local_pids))


def _worker_pids_fallback() -> List[int]:
    """Fallback discovery of candidate worker PIDs inside the container.

    When ``nvidia-smi`` is unavailable or reports no GPU PIDs, we approximate
    "GPU-using PIDs" as the long-lived training workers inside the container.
    In our examples these are the ``python -u /workspace/train.py`` processes.

    We look for any process whose cmdline contains both ``python`` and
    ``train.py``, skipping obvious utilities.
    """
    candidates: List[int] = []
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        pid = int(entry)
        # Skip special/system PIDs.
        if pid <= 1:
            continue
        cmdline_path = os.path.join("/proc", entry, "cmdline")
        try:
            with open(cmdline_path, "rb") as fh:
                raw = fh.read()
        except FileNotFoundError:
            # Process exited.
            continue
        except Exception:
            continue

        # /proc/<pid>/cmdline is NUL-separated.
        parts = [p.decode("utf-8", "ignore") for p in raw.split(b"\0") if p]
        if not parts:
            continue
        cmd = " ".join(parts)

        # Heuristics: skip the launcher itself and common trivial tools.
        if "hegm-launcher" in cmd:
            continue
        if cmd.startswith("ps "):
            continue

        # Treat python train scripts as workers.
        if "python" in cmd and "train.py" in cmd:
            candidates.append(pid)

    return sorted(set(candidates))


# ---------------------------------------------------------------------------
# Checkpoint / resume primitives
# ---------------------------------------------------------------------------


def _launcher_pid_for_worker(pid: int) -> Optional[int]:
    """Return the parent PID (Launcher) for a given worker PID, if available."""
    status_path = os.path.join("/proc", str(pid), "status")
    try:
        with open(status_path, "r", encoding="utf-8") as fh:
            for line in fh:
                if line.startswith("PPid:"):
                    parts = line.split()
                    if len(parts) >= 2:
                        return int(parts[1])
                    break
    except FileNotFoundError:
        # Worker already exited.
        return None
    except Exception:
        return None
    return None


def _wait_for_parent_locks(parent_pids: List[int]) -> Dict[int, str]:
    """Block until lock files appear for all given parent PIDs or timeout.

    Returns a mapping of parent PID -> status string:
        - "lock-ready"            : lock file observed
        - "launcher-exited"       : parent process disappeared without a lock
        - "timeout-waiting-lock"  : timeout reached without lock or exit
    """
    statuses: Dict[int, str] = {}
    pending = set(parent_pids)
    deadline = time.time() + _CHECKPOINT_WAIT_TIMEOUT

    while pending and time.time() < deadline:
        finished: List[int] = []
        for ppid in list(pending):
            lock_path = _lock_path_for_ppid(ppid)
            if os.path.exists(lock_path):
                statuses[ppid] = "lock-ready"
                finished.append(ppid)
                continue

            # If the Launcher is gone and no lock is present, treat as failure.
            if not os.path.exists(os.path.join("/proc", str(ppid))):
                statuses[ppid] = "launcher-exited"
                finished.append(ppid)

        for ppid in finished:
            pending.discard(ppid)

        if pending:
            time.sleep(_CHECKPOINT_WAIT_INTERVAL)

    for ppid in pending:
        statuses[ppid] = "timeout-waiting-lock"

    return statuses


def _lock_path_for_ppid(ppid: int) -> str:
    """Return the lock file path for a given parent PID."""
    return os.path.join(_BASE_DIR, str(ppid), "lock")


def checkpoint_pids(pids: List[int]) -> Dict[int, str]:
    """Send SIGUSR1 to each target PID and wait for checkpoints."""
    results: Dict[int, str] = {}
    for pid in pids:
        try:
            os.kill(pid, signal.SIGUSR1)
            results[pid] = "signalled"
        except ProcessLookupError:
            results[pid] = "no-such-process"
        except PermissionError:
            results[pid] = "permission-denied"
        except Exception as exc:  # pragma: no cover - defensive
            results[pid] = f"error: {exc}"

    # Derive the parent (Launcher) PIDs for all successfully signalled workers.
    parent_pids: List[int] = []
    pid_to_parent: Dict[int, int] = {}
    for pid, status in results.items():
        if status != "signalled":
            continue
        ppid = _launcher_pid_for_worker(pid)
        if ppid is None:
            continue
        parent_pids.append(ppid)
        pid_to_parent[pid] = ppid

    if parent_pids:
        parent_statuses = _wait_for_parent_locks(sorted(set(parent_pids)))
        # Propagate parent lock statuses back to worker PIDs.
        for pid, ppid in pid_to_parent.items():
            p_status = parent_statuses.get(ppid)
            if p_status == "lock-ready":
                results[pid] = "checkpoint-ready"
            elif p_status is not None:
                results[pid] = f"signalled ({p_status})"

    return results


def checkpoint_all_gpu() -> Dict[int, str]:
    """Discover all GPU-using PIDs and send SIGUSR1 to each."""
    pids = _gpu_pids_via_nvidia_smi()
    if not pids:
        # Fallback: approximate workers as children of the Launcher (PID 1).
        pids = _worker_pids_fallback()
    if not pids:
        return {}
    return checkpoint_pids(pids)


def pending_ppids() -> List[int]:
    """Return PPIDs that currently have a checkpoint lock file."""
    pattern = os.path.join(_BASE_DIR, "*", "lock")
    ppids: List[int] = []
    for lock_path in glob.glob(pattern):
        parent = os.path.dirname(lock_path)
        name = os.path.basename(parent)
        try:
            ppids.append(int(name))
        except ValueError:
            continue
    return sorted(set(ppids))


def resume_ppids(ppids: List[int]) -> Dict[int, str]:
    """Remove lock files for the given PPIDs to allow resume."""
    results: Dict[int, str] = {}
    for ppid in ppids:
        lock_path = _lock_path_for_ppid(ppid)
        if not os.path.exists(lock_path):
            results[ppid] = "no-lock"
            continue
        try:
            os.remove(lock_path)
            results[ppid] = "lock-removed"
        except Exception as exc:  # pragma: no cover - defensive
            results[ppid] = f"error: {exc}"
    return results


def resume_all_pending() -> Dict[int, str]:
    """Remove all existing checkpoint locks under the base directory."""
    ppids = pending_ppids()
    if not ppids:
        return {}
    return resume_ppids(ppids)


# ---------------------------------------------------------------------------
# REST API server
# ---------------------------------------------------------------------------


class _CtrlRequestHandler(BaseHTTPRequestHandler):
    """Minimal JSON REST handler for HeGM control."""

    server_version = "HeGM-Ctrl/1.0"

    def _read_json(self) -> Dict:
        length = int(self.headers.get("Content-Length", "0") or "0")
        if not length:
            return {}
        raw = self.rfile.read(length)
        try:
            return json.loads(raw.decode("utf-8"))
        except Exception:
            return {}

    def _send_json(self, status: int, payload: Dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self) -> None:  # noqa: N802
        if self.path == "/checkpoint":
            self._handle_checkpoint()
        elif self.path == "/resume":
            self._handle_resume()
        else:
            self._send_json(404, {"error": "not-found"})

    def log_message(self, format: str, *args) -> None:  # noqa: A003
        # Route access logs through HeGM logger at DEBUG level.
        msg = format % args
        log(f"ctrl-http {self.address_string()}: {msg}", level="DEBUG")

    def _handle_checkpoint(self) -> None:
        payload = self._read_json()
        pids = payload.get("pids") or []
        all_flag = bool(payload.get("all"))

        if pids and all_flag:
            self._send_json(
                400, {"error": "specify either 'pids' or 'all', not both"}
            )
            return

        if all_flag:
            results = checkpoint_all_gpu()
        else:
            try:
                pid_ints = [int(p) for p in pids]
            except Exception:
                self._send_json(400, {"error": "invalid 'pids' payload"})
                return
            results = checkpoint_pids(pid_ints)

        self._send_json(200, {"results": results})

    def _handle_resume(self) -> None:
        payload = self._read_json()
        ppids = payload.get("ppids") or []
        all_flag = bool(payload.get("all"))

        if ppids and all_flag:
            self._send_json(
                400, {"error": "specify either 'ppids' or 'all', not both"}
            )
            return

        if all_flag:
            results = resume_all_pending()
        else:
            try:
                ppid_ints = [int(p) for p in ppids]
            except Exception:
                self._send_json(400, {"error": "invalid 'ppids' payload"})
                return
            results = resume_ppids(ppid_ints)

        self._send_json(200, {"results": results})


def ensure_api_server_started(
    host: str = "0.0.0.0", port: int = _DEFAULT_CTRL_PORT
) -> Tuple[Optional[ThreadingHTTPServer], Optional[threading.Thread]]:
    """Start the control API server in the background if not already running.

    This is idempotent and safe to call from the Launcher. If the port is
    already in use, a warning is logged and no server is started (assumed to
    be running elsewhere in the container).
    """
    global _api_server, _api_thread

    with _api_lock:
        if _api_server is not None:
            return _api_server, _api_thread

        try:
            server = ThreadingHTTPServer((host, port), _CtrlRequestHandler)
        except OSError as exc:
            warn(f"control API server not started (port {port}): {exc}")
            return None, None

        thread = threading.Thread(
            target=server.serve_forever, name="hegm-ctrl-http", daemon=True
        )
        thread.start()

        _api_server = server
        _api_thread = thread

        log(f"control API server listening on {host}:{port}", level="INFO")
        return _api_server, _api_thread


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="hegm-ctrl",
        description="HeGM external control CLI (checkpoint / resume).",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # checkpoint --------------------------------------------------------------
    p_checkpoint = subparsers.add_parser(
        "checkpoint", help="Trigger checkpoint via SIGUSR1."
    )
    p_checkpoint.add_argument(
        "-p",
        "--pid",
        dest="pids",
        type=int,
        nargs="+",
        help="Target worker PID(s) to signal.",
    )
    p_checkpoint.add_argument(
        "--all",
        action="store_true",
        help="Auto-detect all GPU-using PIDs and signal them.",
    )

    # resume ------------------------------------------------------------------
    p_resume = subparsers.add_parser(
        "resume", help="Resume from checkpoints by removing lock files."
    )
    p_resume.add_argument(
        "-P",
        "--ppid",
        dest="ppids",
        type=int,
        nargs="+",
        help="Parent PID(s) whose lock files should be removed.",
    )
    p_resume.add_argument(
        "--all",
        action="store_true",
        help="Resume all pending checkpoints under the base directory.",
    )

    # api-server --------------------------------------------------------------
    p_api = subparsers.add_parser(
        "api-server", help="Run the in-container REST control API server."
    )
    p_api.add_argument(
        "--host",
        default="0.0.0.0",
        help="Bind address for the REST server (default: 0.0.0.0).",
    )
    p_api.add_argument(
        "--port",
        type=int,
        default=_DEFAULT_CTRL_PORT,
        help=f"TCP port for the REST server (default: {_DEFAULT_CTRL_PORT}).",
    )

    return parser


def main_cli(argv: Optional[List[str]] = None) -> int:
    """Entry point for the ``hegm-ctrl`` console script."""
    parser = _build_parser()
    args = parser.parse_args(argv)

    if args.command == "checkpoint":
        if args.pids and args.all:
            parser.error("checkpoint: specify either --pid or --all, not both")
        if args.all:
            results = checkpoint_all_gpu()
            if not results:
                print(
                    "No GPU-using PIDs detected (either no active GPU jobs or "
                    "nvidia-smi is unavailable in this container)."
                )
                return 1
        elif args.pids:
            results = checkpoint_pids(args.pids)
        else:
            parser.error("checkpoint: one of --pid/--all is required")
        for pid, status in results.items():
            print(f"PID {pid}: {status}")
        return 0

    if args.command == "resume":
        if args.ppids and args.all:
            parser.error("resume: specify either --ppid or --all, not both")
        if args.all:
            results = resume_all_pending()
        elif args.ppids:
            results = resume_ppids(args.ppids)
        else:
            parser.error("resume: one of --ppid/--all is required")
        for ppid, status in results.items():
            print(f"PPID {ppid}: {status}")
        return 0

    if args.command == "api-server":
        server, _thread = ensure_api_server_started(args.host, args.port)
        if server is None:
            return 1
        try:
            # Block forever; serve_forever() is already running in a daemon
            # thread. This just keeps the process alive until interrupted.
            threading.Event().wait()
        except KeyboardInterrupt:
            return 0

    return 1


if __name__ == "__main__":  # pragma: no cover - manual invocation
    raise SystemExit(main_cli())

