"""The agent that runs Fluksio nodes on a machine the engine cannot reach.

Point it at the engine on the box with the GPU and it dials in::

    pip install fluksio-worker
    fluksio-worker --url wss://api.example.com/api/v1/workers/attach \\
        --token "$FLUKSIO_WORKER_TOKEN" --labels gpu

Install it into the environment the training code already runs in and node
code runs on that: ``--python`` defaults to the interpreter this was started
with. Point it elsewhere only when the two are meant to differ.

It connects *out*, so the engine needs no route back and nothing has to expose
Redis. What it then does is what the engine's own worker pool does: hold a few
subprocesses running ``worker_main.py``, hand each call to one, and pass back
everything that comes out — including the metrics a training loop reports
while it is still running.

Deliberately two files with one dependency. Nothing of the engine is imported
here; a worker host installs Python, ``websockets``, and whatever the nodes
themselves need. Where pip is not an option, copying this file and
``worker_main.py`` into one directory and running ``python agent.py`` is the
same thing — the engine serves the runner at ``GET /api/v1/workers/runtime``.
"""

from __future__ import annotations

import argparse
import asyncio
import contextlib
import json
import logging
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

try:
    import websockets
except ImportError:  # pragma: no cover - the one dependency, named plainly
    print(
        "This needs the 'websockets' package: pip install websockets", file=sys.stderr
    )
    raise SystemExit(1) from None

log = logging.getLogger("fluksio-worker")

PROTOCOL = 2
#: Sent while a call is running, so the engine can tell working from wedged.
HEARTBEAT_S = 10.0
#: Reconnection backs off to this and no further.
MAX_BACKOFF_S = 30.0


class Refused(RuntimeError):
    """The engine will not have this worker, and retrying will not change it."""


class Idle(Exception):
    """Nothing has run here for long enough that this worker is done."""


#: The user-code runner, beside this file — installed together, or copied
#: together onto a host where pip is not an option.
WORKER_MAIN = Path(__file__).resolve().with_name("worker_main.py")


class Subprocess:
    """One user-code process and the framing of one call over its pipes."""

    def __init__(self, python: str, env: dict[str, str]) -> None:
        self.proc = subprocess.Popen(
            [python, str(WORKER_MAIN)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            close_fds=True,
            env=env,
        )
        self._buffer = bytearray()

    def alive(self) -> bool:
        return self.proc.poll() is None

    def send(self, request: dict[str, Any]) -> None:
        assert self.proc.stdin is not None
        self.proc.stdin.write((json.dumps(request) + "\n").encode())
        self.proc.stdin.flush()

    def read_line(self) -> str:
        """One line, blocking. Empty when the process is gone.

        Read at the file-descriptor level so a report arriving mid-call is
        passed on the moment it is written rather than when a buffer fills.
        """
        assert self.proc.stdout is not None
        fd = self.proc.stdout.fileno()
        while True:
            end = self._buffer.find(b"\n")
            if end >= 0:
                line = bytes(self._buffer[: end + 1])
                del self._buffer[: end + 1]
                return line.decode(errors="replace")
            try:
                chunk = os.read(fd, 65536)
            except OSError:
                chunk = b""
            if not chunk:
                self._buffer.clear()
                return ""
            self._buffer += chunk

    def kill(self) -> None:
        with contextlib.suppress(OSError):
            self.proc.send_signal(signal.SIGKILL)
        with contextlib.suppress(subprocess.TimeoutExpired):
            self.proc.wait(timeout=5)


class Agent:
    """Holds the connection, and one subprocess per call in flight."""

    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.env = dict(os.environ)
        self.env["FLUKSIO_ARTIFACT_URL"] = args.artifact_url or _artifacts_from(
            args.url
        )
        self.env["FLUKSIO_ARTIFACT_TOKEN"] = args.token
        self.running: dict[str, Subprocess] = {}
        self.last_done = time.monotonic()
        self._idled = False

    async def serve_forever(self) -> None:
        backoff = 1.0
        while True:
            try:
                await self._session()
                backoff = 1.0
            except Idle:
                log.info("nothing to do for %.0fs — stopping", self.args.max_idle)
                return
            except Refused:
                raise
            except Exception as exc:
                log.warning("disconnected: %s — retrying in %.0fs", exc, backoff)
                await asyncio.sleep(backoff)
                backoff = min(MAX_BACKOFF_S, backoff * 2)

    async def _session(self) -> None:
        url = f"{self.args.url}?token={self.args.token}"
        async with websockets.connect(url, max_size=None, ping_interval=20) as socket:
            await socket.send(
                json.dumps(
                    {
                        "op": "hello",
                        "protocol": PROTOCOL,
                        "name": self.args.name,
                        "labels": self.args.labels,
                        "python": self.args.python,
                        "max_parallel": self.args.parallel,
                        "venv_digest": _venv_digest(self.args.python),
                        "inventory": _inventory(self.args),
                    }
                )
            )
            welcome = json.loads(await socket.recv())
            if welcome.get("op") != "welcome":
                # Nothing about this machine will change the answer, so
                # reconnecting would only be a quiet loop in the job's log.
                raise Refused(str(welcome.get("reason") or "refused"))
            log.info(
                "attached to %s as '%s' with labels %s",
                self.args.url,
                welcome.get("name"),
                self.args.labels,
            )

            self.last_done = time.monotonic()
            watchdog = (
                asyncio.create_task(self._idle_watch(socket))
                if self.args.max_idle > 0
                else None
            )
            try:
                async for raw in socket:
                    message = json.loads(raw)
                    op = message.get("op")
                    if op in ("run", "compile"):
                        # Compiling is loading the source, which is the same
                        # trip through a subprocess a call is — and has to
                        # happen here rather than on the engine, because "does
                        # this import" is a question about *this* machine's
                        # packages.
                        log.info("%s %s", op, message.get("call_id"))
                        task = asyncio.create_task(self._run(socket, message))
                        # Without this a failure in the task is only noticed
                        # when it is garbage collected, which reads as a call
                        # that vanished.
                        task.add_done_callback(_report_failure)
                    elif op == "cancel":
                        self._cancel(str(message.get("call_id") or ""))
            except Exception:
                # The watchdog closing the socket is a stop, not a disconnect.
                if not self._idled:
                    raise
            finally:
                if watchdog is not None:
                    watchdog.cancel()
        if self._idled:
            raise Idle

    async def _idle_watch(self, socket: Any) -> None:
        """Close the socket once nothing has run here for ``--max-idle``.

        A worker a batch scheduler started for one node has no other way to
        know it is finished, and a job that idles until its walltime is a job
        somebody else was queued behind.
        """
        tick = min(30.0, max(1.0, self.args.max_idle / 4))
        while True:
            await asyncio.sleep(tick)
            if self.running:
                continue
            if time.monotonic() - self.last_done >= self.args.max_idle:
                self._idled = True
                with contextlib.suppress(Exception):
                    await socket.close()
                return

    async def _run(self, socket: Any, request: dict[str, Any]) -> None:
        """Execute one call in a subprocess, streaming what it says back."""
        call_id = str(request.get("call_id") or "")
        loop = asyncio.get_running_loop()
        # What the engine allocated this call: thread caps and the GPUs it may
        # see. Taken off the request because a library reads those once, when
        # it is imported — which for this agent is a process that does not
        # exist yet, so there is nothing to reconfigure and no pool to key.
        extra = {str(k): str(v) for k, v in (request.pop("env", None) or {}).items()}
        mask = self.env.get("CUDA_VISIBLE_DEVICES")
        if mask and extra.get("CUDA_VISIBLE_DEVICES"):
            extra["CUDA_VISIBLE_DEVICES"] = _remap_cuda(
                mask, extra["CUDA_VISIBLE_DEVICES"]
            )
        worker = Subprocess(self.args.python, {**self.env, **extra})
        self.running[call_id] = worker

        async def beat() -> None:
            while True:
                await asyncio.sleep(HEARTBEAT_S)
                with contextlib.suppress(Exception):
                    await socket.send(
                        json.dumps({"call_id": call_id, "event": "heartbeat"})
                    )

        heartbeat = asyncio.create_task(beat())
        try:
            # A subprocess that died before it could be written to is reported
            # by the read below, which says so and ends the call — rather than
            # raising here and leaving the engine waiting out its silence.
            with contextlib.suppress(OSError):
                await loop.run_in_executor(None, worker.send, request)
            while True:
                line = await loop.run_in_executor(None, worker.read_line)
                if not line:
                    await socket.send(
                        json.dumps(
                            {
                                "call_id": call_id,
                                "ok": False,
                                "error": {
                                    "type": "NodeCancelled"
                                    if worker.proc.returncode
                                    else "RemoteError",
                                    "message": "the node process stopped",
                                    "short": "the node process stopped",
                                    "traceback": "",
                                },
                            }
                        )
                    )
                    return
                await socket.send(line.strip())
                # Anything without an `event` is the answer; the call is over.
                if not json.loads(line).get("event"):
                    return
        finally:
            heartbeat.cancel()
            self.running.pop(call_id, None)
            self.last_done = time.monotonic()
            worker.kill()

    def _cancel(self, call_id: str) -> None:
        worker = self.running.get(call_id)
        if worker is not None:
            log.info("cancelling %s", call_id)
            worker.kill()


def _report_failure(task: asyncio.Task[Any]) -> None:
    if not task.cancelled() and task.exception() is not None:
        log.exception("call failed", exc_info=task.exception())


def _artifacts_from(url: str) -> str:
    """The artifact endpoint beside the socket, so one URL configures both."""
    base = url.replace("wss://", "https://").replace("ws://", "http://")
    return base.rsplit("/workers/attach", 1)[0] + "/artifacts"


def _detect_cpus() -> int:
    """Cores this worker may use — what the batch job was given, or the box."""
    given = os.environ.get("SLURM_CPUS_ON_NODE", "")
    return int(given) if given.isdigit() else (os.cpu_count() or 1)


def _detect_ram_mb() -> int | None:
    """Memory in MB, or None where it cannot be asked — which is allowed."""
    given = os.environ.get("SLURM_MEM_PER_NODE", "")
    if given.isdigit():
        return int(given)
    try:
        return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
    except (AttributeError, ValueError, OSError):
        return None


def _detect_gpus() -> int:
    """What this worker was *given*, never what the machine has.

    Nothing is probed: asking a vendor tool would make the one dependency two,
    and the engine does not probe its own GPUs either. A batch scheduler says
    so in the environment; anywhere else it is ``--gpus``.
    """
    for name in ("SLURM_GPUS_ON_NODE", "FLUKSIO_WORKER_GPUS"):
        given = os.environ.get(name, "")
        if given.isdigit():
            return int(given)
    listed = os.environ.get("SLURM_JOB_GPUS", "")
    return len([part for part in listed.split(",") if part.strip()])


def _inventory(args: argparse.Namespace) -> dict[str, Any]:
    """What this machine holds, for the engine to schedule against."""
    return {
        "cpus": args.cpus or _detect_cpus(),
        "gpus": _detect_gpus() if args.gpus is None else args.gpus,
        "ram_mb": args.ram_mb or _detect_ram_mb(),
    }


def _remap_cuda(mask: str, requested: str) -> str:
    """Turn engine-assigned device numbers into the ones this job may see.

    The engine counts GPUs from the inventory reported here, so it asks for
    device 0. Under a scheduler's gres this agent already runs with something
    like ``CUDA_VISIBLE_DEVICES=2,3``, where 0 is a device the job was never
    given — so each index is read as a position in that mask.
    """
    devices = [part.strip() for part in mask.split(",") if part.strip()]
    out = []
    for part in requested.split(","):
        part = part.strip()
        if not part:
            continue
        try:
            out.append(devices[int(part)])
        except (ValueError, IndexError):
            out.append(part)
    return ",".join(out)


def _venv_digest(python: str) -> str:
    """What is installed here, so the engine can say when it has drifted."""
    try:
        listing = subprocess.run(
            [python, "-m", "pip", "freeze"],
            capture_output=True,
            text=True,
            timeout=60,
            check=False,
        ).stdout
    except (OSError, subprocess.SubprocessError):
        return ""
    import hashlib

    return hashlib.sha256(listing.encode()).hexdigest()[:16]


def main(argv: list[str] | None = None) -> int:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser(description="Run Fluksio nodes on this machine.")
    parser.add_argument("--url", required=True, help="wss://…/api/v1/workers/attach")
    parser.add_argument(
        "--token",
        default=os.environ.get("FLUKSIO_WORKER_TOKEN", ""),
        help="issued by POST /api/v1/workers/tokens",
    )
    parser.add_argument("--name", default=os.uname().nodename)
    parser.add_argument(
        "--labels",
        default="",
        help="comma-separated, e.g. gpu,cuda12 — what a node's device matches",
    )
    parser.add_argument(
        "--python",
        default=sys.executable,
        help=(
            "the interpreter node code runs on (default: the one running this, "
            "so installing into the venv with torch in it is enough)"
        ),
    )
    parser.add_argument("--parallel", type=int, default=1)
    parser.add_argument("--artifact-url", default="")
    parser.add_argument(
        "--cpus",
        type=int,
        default=0,
        help="cores to advertise (default: what the job or the machine has)",
    )
    parser.add_argument(
        "--gpus",
        type=int,
        default=None,
        help="GPUs to advertise — nothing is probed, so say so here or in the job",
    )
    parser.add_argument(
        "--ram-mb",
        type=int,
        default=0,
        help="memory in MB to advertise (default: what the job or the machine has)",
    )
    parser.add_argument(
        "--max-idle",
        type=float,
        default=0.0,
        help=(
            "stop after this many seconds with nothing running — for a worker "
            "a batch scheduler started for one job (default: never)"
        ),
    )
    args = parser.parse_args(argv)

    if not args.token:
        parser.error("a token is required (--token or FLUKSIO_WORKER_TOKEN)")
    args.labels = [part.strip() for part in args.labels.split(",") if part.strip()]
    if not WORKER_MAIN.exists():
        parser.error(f"{WORKER_MAIN} is missing — copy it beside this file")

    agent = Agent(args)
    try:
        asyncio.run(agent.serve_forever())
    except KeyboardInterrupt:
        log.info("stopping")
    except Refused as exc:
        log.error("the engine refused this worker: %s", exc)
        return 1
    return 0


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