"""WebSocket real-time hub for otelctx live signal streaming."""

from __future__ import annotations

import asyncio
import json
import threading
from collections import deque
from typing import Any


class RealtimeHub:
    """Async WebSocket broadcast hub with an internal ring-buffer.

    Designed to run inside a FastAPI / Starlette asyncio event loop.
    Signals produced from synchronous ingestion threads are enqueued via
    ``publish_sync()`` (thread-safe); the hub's flush loop broadcasts them
    to all live WebSocket clients every ``flush_interval_ms`` milliseconds.

    Usage::

        hub = RealtimeHub(buffer_size=200, flush_interval_ms=500)

        # In FastAPI startup:
        asyncio.create_task(hub.flush_loop())

        # WebSocket endpoint:
        @app.websocket("/ws/events")
        async def ws_events(websocket: WebSocket):
            await hub.connect(websocket)

        # From ingestion code (sync thread):
        hub.publish_sync({"type": "trace", "trace_id": "abc"})
    """

    def __init__(self, buffer_size: int = 100, flush_interval_ms: int = 500) -> None:
        self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size)
        self._flush_interval = flush_interval_ms / 1000.0
        self._clients: set = set()
        self._lock = threading.Lock()
        self._async_queue: asyncio.Queue[dict[str, Any]] | None = None

    # ------------------------------------------------------------------
    # Async API (call from within an event loop)
    # ------------------------------------------------------------------

    def _ensure_queue(self) -> asyncio.Queue:
        if self._async_queue is None:
            self._async_queue = asyncio.Queue(maxsize=self._buffer.maxlen or 1000)
        return self._async_queue

    async def connect(self, websocket) -> None:
        """Accept a WebSocket connection and keep it alive until disconnect."""
        from starlette.websockets import WebSocketState

        await websocket.accept()
        self._clients.add(websocket)
        try:
            while websocket.client_state == WebSocketState.CONNECTED:
                # Drain any pending messages addressed to this client
                await asyncio.sleep(self._flush_interval)
        except Exception:
            pass
        finally:
            self._clients.discard(websocket)

    async def broadcast(self, message: dict[str, Any]) -> None:
        """Immediately broadcast *message* to all connected clients."""
        text = json.dumps(message, default=str)
        dead = set()
        for ws in list(self._clients):
            try:
                await ws.send_text(text)
            except Exception:
                dead.add(ws)
        self._clients -= dead

    async def flush_loop(self) -> None:
        """Coroutine that drains the hub queue and broadcasts in batches.

        Schedule once at application startup:
        ``asyncio.create_task(hub.flush_loop())``
        """
        queue = self._ensure_queue()
        while True:
            await asyncio.sleep(self._flush_interval)
            batch: list[dict] = []
            while not queue.empty():
                try:
                    batch.append(queue.get_nowait())
                except asyncio.QueueEmpty:
                    break
            if batch and self._clients:
                await self.broadcast({"type": "batch", "events": batch})

    # ------------------------------------------------------------------
    # Sync API (safe to call from non-async threads)
    # ------------------------------------------------------------------

    def publish_sync(self, event: dict[str, Any]) -> None:
        """Enqueue an event from a sync context (thread-safe)."""
        with self._lock:
            self._buffer.append(event)
        queue = self._async_queue
        if queue is not None:
            try:
                queue.put_nowait(event)
            except asyncio.QueueFull:
                pass  # Drop oldest-ish — deque already handles it

    # ------------------------------------------------------------------
    # Observability
    # ------------------------------------------------------------------

    def stats(self) -> dict[str, int]:
        return {
            "connected_clients": len(self._clients),
            "buffer_size": len(self._buffer),
            "buffer_capacity": self._buffer.maxlen or 0,
        }


def build_ws_router(hub: RealtimeHub):
    """Return a FastAPI APIRouter with WebSocket endpoints for the hub."""
    from fastapi import APIRouter
    from fastapi import WebSocket

    router = APIRouter()

    @router.get("/ws/health")
    async def ws_health() -> dict:
        return hub.stats()

    @router.websocket("/ws/events")
    async def ws_events(websocket: WebSocket) -> None:
        await hub.connect(websocket)

    @router.websocket("/ws")
    async def ws_root(websocket: WebSocket) -> None:
        await hub.connect(websocket)

    return router
