"""SSE Event Streamer (Component F15)."""

from __future__ import annotations

import asyncio
import json
from collections import defaultdict

from gps_denied.schemas.events import (
    FlightCompletedEvent,
    FrameProcessedEvent,
    SearchExpandedEvent,
    SSEEventType,
    SSEMessage,
    UserInputNeededEvent,
)


class SSEEventStreamer:
    """Manages real-time SSE connections and event broadcasting."""

    def __init__(self) -> None:
        # Map: flight_id -> Dict[client_id, asyncio.Queue]
        self._streams: dict[str, dict[str, asyncio.Queue[SSEMessage | None]]] = defaultdict(dict)

    def create_stream(self, flight_id: str, client_id: str) -> asyncio.Queue[SSEMessage | None]:
        """Create a new event queue for a client."""
        q: asyncio.Queue[SSEMessage | None] = asyncio.Queue()
        self._streams[flight_id][client_id] = q
        return q

    def close_stream(self, flight_id: str, client_id: str) -> None:
        """Close a client stream by putting a sentinel and removing the queue."""
        if flight_id in self._streams and client_id in self._streams[flight_id]:
            q = self._streams[flight_id].pop(client_id)
            if not self._streams[flight_id]:
                del self._streams[flight_id]
            # Put None to signal generator exit
            try:
                q.put_nowait(None)
            except asyncio.QueueFull:
                pass

    def get_active_connections(self, flight_id: str) -> int:
        return len(self._streams.get(flight_id, {}))

    def _broadcast(self, flight_id: str, msg: SSEMessage) -> bool:
        """Broadcast a message to all clients subscribed to flight_id."""
        if flight_id not in self._streams or not self._streams[flight_id]:
            return False

        for q in self._streams[flight_id].values():
            try:
                q.put_nowait(msg)
            except asyncio.QueueFull:
                pass  # Drop if queue is full rather than blocking
        return True

    # ── Business Event Senders ────────────────────────────────────────────────

    def send_frame_result(self, flight_id: str, event_data: FrameProcessedEvent) -> bool:
        msg = SSEMessage(
            event=SSEEventType.FRAME_PROCESSED,
            data=event_data.model_dump(mode="json"),
            id=f"frame_{event_data.frame_id}",
        )
        return self._broadcast(flight_id, msg)

    def send_refinement(self, flight_id: str, event_data: FrameProcessedEvent) -> bool:
        msg = SSEMessage(
            event=SSEEventType.FRAME_REFINED,
            data=event_data.model_dump(mode="json"),
            id=f"refine_{event_data.frame_id}",
        )
        return self._broadcast(flight_id, msg)

    def send_search_progress(self, flight_id: str, event_data: SearchExpandedEvent) -> bool:
        msg = SSEMessage(
            event=SSEEventType.SEARCH_EXPANDED,
            data=event_data.model_dump(mode="json"),
        )
        return self._broadcast(flight_id, msg)

    def send_user_input_request(self, flight_id: str, event_data: UserInputNeededEvent) -> bool:
        msg = SSEMessage(
            event=SSEEventType.USER_INPUT_NEEDED,
            data=event_data.model_dump(mode="json"),
        )
        return self._broadcast(flight_id, msg)

    def send_flight_completed(self, flight_id: str, event_data: FlightCompletedEvent) -> bool:
        msg = SSEMessage(
            event=SSEEventType.FLIGHT_COMPLETED,
            data=event_data.model_dump(mode="json"),
        )
        return self._broadcast(flight_id, msg)

    def send_heartbeat(self, flight_id: str) -> bool:
        # sse_starlette uses empty string or comment for heartbeat,
        # but we can just send an SSEMessage object that parses as empty event
        if flight_id not in self._streams:
            return False

        # Manually sending a comment via the generator is tricky with strict SSEMessage schema
        # but we'll handle this in the stream generator directly
        return True

    # ── Generic event dispatcher (used by processor.process_frame) ──────────

    async def push_event(self, flight_id: str, event_type: str, data: dict) -> None:
        """Dispatch a generic event to all clients for a flight.

        Maps event_type strings to typed SSE events:
          "frame_result"  → FrameProcessedEvent
          "refinement"    → FrameProcessedEvent (refined)
          Other           → raw broadcast via SSEMessage
        """
        if event_type == "frame_result":
            evt = FrameProcessedEvent(**data) if not isinstance(data, FrameProcessedEvent) else data
            self.send_frame_result(flight_id, evt)
        elif event_type == "refinement":
            evt = FrameProcessedEvent(**data) if not isinstance(data, FrameProcessedEvent) else data
            self.send_refinement(flight_id, evt)
        else:
            msg = SSEMessage(
                event=SSEEventType.FRAME_PROCESSED,
                data=data,
                id=str(data.get("frame_id", "")),
            )
            self._broadcast(flight_id, msg)

    # ── Stream Generator ──────────────────────────────────────────────────────

    async def stream_generator(self, flight_id: str, client_id: str):
        """Yields dicts for sse_starlette EventSourceResponse."""
        q = self.create_stream(flight_id, client_id)

        # Send an immediate connection accepted ping
        yield {"event": "connected", "data": "connected"}

        try:
            while True:
                # Wait for next event or send heartbeat every 15s
                try:
                    msg = await asyncio.wait_for(q.get(), timeout=15.0)
                    if msg is None:
                        # Sentinel for clean shutdown
                        break

                    # Yield dict format for sse_starlette
                    yield {
                        "event": msg.event.value,
                        "id": msg.id if msg.id else "",
                        "data": json.dumps(msg.data)
                    }

                except asyncio.TimeoutError:
                    # Heartbeat format for sse_starlette (empty string generates a comment)
                    yield {"event": "heartbeat", "data": "ping"}

        except asyncio.CancelledError:
            pass  # Client disconnected
        finally:
            self.close_stream(flight_id, client_id)
