"""WebSocket nodes: a socket the engine dials or serves, listening or sending."""

from __future__ import annotations

import asyncio
import contextlib
import hmac
import json
import logging
from collections.abc import Awaitable, Callable, Iterable
from enum import Enum
from typing import TYPE_CHECKING, Any
from urllib.parse import urlsplit

from pydantic import BaseModel, ConfigDict, Field

from fluksio.flow.messages import MessageSpec
from fluksio.flow.nodes.base import Node
from fluksio.flow.nodes.http import HttpNode, _first_catch_all

if TYPE_CHECKING:
    from fastapi import FastAPI
    from starlette.websockets import WebSocket

logger = logging.getLogger(__name__)


class WebSocketNode(Node):
    """
    A long-lived WebSocket, as a client or a server.

    The URL picks the role: ``ws://`` or ``wss://`` dials out and holds the
    connection, a path such as ``/ticker`` is served at
    ``/hooks/<flow id>/<path>[/<secret>]`` while the flow runs. The ports pick the
    direction, as on the mqtt node: outputs only make a **listener**, inputs a
    **sender**.

    - client listener: frames the server sends become outputs
    - client sender: inputs go to the server, one frame per message
    - server listener: frames any connected client sends become outputs
    - server sender: inputs go to every connected client

    Frames are text. An incoming frame is parsed as JSON; an object feeds the
    ports whose names (or ``json_key``) it carries, and a node with one output
    takes anything else whole. Outgoing, a single input is sent as it stands —
    a string raw, anything else as JSON — and several go as one object keyed
    by port.
    """

    class Mode(Enum):
        """Which way the frames go."""

        LISTENER = "listener"
        SENDER = "sender"

    # A frame sent twice is two commands to whatever reads it.
    idempotent = False

    __slots__ = (
        "url",
        "mode",
        "serves",
        "headers",
        "on_connect",
        "secret",
        "timeout",
        "send_queue_size",
        "json_keys",
        "_clients",
        "_route_registered",
        "_task",
        "_queue",
        "_loop",
    )

    class Params(BaseModel):
        model_config = ConfigDict(extra="allow")

        url: str = Field(
            description=(
                "ws:// or wss:// to connect to a server, or a path such as "
                "/ticker to serve one at /hooks/<flow id>/<path>."
            )
        )
        headers: dict[str, str] = Field(
            default_factory=dict,
            description=(
                "Handshake headers, client only. A value may be a secret reference."
            ),
        )
        on_connect: list[Any] = Field(
            default_factory=list,
            description=(
                "Frames sent after every (re)connect, client only: the "
                "subscribe or auth message many APIs expect. A string goes "
                "raw, anything else as JSON; an item may be a secret reference."
            ),
        )
        secret: str = Field(
            default="",
            description=(
                "Shared secret for a served socket, appended to its URL: "
                "/hooks/<flow id>/<url>/<secret>. Empty leaves it open to anyone."
            ),
            json_schema_extra={"x-secret": True},
        )
        timeout: float = Field(
            default=10.0,
            gt=0,
            description=(
                "Seconds allowed for the handshake, and for a send to one "
                "client of a served socket."
            ),
        )
        send_queue_size: int = Field(
            default=256,
            gt=0,
            description=(
                "How many messages may wait to be sent. Past this the oldest "
                "is dropped and the node reports degraded."
            ),
        )
        # Same as the mqtt node: a device that wraps its reading in an object.
        json_key: str | dict[str, str] = Field(
            default="",
            description=(
                "Key to lift out of a JSON object frame, for every port or as "
                '{"port": "key"}.'
            ),
        )

    @classmethod
    def target_key(cls, params: dict[str, Any]) -> str | None:
        # A URL, credentials stripped, exactly as for the http node.
        return HttpNode.target_key(params)

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        cfg = self.Params.model_validate(params or {})

        requires = Node._normalize_ports(requires)
        provides = Node._normalize_ports(provides)
        if not requires and not provides:
            raise ValueError(
                "A WebSocket node needs either inputs (to send) or outputs (to listen)"
            )
        # ponytail: one direction per node, as mqtt and http. Inputs and
        # outputs over one socket is not supported yet.
        self.mode = (
            WebSocketNode.Mode.SENDER if requires else WebSocketNode.Mode.LISTENER
        )

        self.serves = cfg.url.startswith("/")
        if not self.serves and urlsplit(cfg.url).scheme not in ("ws", "wss"):
            raise ValueError(
                "A WebSocket url is ws://…, wss://… or a path to serve, "
                f"not '{cfg.url}'"
            )

        ports = requires or provides
        if isinstance(cfg.json_key, dict):
            self.json_keys: dict[str, str] = dict(cfg.json_key)
        else:
            self.json_keys = (
                {spec.port: cfg.json_key for spec in ports} if cfg.json_key else {}
            )

        self.url = cfg.url
        self.headers = cfg.headers
        self.on_connect = [_frame(item) for item in cfg.on_connect]
        self.secret = cfg.secret
        self.timeout = cfg.timeout
        self.send_queue_size = cfg.send_queue_size

        self._clients: set[WebSocket] = set()
        self._route_registered = False
        self._task: asyncio.Task[None] | None = None
        self._queue: asyncio.Queue[dict[str, Any]] | None = None
        self._loop: asyncio.AbstractEventLoop | None = None

        if name is None:
            safe_url = (
                self.url.split("://", 1)[-1].replace("/", "_").replace(":", "")
            ).strip("_")
            name = f"websocket_{self.mode.value}_{safe_url}"

        super().__init__(
            f=(
                self._sender_handler
                if self.mode is WebSocketNode.Mode.SENDER
                else self._noop_listener
            ),
            requires=requires,
            provides=provides,
            params=params,
            name=name,
        )

    @staticmethod
    def _noop_listener(params: dict[str, Any], **kwargs: Any) -> dict[str, Any] | None:
        """A listener injects what arrives; it is never called with inputs."""
        return None

    # -- frames ---------------------------------------------------------------

    def _encode(self, data: dict[str, Any]) -> str:
        """One message as one frame."""
        if len(data) == 1:
            return _frame(next(iter(data.values())))
        return json.dumps(data, default=str)

    def _decode(self, text: str) -> dict[str, Any]:
        """One frame as outputs, keyed by port. Empty when nothing matched."""
        try:
            parsed = json.loads(text)
        except json.JSONDecodeError:
            parsed = text
        outputs: dict[str, Any] = {}
        if isinstance(parsed, dict):
            for spec in self.output_ports:
                key = self.json_keys.get(spec.port, spec.port)
                if key in parsed:
                    outputs[spec.port] = spec.coerce(parsed[key])
        if not outputs and len(self.output_ports) == 1:
            spec = self.output_ports[0]
            outputs[spec.port] = spec.coerce(parsed)
        return outputs

    async def _receive(self, text: str) -> None:
        """Inject one frame. A bad one is logged, never fatal to the socket."""
        try:
            outputs = self._decode(text)
            if outputs:
                await asyncio.to_thread(self.inject, outputs)
            else:
                logger.debug("[%s] Frame matched no port: %s", self.name, text)
        except Exception as exc:
            logger.error("[%s] Error processing frame: %s", self.name, exc)

    # -- sending --------------------------------------------------------------

    def _sender_handler(
        self, params: dict[str, Any], **kwargs: Any
    ) -> dict[str, Any] | None:
        """Hand the inputs to the node's socket task, which owns the send."""
        loop, queue = self._loop, self._queue
        if loop is not None and queue is not None:
            loop.call_soon_threadsafe(self._enqueue, queue, dict(kwargs))
            return None

        # Not started: a draft preview or a test. A client connects once and
        # sends; a served socket has nobody to send to.
        if not self.serves:
            from websockets.sync.client import connect

            with connect(
                self.url,
                additional_headers=self.headers,
                open_timeout=self.timeout,
            ) as ws:
                for frame in self.on_connect:
                    ws.send(frame)
                ws.send(self._encode(dict(kwargs)))
        return None

    def _enqueue(
        self, queue: asyncio.Queue[dict[str, Any]], data: dict[str, Any]
    ) -> None:
        """Queue a message, dropping the oldest when the socket cannot keep up."""
        if queue.full():
            with contextlib.suppress(asyncio.QueueEmpty):
                queue.get_nowait()
                logger.warning("Send queue full for node '%s', dropped", self.name)
                self.report_health("degraded", "send queue full")
        queue.put_nowait(data)

    async def _send_queued(self, send: Callable[[str], Awaitable[None]]) -> None:
        """Send whatever is queued, in order, for as long as the socket lives."""
        queue = self._queue
        if queue is None:
            return
        while True:
            await send(self._encode(await queue.get()))

    async def _broadcast(self, text: str) -> None:
        """One frame to every connected client.

        ponytail: sequential, so a stalled client holds the others up for
        `timeout`. Per-client queues if a served socket ever has many readers.
        """
        for ws in list(self._clients):
            try:
                await asyncio.wait_for(ws.send_text(text), self.timeout)
            except Exception:
                self._clients.discard(ws)
                with contextlib.suppress(Exception):
                    await asyncio.wait_for(ws.close(code=1011), self.timeout)

    # -- client ---------------------------------------------------------------

    async def _client_loop(self) -> None:
        """Hold one connection. A drop raises, and the supervisor reconnects."""
        from websockets.asyncio.client import connect

        try:
            async with connect(
                self.url,
                additional_headers=self.headers,
                open_timeout=self.timeout,
            ) as ws:
                for frame in self.on_connect:
                    await ws.send(frame)
                self.report_health("ok")

                if self.mode is WebSocketNode.Mode.LISTENER:
                    async for received in ws:
                        if isinstance(received, str):
                            await self._receive(received)
                else:
                    sender = asyncio.create_task(self._send_queued(ws.send))
                    try:
                        # Replies are not a sender's to publish, but they are
                        # read: an unread buffer stops the pings being answered.
                        async for _ in ws:
                            pass
                    finally:
                        await self._cancel_task(sender)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            self.report_health("down", str(exc))
            raise
        self.report_health("down", "connection closed")
        raise ConnectionError(f"WebSocket for '{self.name}' closed")

    # -- server ---------------------------------------------------------------

    async def _serve(self, websocket: WebSocket) -> None:
        """One client of a served socket, for as long as it stays."""
        from fluksio.cloud.config import via_portal

        if (
            self.secret
            and not hmac.compare_digest(
                str(websocket.path_params.get("secret", "")).encode(),
                self.secret.encode(),
            )
        ) or (not self.secret and via_portal(websocket.headers)):
            # Closed before accepting: the same 403 as a path with no route.
            # Through the portal a socket needs a secret, as a webhook does.
            await websocket.close(code=1008)
            return
        await websocket.accept()
        self._clients.add(websocket)
        try:
            while True:
                message = await websocket.receive()
                if message["type"] == "websocket.disconnect":
                    break
                text = message.get("text")
                # A sender's clients are read only to notice them leave.
                if text is not None and self.mode is WebSocketNode.Mode.LISTENER:
                    await self._receive(text)
        finally:
            self._clients.discard(websocket)

    def register_route(self, app: FastAPI) -> None:
        """Mount the served socket, ahead of any catch-all, as webhooks are."""
        if self._route_registered:
            return
        from starlette.routing import WebSocketRoute

        # The secret is a path parameter, so the route table never holds it.
        path = f"{self.url}/{{secret}}" if self.secret else self.url
        app.routes.insert(
            _first_catch_all(app), WebSocketRoute(path, self._serve, name=self.id)
        )
        self._route_registered = True
        logger.info("Registered websocket route for node '%s': %s", self.name, path)

    def unregister_route(self, app: FastAPI) -> None:
        if not self._route_registered:
            return
        app.routes[:] = [
            route for route in app.routes if getattr(route, "name", None) != self.id
        ]
        self._route_registered = False

    # -- lifecycle ------------------------------------------------------------

    async def start(self, app: FastAPI | None = None) -> None:
        if self.serves and app is not None:
            self.register_route(app)
        if self.mode is WebSocketNode.Mode.SENDER and self._queue is None:
            # Outlives a reconnect, so what was queued meanwhile still goes.
            self._queue = asyncio.Queue(maxsize=self.send_queue_size)
            self._loop = asyncio.get_running_loop()
        # Spawned on every start: enabling a flow cancels its supervised loops
        # and starts the nodes again, and the supervisor ignores a duplicate.
        if self._task is not None and not self._task.done():
            return
        if not self.serves:
            self._task = self._run_supervised("websocket", self._client_loop)
        elif self.mode is WebSocketNode.Mode.SENDER:
            self._task = self._run_supervised(
                "websocket-out", lambda: self._send_queued(self._broadcast)
            )

    async def stop(self, app: FastAPI | None = None) -> None:
        if self.serves and app is not None:
            self.unregister_route(app)
        for ws in list(self._clients):
            with contextlib.suppress(Exception):
                await ws.close(code=1001)
        self._clients.clear()
        if self._task is not None:
            await self._cancel_task(self._task)
        self._task = None
        self._queue = None
        self._loop = None


def _frame(value: Any) -> str:
    """A value as frame text: a string as it stands, anything else as JSON."""
    # default=str: a value JSON cannot hold must not kill the send loop.
    return value if isinstance(value, str) else json.dumps(value, default=str)