"""Live SSE session tracking + forced-close on token revoke.

Revoking a token via ``DELETE /api/mcp/tokens/{id}`` previously only
blocked NEW connections — established SSE streams kept flowing until the
client noticed the next POST got 401 (which could be never on a long
read-only session). That's a real security gap: a leaked token in active
use would keep its grip until the server restarted.

This module closes the gap. Every SSE handler registers its session here
on entry; the revoke API signals the matching ``asyncio.Event`` and the
handler's ``asyncio.wait()`` race exits, cancelling the MCP server loop.

The tracker is in-process (per-server). Sessions don't survive a server
restart — but restart kills SSE streams anyway, so that's the safe state.
"""

from __future__ import annotations

import asyncio
import logging
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import UTC, datetime

log = logging.getLogger(__name__)


@dataclass(eq=False)
class _SessionHandle:
    """One live SSE session. ``revoke_event`` flips when the bearer
    token is revoked, signalling the handler to exit.

    ``eq=False`` keeps identity-based equality + ``object.__hash__`` so
    handles can live in a ``set`` even though they hold a mutable
    ``asyncio.Event``. Each handle is a unique instance representing
    one connection — value equality wouldn't make sense anyway.
    """

    token_id: str
    opened_at: str
    revoke_event: asyncio.Event = field(default_factory=asyncio.Event)


class SseSessionTracker:
    """Process-wide registry of live SSE sessions, keyed by token id.

    Multiple sessions per token are allowed — one user connecting from
    two devices with the same token shows up as two entries, both
    closed in lockstep when the token is revoked.
    """

    def __init__(self) -> None:
        # token_id → set of handles. Using a set lets us add / remove
        # individual sessions without re-keying.
        self._sessions: dict[str, set[_SessionHandle]] = defaultdict(set)
        self._lock = asyncio.Lock()

    async def register(self, token_id: str) -> _SessionHandle:
        """Register a fresh session. Returns the handle the handler must
        ``await`` on its ``revoke_event`` to detect forced close.
        """
        handle = _SessionHandle(
            token_id=token_id,
            opened_at=datetime.now(UTC).isoformat(),
        )
        async with self._lock:
            self._sessions[token_id].add(handle)
        return handle

    async def unregister(self, handle: _SessionHandle) -> None:
        """Drop a session from the tracker. Idempotent — calling twice
        on the same handle (e.g., handler exit AFTER revoke fired) is
        a no-op.
        """
        async with self._lock:
            sessions = self._sessions.get(handle.token_id)
            if not sessions:
                return
            sessions.discard(handle)
            if not sessions:
                # Clean up the empty bucket so ``active_count`` stays honest.
                del self._sessions[handle.token_id]

    async def revoke(self, token_id: str) -> int:
        """Signal every session for ``token_id`` to close. Returns the
        number of sessions signalled (0 if no live sessions).

        The handler is responsible for un-registering — we DON'T pop the
        bucket here, because doing so would break the ``unregister``
        idempotence path.
        """
        async with self._lock:
            sessions = list(self._sessions.get(token_id, set()))
        for handle in sessions:
            handle.revoke_event.set()
        return len(sessions)

    def active_count(self, token_id: str | None = None) -> int:
        """Count of live sessions. Lock-free read for cheap polling.

        Reading without the lock is safe for a counter — worst case we
        see a stale value during a register/unregister race, which is
        fine for a UI "currently streaming" indicator.
        """
        if token_id is not None:
            return len(self._sessions.get(token_id, set()))
        return sum(len(s) for s in self._sessions.values())

    def snapshot(self) -> dict[str, dict[str, int | list[str]]]:
        """Per-token counts + opened_at timestamps. For the
        ``GET /api/mcp/sessions`` admin endpoint. Lock-free; the data
        is small and tearing during a register race is harmless.
        """
        return {
            token_id: {
                "count": len(sessions),
                "opened_at": sorted(h.opened_at for h in sessions),
            }
            for token_id, sessions in self._sessions.items()
        }
