import asyncio
import json
from collections.abc import AsyncIterator
from typing import Any

from fastmcp import FastMCP
from sse_starlette.sse import EventSourceResponse, ServerSentEvent
from starlette.requests import Request
from starlette.responses import Response

from relay.server.state import RelayState


async def notification_stream(
    state: RelayState, group_id: str, peer_id: str
) -> AsyncIterator[ServerSentEvent]:
    queue = await state.connect_sse(peer_id)
    try:
        while True:
            payload: dict[str, Any] = await queue.get()
            yield ServerSentEvent(
                data=json.dumps(payload),
                event="notification",
                id=payload["message_id"],
            )
    except asyncio.CancelledError:
        pass
    finally:
        await state.disconnect_sse(peer_id)


def register_sse_routes(mcp: FastMCP, state: RelayState) -> None:
    @mcp.custom_route("/notifications/{group_id}/{peer_id}", methods=["GET"])
    async def notifications_endpoint(request: Request) -> Response:
        group_id = request.path_params["group_id"]
        peer_id = request.path_params["peer_id"]
        return EventSourceResponse(notification_stream(state, group_id, peer_id))
