from __future__ import annotations

import argparse
import json
import os
import subprocess
import time
from dataclasses import dataclass
from typing import Callable

from .discovery import HeadRecord, load_head_record


Runner = Callable[[list[str]], int]


def _default_runner(argv: list[str]) -> int:
    proc = subprocess.run(argv, check=False)
    return int(proc.returncode)


def _ray_stop_cmd() -> list[str]:
    return ["ray", "stop", "--force"]


def _ray_start_cmd(*, head_addr: str, node_ip: str | None, resources_json: str, block: bool) -> list[str]:
    argv = ["ray", "start", f"--address={head_addr}", f"--resources={resources_json}", "--disable-usage-stats"]
    if node_ip:
        argv.append(f"--node-ip-address={node_ip}")
    if block:
        argv.append("--block")
    return argv


@dataclass
class Watchdog:
    head_ip_file: str
    node_ip: str | None
    resources_json: str
    poll_s: int
    runner: Runner
    ray_start_block: bool

    _current_head_addr: str | None = None
    _ray_proc: subprocess.Popen[bytes] | None = None

    def _kill_ray_proc(self) -> None:
        p = self._ray_proc
        self._ray_proc = None
        if not p:
            return
        try:
            p.terminate()
        except Exception:
            pass
        try:
            p.wait(timeout=10)
        except Exception:
            try:
                p.kill()
            except Exception:
                pass
            try:
                p.wait(timeout=5)
            except Exception:
                pass

    def _restart_ray(self, desired: str) -> None:
        # Best-effort stop then start.
        self._kill_ray_proc()
        self.runner(_ray_stop_cmd())
        if self.ray_start_block:
            argv = _ray_start_cmd(
                head_addr=desired,
                node_ip=self.node_ip,
                resources_json=self.resources_json,
                block=True,
            )
            # Use Popen so watchdog can notice exits and reconnect if needed.
            self._ray_proc = subprocess.Popen(argv, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        else:
            self.runner(
                _ray_start_cmd(
                    head_addr=desired,
                    node_ip=self.node_ip,
                    resources_json=self.resources_json,
                    block=False,
                )
            )
        self._current_head_addr = desired

    def tick_once(self) -> HeadRecord | None:
        rec = load_head_record(self.head_ip_file)
        if rec is None:
            return None
        desired = rec.head_addr()
        if self._current_head_addr != desired:
            self._restart_ray(desired)
            return rec

        # If we're in block mode and the ray process died, restart it against the same head.
        if self.ray_start_block and self._ray_proc is not None:
            if self._ray_proc.poll() is not None:
                self._restart_ray(desired)
        return rec

    def run_forever(self) -> None:
        while True:
            try:
                self.tick_once()
            except Exception:
                # keep watchdog alive
                pass
            time.sleep(max(1, int(self.poll_s)))


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Stateless Ray worker watchdog (auto-connect + reconnect).")
    ap.add_argument("--head-ip-file", required=True)
    ap.add_argument("--node-ip", default=None)
    ap.add_argument(
        "--resources-kv",
        action="append",
        default=[],
        help='Repeatable k=v resources, e.g. --resources-kv worker_node=100',
    )
    ap.add_argument("--poll-s", type=int, default=5)
    ap.add_argument("--ray-start-block", action="store_true", help="Run `ray start ... --block` and self-heal on exit.")
    ap.add_argument("--once", action="store_true", help="Run one tick then exit (for testing/debug).")
    args = ap.parse_args(argv)

    resources: dict[str, float] = {"worker_node": 100.0}
    if args.resources_kv:
        resources = {}
        for item in args.resources_kv:
            if "=" not in item:
                raise SystemExit(f"invalid --resources-kv (expected k=v): {item!r}")
            k, v = item.split("=", 1)
            k = k.strip()
            v = v.strip()
            if not k:
                raise SystemExit(f"invalid --resources-kv key: {item!r}")
            resources[k] = float(v)

    node_ip = args.node_ip
    if node_ip in (None, ""):
        node_ip = os.environ.get("ARGUS_NODE_IP") or None

    wd = Watchdog(
        head_ip_file=args.head_ip_file,
        node_ip=node_ip,
        resources_json=json.dumps(resources),
        poll_s=int(args.poll_s),
        runner=_default_runner,
        ray_start_block=bool(args.ray_start_block),
    )
    if args.once:
        wd.tick_once()
        return 0

    wd.run_forever()
    return 0


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