"""Broker-agnostic SSE streaming primitive.

`SSEStreamService.open_stream` acquires a slot (raising
:class:`CapExceededError` *synchronously* if the cap is full — so a
consumer maps it to 429 BEFORE opening the stream), subscribes via the
broker, and returns an :class:`SSEStream` (headers + an async byte
iterator). The iterator emits the ``retry:`` directive, heartbeat
comments, and ``data:`` frames, and in a ``finally`` block releases the
slot + closes the subscription on disconnect / cancel / error.

Framework-agnostic: the consumer supplies an async ``is_disconnected``
callable (FastAPI's ``request.is_disconnected``) and wraps the iterator
in whatever streaming response its framework uses.
:mod:`mgf.livepush.fastapi` ships a one-line wrapper.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING

from mgf.livepush._errors import CapExceededError

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Awaitable, Callable

    from mgf.livepush._broker import Broker, Subscription
    from mgf.livepush._slotcap import SlotCap

__all__ = ["SSEStream", "SSEStreamService", "sse_headers"]

LOG = logging.getLogger(__name__)

# EventSource reconnect cadence (ms). Pairs with a 429 Retry-After: 5 so
# a flapping client doesn't tighten into a hot loop against the cap.
_RETRY_MS = 5000
# Heartbeat envelope — survives AWS ALB / Cloudflare 30 s idle-close.
_DEFAULT_HEARTBEAT_SECONDS = 15.0


def sse_headers() -> dict[str, str]:
    """Headers every SSE response needs so proxies pass the wire
    untouched (disable caching/transform + Nginx-style buffering)."""
    return {"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"}


@dataclass
class SSEStream:
    """The result of :meth:`SSEStreamService.open_stream` — the headers to
    set on the response + the byte iterator to stream as its body."""

    headers: dict[str, str]
    body: AsyncIterator[bytes]


class SSEStreamService:
    """Per-app SSE primitive over a :class:`Broker` + :class:`SlotCap`."""

    def __init__(
        self,
        broker: Broker,
        slot_cap: SlotCap,
        *,
        heartbeat_seconds: float = _DEFAULT_HEARTBEAT_SECONDS,
    ) -> None:
        self._broker = broker
        self._slot_cap = slot_cap
        self._heartbeat_s = heartbeat_seconds

    async def open_stream(
        self,
        *,
        channel: str,
        slot_key: str,
        is_disconnected: Callable[[], Awaitable[bool]],
    ) -> SSEStream:
        """Acquire a slot, subscribe to ``channel``, return an
        :class:`SSEStream`. Raises :class:`CapExceededError` synchronously
        when the cap is full — no subscription / slot leak (the acquire
        happens before any pubsub state)."""
        acquired = await self._slot_cap.acquire(slot_key)
        if not acquired:
            raise CapExceededError(slot_key, self._slot_cap.cap)
        sub = await self._broker.subscribe(channel)
        return SSEStream(
            headers=sse_headers(),
            body=self._event_gen(sub, slot_key, is_disconnected),
        )

    async def _event_gen(
        self,
        sub: Subscription,
        slot_key: str,
        is_disconnected: Callable[[], Awaitable[bool]],
    ) -> AsyncIterator[bytes]:
        try:
            yield f"retry: {_RETRY_MS}\n\n".encode()
            while True:
                data = await sub.get_message(timeout=self._heartbeat_s)
                if await is_disconnected():
                    return
                if data is None:
                    yield b": keep-alive\n\n"
                    continue
                yield b"data: " + data + b"\n\n"
        except asyncio.CancelledError:
            return
        finally:
            # Close the subscription BEFORE releasing the slot: if release
            # came first and close hung, the slot would be reusable before
            # the underlying resource was freed (slot-leak adjacent).
            with contextlib.suppress(Exception):
                await sub.close()
            await self._slot_cap.release(slot_key)
