"""
The WebSocket endpoint — channel subscriptions and RPC over one connection.

    Client sends:
        {"action": "subscribe",   "channel": "chat", "params": {...}}
        {"action": "unsubscribe", "channel": "chat", "params": {...}}
        {"action": "message",     "channel": "chat", "params": {...}, "data": {...}}
        {"action": "rpc", "id": "request-id", "fn": "function_name", "args": {...}}
        {"action": "ctx", "id": "request-id", "context": "name", "params": {...}}

    Server sends:
        {"channel": "chat", "params": {...}, "type": "...", "data": {...}}
        {"id": "request-id", "ok": true,  "data": {...}}
        {"id": "request-id", "ok": false, "error": {"code": "...", "message": "..."}}
        {"error": "..."}

The envelope is mizan-django's, so a client speaks to either backend unchanged. RPC
dispatches through the same `execute_function` the HTTP route calls — one dispatch path,
so the two transports cannot disagree about what a function does.

`data` is the `{result, invalidate, merge}` envelope every non-HTTP transport hands the
kernel, so `mizanCall` applies a socket mutation's invalidation exactly as it applies an
HTTP one. mizan-django's socket sends a bare result and drops it.
"""

from __future__ import annotations

import logging
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from mizan_core.registry import get_context_groups, get_function
from mizan_fastapi import channels
from mizan_fastapi.executor import (
    MizanError,
    compute_invalidation,
    compute_merges,
    execute_function,
)

logger = logging.getLogger(__name__)

ws_router = APIRouter()


class _SocketRequest:
    """What a server function receives when the call arrived over the socket.

    The HTTP path hands functions a Starlette `Request`; there is none here, so this
    carries the surface a function actually reads — `state`, headers, and a method, since
    a socket RPC sends data and expects an answer.
    """

    method = "POST"

    def __init__(self, socket: WebSocket) -> None:
        self.state = socket.state
        self.scope = socket.scope
        self.headers = socket.headers
        self.query_params = socket.query_params
        self.socket = socket


def _params_model(channel_cls: Any, raw: dict[str, Any] | None) -> Any:
    """Params as the channel's declared model, or a bare holder when it declares none."""
    model = getattr(channel_cls, "Params", None)
    if model is not None and raw:
        return model(**raw)
    return channels._Params(raw) if raw else None


def _resolve(body: dict[str, Any]) -> tuple[Any, Any, str] | None:
    """The channel class, its params, and the name — or None when the name is unknown."""
    name = body.get("channel") or ""
    channel_cls = channels.get_channel(name)
    if channel_cls is None:
        return None
    return channel_cls, _params_model(channel_cls, body.get("params")), name


async def _subscribe(socket: WebSocket, body: dict[str, Any]) -> None:
    found = _resolve(body)
    if found is None:
        await socket.send_json({"error": f"unknown channel {body.get('channel')!r}"})
        return
    channel_cls, params, name = found

    channel = channel_cls()
    if not channel.authorize(params):
        await socket.send_json({"error": f"not authorized for channel {name!r}"})
        return

    await channels.join(channel.group(params), socket)
    hook = getattr(channel, "on_connect", None)
    if hook is not None:
        await hook(params)


async def _unsubscribe(socket: WebSocket, body: dict[str, Any]) -> None:
    found = _resolve(body)
    if found is None:
        return
    channel_cls, params, _ = found
    channel = channel_cls()
    await channels.leave(channel.group(params), socket)
    hook = getattr(channel, "on_disconnect", None)
    if hook is not None:
        await hook()


async def _message(socket: WebSocket, body: dict[str, Any]) -> None:
    found = _resolve(body)
    if found is None:
        await socket.send_json({"error": f"unknown channel {body.get('channel')!r}"})
        return
    channel_cls, params, name = found

    channel = channel_cls()
    if not channel.authorize(params):
        await socket.send_json({"error": f"not authorized for channel {name!r}"})
        return

    outgoing = channel.receive(params, body.get("data") or {})
    if outgoing is None:
        return  # the channel dropped it
    await channels.broadcast(
        channel.group(params), type(outgoing).__name__, outgoing, body.get("params")
    )


async def _rpc(socket: WebSocket, body: dict[str, Any]) -> None:
    request_id = body.get("id")
    fn_name = body.get("fn")
    if not fn_name:
        await socket.send_json(
            {
                "id": request_id,
                "ok": False,
                "error": {"code": "BAD_REQUEST", "message": "rpc requires 'fn'"},
            }
        )
        return

    args = body.get("args") or {}
    try:
        fn_class = get_function(fn_name)
        result = await execute_function(_SocketRequest(socket), fn_name, args)
    except MizanError as e:
        payload: dict[str, Any] = {"code": e.code.value, "message": e.message}
        if e.details:
            payload["details"] = e.details
        await socket.send_json({"id": request_id, "ok": False, "error": payload})
        return

    data: dict[str, Any] = {
        "result": result,
        "invalidate": compute_invalidation(fn_class, args),
    }
    merges = compute_merges(fn_class, args, result)
    if merges:
        data["merge"] = merges
    await socket.send_json({"id": request_id, "ok": True, "data": data})


async def _ctx(socket: WebSocket, body: dict[str, Any]) -> None:
    """A context bundle over the socket, so a client needs no second connection to read."""
    request_id = body.get("id")
    name = body.get("context") or ""
    fn_names = get_context_groups().get(name)
    if not fn_names:
        await socket.send_json(
            {
                "id": request_id,
                "ok": False,
                "error": {"code": "NOT_FOUND", "message": f"Context '{name}' not found"},
            }
        )
        return

    params = body.get("params") or {}
    request = _SocketRequest(socket)
    try:
        bundled = {fn: await execute_function(request, fn, params) for fn in fn_names}
    except MizanError as e:
        payload: dict[str, Any] = {"code": e.code.value, "message": e.message}
        if e.details:
            payload["details"] = e.details
        await socket.send_json({"id": request_id, "ok": False, "error": payload})
        return

    await socket.send_json({"id": request_id, "ok": True, "data": bundled})


_ACTIONS = {
    "subscribe": _subscribe,
    "unsubscribe": _unsubscribe,
    "message": _message,
    "rpc": _rpc,
    "ctx": _ctx,
}


@ws_router.websocket("/ws/")
async def mizan_socket(socket: WebSocket) -> None:
    """One connection, every action. A disconnect ends the loop and clears the membership.

    A close is how a socket ends, so the disconnect is reported as the departure it is
    rather than raised as a fault; `finally` clears the membership either way.
    """
    await socket.accept()
    try:
        while True:
            body = await socket.receive_json()
            handler = _ACTIONS.get(body.get("action"))
            if handler is None:
                await socket.send_json({"error": f"unknown action {body.get('action')!r}"})
                continue
            await handler(socket, body)
    except WebSocketDisconnect as e:
        logger.debug("mizan socket closed: code=%s reason=%s", e.code, e.reason)
    finally:
        await channels.leave_all(socket)
