from __future__ import annotations

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

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

from relay.server.state import RelayState

if TYPE_CHECKING:
    from relay.server.db import RelayDB


async def notification_stream(
    state: RelayState, user_id: str
) -> AsyncIterator[ServerSentEvent]:
    queue = await state.connect_sse(user_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(user_id)


def register_sse_routes(
    mcp: FastMCP, state: RelayState, db: RelayDB | None = None
) -> None:
    @mcp.custom_route("/notifications/{user_id}", methods=["GET"])
    async def notifications_endpoint(request: Request) -> Response:
        user_id = request.path_params["user_id"]
        if db is not None:
            api_key = request.headers.get("x-api-key", "")
            if not api_key:
                return JSONResponse({"error": "unauthorized"}, status_code=401)
            account = await db.get_account_by_api_key(api_key)
            if account is None:
                return JSONResponse({"error": "unauthorized"}, status_code=401)
            user = await db.get_user(user_id)
            if user is None or user.account_email != account.email:
                return JSONResponse({"error": "forbidden"}, status_code=403)
        return EventSourceResponse(notification_stream(state, user_id))
