import argparse
import json
import os
import queue
import signal
import subprocess
import sys
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Dict, List, Optional, Sequence, Tuple


_DEFAULT_SHUTDOWN_TIMEOUT = 10.0
_POLL_INTERVAL = 0.1
_WORKER_ENTRYPOINT = Path(__file__).with_name("supervisor_worker.py")


@dataclass(frozen=True)
class RestartPolicy:
    mode: str
    max_restarts: Optional[int] = None

    @classmethod
    def parse(cls, raw: str) -> "RestartPolicy":
        value = raw.strip().lower()
        if value == "never":
            return cls(mode="never", max_restarts=0)
        if value == "always":
            return cls(mode="always", max_restarts=None)

        try:
            count = int(value)
        except ValueError as exc:
            raise argparse.ArgumentTypeError(
                "restart must be 'never', 'always', or a non-negative integer"
            ) from exc

        if count < 0:
            raise argparse.ArgumentTypeError(
                "restart count must be a non-negative integer"
            )
        return cls(mode="count", max_restarts=count)

    def allows_restart(self, restarts_so_far: int) -> bool:
        if self.mode == "always":
            return True
        if self.mode == "never":
            return False
        return restarts_so_far < (self.max_restarts or 0)

    def describe(self) -> str:
        if self.mode == "always":
            return "always"
        if self.mode == "never":
            return "never"
        return str(self.max_restarts)


@dataclass
class WorkerProcess:
    index: int
    process: subprocess.Popen[str]
    pgid: Optional[int]
    actual_pid: Optional[int]
    guard: str
    launch_id: int
    restart_count: int = 0
    stream_threads: List[threading.Thread] = field(default_factory=list)

    @property
    def name(self) -> str:
        return f"worker-{self.index}"


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="sdrun",
        description="Run and supervise multiple SandAI operator worker processes.",
    )
    parser.add_argument(
        "-w",
        "--worker",
        type=int,
        default=1,
        help="number of worker processes to run (default: 1)",
    )
    parser.add_argument(
        "--restart",
        type=RestartPolicy.parse,
        default=RestartPolicy.parse("never"),
        help="worker restart policy for non-zero exits: never, always, or a non-negative integer",
    )
    parser.add_argument(
        "--success-exit",
        choices=("ignore", "shutdown"),
        default="ignore",
        help="behavior when a worker exits with code 0 (default: ignore)",
    )
    parser.add_argument(
        "--failure-exit",
        choices=("ignore", "shutdown"),
        default="ignore",
        help="behavior when a worker exits non-zero and will not be restarted (default: ignore)",
    )
    parser.add_argument(
        "--startup-stagger",
        type=float,
        default=0.0,
        help="delay in seconds between sequential worker starts (default: 0)",
    )
    parser.add_argument(
        "command",
        nargs=argparse.REMAINDER,
        help="worker command after '--', for example: -- python main.py",
    )
    return parser


class Supervisor:
    def __init__(
        self,
        *,
        worker_count: int,
        restart_policy: RestartPolicy,
        command: Sequence[str],
        success_exit_policy: str = "ignore",
        failure_exit_policy: str = "ignore",
        startup_stagger: float = 0.0,
        stdout: Optional[IO[str]] = None,
        stderr: Optional[IO[str]] = None,
        shutdown_timeout: float = _DEFAULT_SHUTDOWN_TIMEOUT,
        install_signal_handlers: bool = True,
    ) -> None:
        if worker_count < 1:
            raise ValueError("worker_count must be >= 1")
        if not command:
            raise ValueError("command must not be empty")

        self.worker_count = worker_count
        self.restart_policy = restart_policy
        self.command = list(command)
        self.success_exit_policy = success_exit_policy
        self.failure_exit_policy = failure_exit_policy
        self.startup_stagger = startup_stagger
        self.stdout = stdout or sys.stdout
        self.stderr = stderr or sys.stderr
        self.shutdown_timeout = shutdown_timeout
        self.install_signal_handlers = install_signal_handlers

        self._workers: Dict[int, WorkerProcess] = {}
        self._launch_ids: Dict[int, int] = {}
        self._log_lock = threading.Lock()
        self._stopping = False
        self._shutdown_deadline: Optional[float] = None
        self._exit_code = 0
        self._completed_exit_codes: List[int] = []
        self._forwarded_signal: Optional[int] = None

    def run(self) -> int:
        if self.install_signal_handlers:
            self._install_signal_handlers()

        for index in range(self.worker_count):
            self._sleep_before_start(index, is_restart=False)
            self._workers[index] = self._start_worker(index, restart_count=0)

        while self._workers:
            for index in list(self._workers):
                worker = self._workers.get(index)
                if worker is None:
                    continue

                return_code = worker.process.poll()
                if return_code is None:
                    continue

                self._handle_worker_exit(worker, return_code)

            if self._stopping:
                self._maybe_escalate_shutdown()

            time.sleep(_POLL_INTERVAL)

        if self._stopping or self._forwarded_signal is not None:
            return self._exit_code
        return sum(self._completed_exit_codes)

    def _install_signal_handlers(self) -> None:
        for signame in ("SIGINT", "SIGTERM", "SIGHUP", "SIGQUIT"):
            signum = getattr(signal, signame, None)
            if signum is not None:
                signal.signal(signum, self._handle_signal)

    def _handle_signal(self, signum: int, _frame: object) -> None:
        if self._stopping:
            self._log_event(
                f"received second signal {signal.Signals(signum).name}; escalating to SIGKILL"
            )
            self._kill_all(signal.SIGKILL)
            return

        self._forwarded_signal = signum
        self._exit_code = 128 + signum
        self._log_event(
            f"received signal {signal.Signals(signum).name}; forwarding to workers"
        )
        self._begin_shutdown(signum)

    def _begin_shutdown(self, signum: int) -> None:
        if self._stopping:
            self._kill_all(signum)
            return

        self._stopping = True
        self._shutdown_deadline = time.monotonic() + self.shutdown_timeout
        self._kill_all(signum)

    def _maybe_escalate_shutdown(self) -> None:
        if self._shutdown_deadline is None:
            return
        if time.monotonic() < self._shutdown_deadline:
            return

        self._shutdown_deadline = None
        self._log_event("workers did not stop before timeout; sending SIGKILL")
        self._kill_all(signal.SIGKILL)

    def _handle_worker_exit(self, worker: WorkerProcess, return_code: int) -> None:
        self._join_stream_threads(worker)
        self._workers.pop(worker.index, None)

        if self._stopping:
            self._log_event(f"{worker.name} exited with code {return_code} during shutdown")
            return

        if return_code != 0 and self.restart_policy.allows_restart(worker.restart_count):
            next_restart_count = worker.restart_count + 1
            self._log_event(
                f"{worker.name} exited with code {return_code}; restarting "
                f"({next_restart_count}/{self.restart_policy.describe()})"
            )
            self._sleep_before_start(worker.index, is_restart=True)
            self._workers[worker.index] = self._start_worker(
                worker.index,
                restart_count=next_restart_count,
            )
            return

        self._completed_exit_codes.append(return_code)
        if return_code == 0:
            if self.success_exit_policy == "shutdown" and self._workers:
                self._exit_code = 0
                self._log_event(
                    f"{worker.name} exited cleanly; success policy is shutdown, stopping remaining workers"
                )
                self._begin_shutdown(signal.SIGTERM)
                return

            self._log_event(
                f"{worker.name} exited cleanly; success policy is {self.success_exit_policy}"
            )
            return

        if self.failure_exit_policy == "shutdown" and self._workers:
            self._exit_code = return_code
            self._log_event(
                f"{worker.name} exited with code {return_code}; failure policy is shutdown, stopping remaining workers"
            )
            self._begin_shutdown(signal.SIGTERM)
            return

        self._log_event(
            f"{worker.name} exited with code {return_code}; failure policy is {self.failure_exit_policy}, waiting for remaining workers"
        )

    def _start_worker(self, index: int, *, restart_count: int) -> WorkerProcess:
        self._launch_ids[index] = self._launch_ids.get(index, 0) + 1
        launch_id = self._launch_ids[index]

        ready_r, ready_w = os.pipe()
        env = os.environ.copy()
        env["SANDAI_SUPERVISOR_READY_FD"] = str(ready_w)
        env["SDRUN_MODE"] = "true"
        env["SDRUN_WORLD_SIZE"] = str(self.worker_count)
        env["SDRUN_RANK"] = str(index)
        env["SDRUN_LOCAL_RANK"] = str(index)

        process = subprocess.Popen(
            [
                sys.executable,
                str(_WORKER_ENTRYPOINT),
                "--",
                *self.command,
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            stdin=subprocess.DEVNULL,
            text=True,
            bufsize=1,
            env=env,
            pass_fds=(ready_w,),
            close_fds=True,
        )
        os.close(ready_w)

        metadata = self._read_ready_payload(ready_r)

        worker = WorkerProcess(
            index=index,
            process=process,
            pgid=self._as_int(metadata.get("pgid")),
            actual_pid=self._as_int(metadata.get("pid")),
            guard=str(metadata.get("guard", "unknown")) if metadata else "unknown",
            launch_id=launch_id,
            restart_count=restart_count,
        )
        worker.stream_threads.extend(self._start_log_threads(worker))

        self._log_event(
            f"started {worker.name} launch={launch_id} pid={worker.actual_pid or process.pid} "
            f"pgid={worker.pgid or 'n/a'} guard={worker.guard}"
        )
        return worker

    def _sleep_before_start(self, index: int, *, is_restart: bool) -> None:
        if is_restart:
            return
        delay = index * self.startup_stagger
        if delay <= 0:
            return
        self._log_event(f"delaying start of worker-{index} by {delay:.3f}s")
        time.sleep(delay)

    def _read_ready_payload(self, fd: int) -> Dict[str, object]:
        with os.fdopen(fd, "r", encoding="utf-8", errors="replace") as stream:
            line = stream.readline().strip()
        if not line:
            return {}
        try:
            payload = json.loads(line)
        except json.JSONDecodeError:
            return {}
        if not isinstance(payload, dict):
            return {}
        return payload

    def _start_log_threads(self, worker: WorkerProcess) -> List[threading.Thread]:
        threads: List[threading.Thread] = []
        for stream_name, stream, sink in (
            ("stdout", worker.process.stdout, self.stdout),
            ("stderr", worker.process.stderr, self.stderr),
        ):
            if stream is None:
                continue
            thread = threading.Thread(
                target=self._pump_stream,
                args=(worker.name, worker.launch_id, stream_name, stream, sink),
                daemon=True,
            )
            thread.start()
            threads.append(thread)
        return threads

    def _pump_stream(
        self,
        worker_name: str,
        launch_id: int,
        stream_name: str,
        stream: IO[str],
        sink: IO[str],
    ) -> None:
        try:
            for raw_line in iter(stream.readline, ""):
                self._write_line(
                    sink,
                    f"[{worker_name}#{launch_id}][{stream_name}] {raw_line.rstrip()}\n",
                )
        finally:
            stream.close()

    def _join_stream_threads(self, worker: WorkerProcess) -> None:
        for thread in worker.stream_threads:
            thread.join(timeout=1.0)

    def _kill_all(self, signum: int) -> None:
        for worker in list(self._workers.values()):
            self._signal_worker(worker, signum)

    def _signal_worker(self, worker: WorkerProcess, signum: int) -> None:
        if worker.process.poll() is not None:
            return

        try:
            if os.name == "posix" and worker.pgid:
                os.killpg(worker.pgid, signum)
            else:
                worker.process.send_signal(signum)
        except ProcessLookupError:
            return
        except OSError as exc:
            self._log_event(f"failed to signal {worker.name}: {exc}")

    def _log_event(self, message: str) -> None:
        self._write_line(self.stderr, f"[sdrun] {message}\n")

    def _write_line(self, sink: IO[str], line: str) -> None:
        with self._log_lock:
            sink.write(line)
            sink.flush()

    @staticmethod
    def _as_int(value: object) -> Optional[int]:
        if value is None:
            return None
        try:
            return int(value)
        except (TypeError, ValueError):
            return None


def run_supervisor_in_thread(
    supervisor: Supervisor,
) -> Tuple[threading.Thread, "queue.Queue[object]"]:
    result_queue: "queue.Queue[object]" = queue.Queue(maxsize=1)

    def _target() -> None:
        try:
            result_queue.put(supervisor.run())
        except BaseException as exc:  # pragma: no cover
            result_queue.put(exc)

    thread = threading.Thread(target=_target, daemon=True)
    thread.start()
    return thread, result_queue


def _normalize_command(command: Sequence[str]) -> List[str]:
    items = list(command)
    if items and items[0] == "--":
        items = items[1:]
    return items


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)

    command = _normalize_command(args.command)
    if args.worker < 1:
        parser.error("worker count must be >= 1")
    if not command:
        parser.error("worker command is required after '--'")

    supervisor = Supervisor(
        worker_count=args.worker,
        restart_policy=args.restart,
        command=command,
        success_exit_policy=args.success_exit,
        failure_exit_policy=args.failure_exit,
        startup_stagger=args.startup_stagger,
    )
    return supervisor.run()


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