"""Server-sent events helpers shared by the markdown and API frontends."""

from __future__ import annotations

import asyncio
import contextlib
import json
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

__all__ = ["HEARTBEAT_INTERVAL", "SSE_HEADERS", "events_with_heartbeat", "sse_pack"]

HEARTBEAT_INTERVAL = 15.0
"""Seconds of backend silence before a comment frame keeps the socket warm."""

SSE_HEADERS = {
    "Cache-Control": "no-cache, no-transform",
    "Connection": "keep-alive",
    "X-Accel-Buffering": "no",
}

KEEPALIVE = b": keepalive\n\n"


async def events_with_heartbeat(
    events: AsyncIterator[Any], interval: float = HEARTBEAT_INTERVAL
) -> AsyncIterator[Any]:
    """Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence.

    Only one consumer may iterate the result at a time; cancelling the
    outer scope cancels the in-flight upstream fetch instead of leaking it.
    """
    src = events.__aiter__()
    next_task: asyncio.Task[Any] | None = None
    try:
        while True:
            if next_task is None:
                next_task = asyncio.ensure_future(src.__anext__())
            done, _pending = await asyncio.wait({next_task}, timeout=interval)
            if not done:
                yield None
                continue
            task = next_task
            next_task = None
            try:
                result = task.result()
            except StopAsyncIteration:
                return
            yield result
    finally:
        if next_task is not None and not next_task.done():
            next_task.cancel()
            with contextlib.suppress(BaseException):
                await next_task


def sse_pack(event: str, data: dict[str, Any]) -> bytes:
    body = json.dumps(data, ensure_ascii=False)
    return f"event: {event}\ndata: {body}\n\n".encode()