from __future__ import annotations

import asyncio
from collections import defaultdict

from app.core.api.domain.events.client_event import ClientEventRecord
from app.core.api.infrastructure.streaming.replay_buffer import ReplayBuffer


class SseEventChannel:
    def __init__(self, replay_buffer: ReplayBuffer | None = None) -> None:
        self._channels: dict[str, list[asyncio.Queue[ClientEventRecord]]] = defaultdict(list)
        self._replay = replay_buffer or ReplayBuffer()
        self._lock = asyncio.Lock()

    async def subscribe(self, request_id: str, replay: bool = True) -> asyncio.Queue[ClientEventRecord]:
        queue: asyncio.Queue[ClientEventRecord] = asyncio.Queue()
        snapshot: list[ClientEventRecord] = []
        async with self._lock:
            self._channels[request_id].append(queue)
            if replay:
                snapshot = self._replay.list(request_id)
        for item in snapshot:
            await queue.put(item)
        return queue

    async def unsubscribe(self, request_id: str, queue: asyncio.Queue[ClientEventRecord]) -> None:
        async with self._lock:
            queues = self._channels.get(request_id)
            if not queues:
                return
            if queue in queues:
                queues.remove(queue)
            if not queues:
                self._channels.pop(request_id, None)

    async def publish(self, event: ClientEventRecord) -> None:
        async with self._lock:
            self._replay.append(event.request_id, event)
            queues = list(self._channels.get(event.request_id, []))
        for queue in queues:
            await queue.put(event)
