"""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 --python /opt/venv/bin/python

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
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 = 1
#: 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


#: 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] = {}

    async def serve_forever(self) -> None:
        backoff = 1.0
        while True:
            try:
                await self._session()
                backoff = 1.0
            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),
                    }
                )
            )
            welcome = json.loads(await socket.recv())
            if welcome.get("op") != "welcome":
                raise RuntimeError(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,
            )

            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 ""))

    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()
        worker = Subprocess(self.args.python, self.env)
        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:
            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)
            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 _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; point it at the venv with torch",
    )
    parser.add_argument("--parallel", type=int, default=1)
    parser.add_argument("--artifact-url", default="")
    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")
    return 0


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