"""
Channels for FastAPI — multiplexed pub/sub over the one WebSocket connection.

A channel names a group of subscribers and decides who may join it. `group(params)`
is the fan-out key, so two subscribers with the same params share a group and a push
addressed to those params reaches both.

Membership is held in this process. Django's channel layer carries groups across
workers through Redis; nothing here does, so a push reaches only the subscribers
connected to the process that sent it. One process — a desktop shell, a single
uvicorn worker — is the shape this serves.
"""

from __future__ import annotations

import asyncio
from collections import defaultdict
from typing import Any, ClassVar

from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel

# group name -> the live sockets subscribed to it
_groups: dict[str, set[Any]] = defaultdict(set)
_registry: dict[str, type["ReactChannel"]] = {}
_lock = asyncio.Lock()


class ReactChannel:
    """A named fan-out. Subclass, override what the channel decides, and register it.

    A subclass may define `on_connect(params)` / `on_disconnect()`; the socket handler
    calls them when they exist.
    """

    name: ClassVar[str] = ""

    def authorize(self, params: BaseModel | None = None) -> bool:
        """Whether this subscriber may join. Default: anyone may."""
        return True

    def group(self, params: BaseModel | None = None) -> str:
        """The fan-out key. Subscribers sharing it share every message sent to it."""
        if params is None:
            return self.name
        parts = sorted(f"{k}={v}" for k, v in params.model_dump().items())
        return f"{self.name}:{':'.join(parts)}" if parts else self.name

    def receive(self, params: BaseModel | None, msg: BaseModel) -> BaseModel | None:
        """What a client-sent message becomes for the group. None drops it."""
        return msg

    @classmethod
    async def push(cls, message: BaseModel | dict, **params: Any) -> None:
        """Send to every subscriber whose params key this group, from anywhere in the app."""
        channel = cls()
        key = channel.group(_Params(params) if params else None)
        await broadcast(key, cls.__name__, message, params)


class _Params:
    """Params given as keywords rather than a model, so `group` can read them uniformly."""

    def __init__(self, values: dict[str, Any]) -> None:
        self._values = values

    def model_dump(self) -> dict[str, Any]:
        return self._values


def register(channel_class: type[ReactChannel], name: str) -> None:
    channel_class.name = name
    _registry[name] = channel_class


def get_channel(name: str) -> type[ReactChannel] | None:
    return _registry.get(name)


def registered() -> dict[str, type[ReactChannel]]:
    return dict(_registry)


async def join(group: str, socket: Any) -> None:
    async with _lock:
        _groups[group].add(socket)


async def leave(group: str, socket: Any) -> None:
    async with _lock:
        _groups[group].discard(socket)
        if not _groups[group]:
            del _groups[group]


async def leave_all(socket: Any) -> None:
    async with _lock:
        for group in [g for g, sockets in _groups.items() if socket in sockets]:
            _groups[group].discard(socket)
            if not _groups[group]:
                del _groups[group]


async def members(group: str) -> set[Any]:
    async with _lock:
        return set(_groups.get(group, ()))


async def broadcast(
    group: str, type_name: str, message: BaseModel | dict, params: dict[str, Any] | None = None
) -> None:
    """Deliver to the group. A socket that fails to take it has departed, and is dropped."""
    payload = {
        "channel": group.split(":", 1)[0],
        "params": params or {},
        "type": type_name,
        "data": jsonable_encoder(message),
    }
    for socket in await members(group):
        try:
            await socket.send_json(payload)
        except Exception as e:
            print(f"mizan.channels: dropping subscriber from {group}: {type(e).__name__}: {e}")
            await leave(group, socket)
