"""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.nodes.base import NodeOutputError
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 at all for this long — not even a heartbeat —
#: is treated as gone. It beats every ten seconds while it is executing, so this
#: catches a dead socket rather than a slow node, and it is what bounds a call
#: whose node has no timeout of its own.
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)

            # The node's own deadline measures silence, so a node reporting its
            # progress is never mistaken for a hung one. A heartbeat is not
            # progress: it says the agent is alive, which is what SILENCE_S
            # asks, and says nothing about the node — so it feeds the liveness
            # bound below and never the node's own.
            deadline = time.monotonic() + timeout if timeout > 0 else None
            while True:
                wait = SILENCE_S
                if deadline is not None:
                    wait = min(SILENCE_S, deadline - time.monotonic())
                try:
                    message = inbox.get(timeout=max(wait, 0.0))
                except queue.Empty:
                    self.cancel(call_id)
                    if deadline is not None and time.monotonic() >= deadline:
                        raise NodeTimeout(
                            f"'{self.name}' was silent for {timeout}s"
                        ) from None
                    raise RemoteError(
                        f"worker '{self.name}' sent nothing for "
                        f"{SILENCE_S:.0f}s and is presumed gone"
                    ) 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 NodeOutputError:
                            # A port the node never declared. Stop the call
                            # rather than let the rest of its emissions arrive.
                            self.cancel(call_id)
                            raise
                        except Exception:
                            logger.exception("Could not record a worker event")
                    if deadline is not None:
                        deadline = time.monotonic() + timeout
                    continue
                return message
        except Exception as exc:
            if isinstance(exc, (NodeTimeout, RemoteError, NodeOutputError)):
                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