"""In-process SSE pub/sub keyed by run id (CONTRACT §4).

A single-process broker holding one fan-out set of ``asyncio.Queue`` subscribers
per ``run_id``. Producers (the run pipeline) call :meth:`SSEBroker.publish`;
the ``GET /runs/{id}/stream`` endpoint subscribes and forwards formatted events.

Event types (CONTRACT §4): ``run.status``, ``agent.spawned``, ``agent.state``,
``line.decided``, ``flow.span``, ``run.done``, ``run.error``.
"""

from __future__ import annotations

import asyncio
import json
from typing import Any, AsyncIterator

# Sentinel pushed to subscribers to signal the stream is finished.
_CLOSE = object()


def format_sse_event(event: str, data: Any) -> str:
    """Render one ``event:`` / ``data:`` SSE frame."""
    payload = json.dumps(data, default=str)
    return f"event: {event}\ndata: {payload}\n\n"


class SSEBroker:
    """Fan-out pub/sub of SSE events per run id."""

    def __init__(self) -> None:
        self._subscribers: dict[int, set[asyncio.Queue]] = {}
        self._history: dict[int, list[tuple[str, Any]]] = {}
        self._closed: set[int] = set()
        self._lock = asyncio.Lock()

    async def publish(self, run_id: int, event: str, data: Any) -> None:
        """Publish an event to all subscribers of ``run_id`` (and history)."""
        async with self._lock:
            self._history.setdefault(run_id, []).append((event, data))
            queues = list(self._subscribers.get(run_id, ()))
        for queue in queues:
            queue.put_nowait((event, data))

    def publish_threadsafe(
        self, loop: asyncio.AbstractEventLoop, run_id: int, event: str, data: Any
    ) -> None:
        """Publish from a worker thread onto the service event loop."""
        asyncio.run_coroutine_threadsafe(self.publish(run_id, event, data), loop)

    async def close(self, run_id: int) -> None:
        """Mark a run finished and release its subscribers."""
        async with self._lock:
            self._closed.add(run_id)
            queues = list(self._subscribers.get(run_id, ()))
        for queue in queues:
            queue.put_nowait(_CLOSE)

    async def subscribe(self, run_id: int) -> AsyncIterator[tuple[str, Any]]:
        """Yield events for ``run_id``: buffered history first, then live."""
        queue: asyncio.Queue = asyncio.Queue()
        async with self._lock:
            replay = list(self._history.get(run_id, ()))
            already_closed = run_id in self._closed
            self._subscribers.setdefault(run_id, set()).add(queue)
        try:
            for event, data in replay:
                yield event, data
            if already_closed:
                return
            while True:
                item = await queue.get()
                if item is _CLOSE:
                    return
                yield item  # type: ignore[misc]
        finally:
            async with self._lock:
                subs = self._subscribers.get(run_id)
                if subs is not None:
                    subs.discard(queue)
                    if not subs:
                        self._subscribers.pop(run_id, None)


# Process-wide broker instance shared by the pipeline and the API.
broker = SSEBroker()