"""Workers on other hosts, reached over a socket they opened themselves.

The engine runs where the automations are; a GPU sits somewhere else. Those
two are usually not on the same network, and the one that can be dialled is
the engine — so a worker connects *out* to it and the engine answers over the
connection it was given. That also means nothing has to expose Redis, which is
the thing a remote worker must never be handed.

What travels is the protocol the local worker pool already speaks: one JSON
object per line becomes one JSON frame, the node's source rides along with
every call so no code has to be distributed, and the reports a node makes
while it runs arrive the same way they do over a pipe. A node cannot tell
which kind of worker it is running on, which is the point — the same flow runs
in both places.

The awkward part is that the socket lives on the event loop and a node
executes on a worker thread. A call therefore hands its frame to the loop and
blocks on a queue of its own until the loop puts the answer there.
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
import queue
import sys
import threading
import time
from collections.abc import Callable
from typing import Any

from fluksio.flow.workers import NodeTimeout, RemoteError, _remote_class

logger = logging.getLogger(__name__)

#: How long the loop is given to accept a frame we are handing it.
SEND_TIMEOUT_S = 30.0
#: A worker that has said nothing for this long is treated as gone. It sends a
#: heartbeat while it is executing, so this only ever catches a dead socket.
SILENCE_S = 90.0
#: Protocol version this engine speaks. A worker announcing anything else is
#: refused rather than half-understood.
PROTOCOL = 1


class NoWorker(RemoteError):
    """Nothing is attached that carries the label this node asked for."""


class RemoteWorker:
    """One attached worker, and the calls it has in flight."""

    def __init__(
        self,
        name: str,
        labels: list[str],
        send: Callable[[dict[str, Any]], Any],
        loop: asyncio.AbstractEventLoop,
        max_parallel: int = 1,
        info: dict[str, Any] | None = None,
    ) -> None:
        self.name = name
        self.labels = set(labels)
        self.info = info or {}
        self.attached_at = time.time()
        self.last_seen = time.time()
        self._send = send
        self._loop = loop
        self._slots = threading.Semaphore(max_parallel)
        self.max_parallel = max_parallel
        self._pending: dict[str, queue.Queue[dict[str, Any] | None]] = {}
        self._lock = threading.Lock()
        self._gone = False
        # Source this worker has already loaded, by digest. A rebuild asks
        # every device-bound node whether it compiles, and without this that
        # is a network round trip each time — which is how attaching a GPU box
        # makes rebuilding the flows slow enough to fail a health check. A
        # worker that reconnects is a new object, so this empties with it.
        self._compiled: set[str] = set()

    # -------------------------------------------------------------------------
    # From the socket's side, on the event loop
    # -------------------------------------------------------------------------

    def deliver(self, message: dict[str, Any]) -> None:
        """Hand a frame to whichever call is waiting for it."""
        self.last_seen = time.time()
        call_id = str(message.get("call_id") or "")
        with self._lock:
            inbox = self._pending.get(call_id)
        if inbox is not None:
            inbox.put(message)

    def detach(self) -> None:
        """The socket closed: wake everything still waiting on it."""
        self._gone = True
        with self._lock:
            inboxes = list(self._pending.values())
        for inbox in inboxes:
            # None is "no more answers are coming", which the caller turns into
            # a failed node rather than a wait that never ends.
            inbox.put(None)

    # -------------------------------------------------------------------------
    # From a node's side, on a worker thread
    # -------------------------------------------------------------------------

    def request(
        self,
        payload: dict[str, Any],
        timeout: float,
        on_event: Callable[[dict[str, Any]], None] | None = None,
    ) -> dict[str, Any]:
        if self._gone:
            raise RemoteError(f"worker '{self.name}' is no longer attached")
        call_id = str(payload["call_id"])
        inbox: queue.Queue[dict[str, Any] | None] = queue.Queue()
        # Blocking here is the backpressure, exactly as taking a slot is in the
        # local pool.
        self._slots.acquire()
        with self._lock:
            self._pending[call_id] = inbox
        try:
            future = asyncio.run_coroutine_threadsafe(self._send(payload), self._loop)
            future.result(timeout=SEND_TIMEOUT_S)

            while True:
                try:
                    # Reset per frame: the deadline measures silence, so a node
                    # reporting its progress is never mistaken for a hung one.
                    message = inbox.get(timeout=timeout)
                except queue.Empty:
                    self.cancel(call_id)
                    raise NodeTimeout(
                        f"'{self.name}' was silent for {timeout}s"
                    ) from None
                if message is None:
                    raise RemoteError(f"worker '{self.name}' went away mid-call")
                kind = message.get("event")
                if kind == "heartbeat":
                    continue
                if kind:
                    if on_event is not None:
                        try:
                            on_event(message)
                        except Exception:
                            logger.exception("Could not record a worker event")
                    continue
                return message
        except Exception as exc:
            if isinstance(exc, (NodeTimeout, RemoteError)):
                raise
            raise RemoteError(f"worker '{self.name}': {exc}") from exc
        finally:
            with self._lock:
                self._pending.pop(call_id, None)
            self._slots.release()

    @property
    def compiled(self) -> set[str]:
        """Digests of the source this worker has already loaded."""
        return self._compiled

    @property
    def gone(self) -> bool:
        return self._gone

    @property
    def in_flight(self) -> int:
        """Calls this worker has not answered yet."""
        with self._lock:
            return len(self._pending)

    def calls_of(self, run_id: str) -> list[str]:
        with self._lock:
            return [call for call in self._pending if call.startswith(f"{run_id}:")]

    def cancel(self, call_id: str) -> None:
        """Ask the worker to kill what it is running for this call."""
        try:
            asyncio.run_coroutine_threadsafe(
                self._send({"op": "cancel", "call_id": call_id}), self._loop
            ).result(timeout=SEND_TIMEOUT_S)
        except Exception:
            logger.warning("Could not cancel '%s' on '%s'", call_id, self.name)


class RemoteWorkerHub:
    """Every attached worker, and which of them a node may run on."""

    def __init__(self) -> None:
        self._workers: dict[str, RemoteWorker] = {}
        self._lock = threading.Lock()

    # -------------------------------------------------------------------------
    # Attachment
    # -------------------------------------------------------------------------

    def attach(self, worker: RemoteWorker) -> None:
        with self._lock:
            existing = self._workers.get(worker.name)
            if existing is not None:
                # A worker that reconnects after a network drop: the old socket
                # is dead whether or not it has noticed yet.
                existing.detach()
            self._workers[worker.name] = worker
        logger.info(
            "Worker '%s' attached with labels %s", worker.name, sorted(worker.labels)
        )

    def detach(self, name: str) -> None:
        with self._lock:
            worker = self._workers.pop(name, None)
        if worker is not None:
            worker.detach()
            logger.info("Worker '%s' detached", name)

    def workers(self) -> list[RemoteWorker]:
        with self._lock:
            return list(self._workers.values())

    def labels(self) -> set[str]:
        """Every label something attached right now carries."""
        with self._lock:
            return {
                label for worker in self._workers.values() for label in worker.labels
            }

    def pick(self, label: str) -> RemoteWorker | None:
        """A worker carrying this label, least busy first.

        Resolved per call rather than when the flow was built, so a worker that
        attaches after a run was submitted picks the work up without anything
        being rebuilt.
        """
        with self._lock:
            candidates = [
                worker
                for worker in self._workers.values()
                if not worker.gone and (label == worker.name or label in worker.labels)
            ]
        if not candidates:
            return None
        return min(candidates, key=lambda worker: worker.in_flight)

    # -------------------------------------------------------------------------
    # Running a node on one
    # -------------------------------------------------------------------------

    def run(
        self,
        label: str,
        flow: str,
        node: str,
        source: str,
        kwargs: dict[str, Any],
        node_id: str,
        timeout: float,
        run_id: str = "",
        on_event: Callable[[dict[str, Any]], None] | None = None,
    ) -> Any:
        worker = self.pick(label)
        if worker is None:
            raise NoWorker(f"no worker labelled '{label}' is attached")
        response = worker.request(
            {
                "op": "run",
                "call_id": f"{run_id}:{node_id}" if run_id else node_id,
                "flow": flow,
                "node": node,
                "source": source,
                "kwargs": kwargs,
                "run": {"id": run_id} if run_id else None,
                "timeout": timeout,
            },
            timeout=timeout,
            on_event=on_event,
        )
        logs = response.get("logs")
        if logs:
            sys.stdout.write(logs)
        if response.get("ok"):
            return response.get("result")
        error = response.get("error") or {}
        raise _remote_class(str(error.get("type") or "RemoteError"))(
            str(error.get("message") or "the node failed"),
            str(error.get("traceback") or ""),
        )

    def compile(
        self, label: str, flow: str, node: str, source: str, timeout: float = 60.0
    ) -> str | None:
        """Load this source on the worker that will run it.

        Which machine compiles matters: a node importing torch is fine on the
        GPU box and a ``ModuleNotFoundError`` on the engine, so checking it
        here would fail a node that is perfectly correct. When nothing is
        attached there is nothing to check against, and ``None`` says so — the
        node is not broken, it is waiting.
        """
        worker = self.pick(label)
        if worker is None:
            return None
        digest = hashlib.md5(f"{flow}.{node}:{source}".encode()).hexdigest()
        if digest in worker.compiled:
            return None
        try:
            response = worker.request(
                {
                    "op": "compile",
                    "call_id": f"compile:{flow}.{node}",
                    "flow": flow,
                    "node": node,
                    "source": source,
                },
                timeout=timeout,
            )
        except RemoteError as exc:
            return f"{type(exc).__name__}: {exc}"
        if response.get("ok"):
            worker.compiled.add(digest)
            return None
        error = response.get("error") or {}
        return str(error.get("short") or "The node could not be loaded.")

    def proxy(
        self,
        label: str,
        flow: str,
        node: str,
        source: str,
        node_id: str,
        timeout: float,
        run_id: str = "",
        on_event: Callable[[dict[str, Any]], None] | None = None,
        fallback: Callable[..., Any] | None = None,
    ) -> Callable[..., Any]:
        """What a node with a device runs instead of its own function.

        ``fallback`` is the local pool's proxy, used when the node only prefers
        the label rather than requiring it.
        """

        def call(**kwargs: Any) -> Any:
            if fallback is not None and self.pick(label) is None:
                return fallback(**kwargs)
            return self.run(
                label,
                flow,
                node,
                source,
                kwargs,
                node_id,
                timeout,
                run_id=run_id,
                on_event=on_event,
            )

        return call

    def cancel_run(self, run_id: str) -> int:
        """Kill whatever this run has executing on any attached worker."""
        stopped = 0
        for worker in self.workers():
            for call in worker.calls_of(run_id):
                worker.cancel(call)
                stopped += 1
        return stopped