import asyncio
import json
import uuid
from collections.abc import AsyncIterator
from contextlib import suppress
from typing import Any

from redis.asyncio import Redis

DIALOG_CHANNEL_PREFIX = "han:rt:dialog:"
USER_CHANNEL_PREFIX = "han:rt:user:"
# Backward-compatible name used by existing chat integrations.
CHANNEL_PREFIX = DIALOG_CHANNEL_PREFIX


class LocalFanout:
    def __init__(self) -> None:
        self._queues: set[asyncio.Queue[dict[str, Any]]] = set()

    async def publish(self, event: dict[str, Any]) -> None:
        for queue in tuple(self._queues):
            with suppress(asyncio.QueueFull):
                queue.put_nowait(event)

    async def subscribe(self) -> AsyncIterator[dict[str, Any]]:
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=256)
        self._queues.add(queue)
        try:
            while True:
                yield await queue.get()
        finally:
            self._queues.discard(queue)


class RealtimeFanout:
    def __init__(self, redis: Redis, local: LocalFanout | None = None) -> None:
        self.redis = redis
        self.local = local or LocalFanout()

    async def publish(self, event: dict[str, Any]) -> None:
        event = {"event_id": str(uuid.uuid4()), **event}
        channel = DIALOG_CHANNEL_PREFIX + str(event["dialog_id"])
        try:
            await self.redis.publish(channel, json.dumps(event, default=str, separators=(",", ":")))
        except Exception:
            await self.local.publish(event)

    async def publish_user(self, user_id: uuid.UUID, event: dict[str, Any]) -> None:
        event = {"event_id": str(uuid.uuid4()), "_user_id": str(user_id), **event}
        channel = USER_CHANNEL_PREFIX + str(user_id)
        try:
            await self.redis.publish(channel, json.dumps(event, default=str, separators=(",", ":")))
        except Exception:
            await self.local.publish(event)

    async def events(
        self,
        dialog_ids: set[uuid.UUID],
        user_id: uuid.UUID | None = None,
        notifications: bool = False,
    ) -> AsyncIterator[dict[str, Any]]:
        channels = [DIALOG_CHANNEL_PREFIX + str(dialog_id) for dialog_id in dialog_ids]
        if notifications and user_id is not None:
            channels.append(USER_CHANNEL_PREFIX + str(user_id))
        if not channels:
            await asyncio.Event().wait()
            return
        pubsub = self.redis.pubsub()
        try:
            await pubsub.subscribe(*channels)
        except Exception:
            await pubsub.aclose()
            async for event in self.local.subscribe():
                dialog_match = event.get("dialog_id") and uuid.UUID(
                    str(event["dialog_id"])
                ) in dialog_ids
                user_match = (
                    notifications
                    and user_id is not None
                    and event.get("_user_id") == str(user_id)
                )
                if dialog_match or user_match:
                    payload = dict(event)
                    payload.pop("_user_id", None)
                    yield payload
            return
        try:
            while True:
                message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)
                if message:
                    payload = json.loads(message["data"])
                    payload.pop("_user_id", None)
                    yield payload
                else:
                    await asyncio.sleep(0)
        finally:
            await pubsub.unsubscribe(*channels)
            await pubsub.aclose()