"""In-process publish/subscribe for the SSE status stream.

Works within one process. With several replicas, publish through Redis pub/sub or
PostgreSQL LISTEN/NOTIFY so a subscriber on any replica sees every update.
"""
import asyncio
from collections import defaultdict
from collections.abc import Iterator
from contextlib import contextmanager


class OrderEvents:
    def __init__(self) -> None:
        self._subscribers: dict[int, set[asyncio.Queue]] = defaultdict(set)

    @contextmanager
    def subscribe(self, order_id: int) -> Iterator[asyncio.Queue]:
        queue: asyncio.Queue = asyncio.Queue(maxsize=100)
        self._subscribers[order_id].add(queue)
        try:
            yield queue
        finally:
            self._subscribers[order_id].discard(queue)
            if not self._subscribers[order_id]:
                del self._subscribers[order_id]

    def publish(self, order_id: int, event: dict) -> None:
        for queue in self._subscribers.get(order_id, ()):
            if queue.full():                     # a slow consumer loses the oldest update, not the process
                queue.get_nowait()
            queue.put_nowait(event)

    def subscriber_count(self, order_id: int) -> int:
        return len(self._subscribers.get(order_id, ()))