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

from redis.asyncio import Redis

CHANNEL_PREFIX = "han:rt:dialog:"


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 = 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 events(self, dialog_ids: set[uuid.UUID]) -> AsyncIterator[dict[str, Any]]:
        channels = [CHANNEL_PREFIX + str(dialog_id) for dialog_id in dialog_ids]
        pubsub = self.redis.pubsub()
        try:
            await pubsub.subscribe(*channels)
        except Exception:
            await pubsub.aclose()
            async for event in self.local.subscribe():
                if uuid.UUID(str(event["dialog_id"])) in dialog_ids:
                    yield event
            return
        try:
            while True:
                message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)
                if message:
                    yield json.loads(message["data"])
                else:
                    await asyncio.sleep(0)
        finally:
            await pubsub.unsubscribe(*channels)
            await pubsub.aclose()