"""WebSocket server for real-time NEUROVAULT event streaming."""

from __future__ import annotations

import asyncio
import hmac
import json
import logging
import time
from typing import Any, Set

logger = logging.getLogger(__name__)


class EventStream:
    """Central event bus for pub/sub streaming."""

    def __init__(self):
        self._subscribers: Set[asyncio.Queue] = set()

    async def publish(self, event_type: str, data: dict[str, Any]) -> None:
        event = {
            "type": event_type,
            "timestamp": int(time.time() * 1000),
            "data": data,
        }

        stale: set[asyncio.Queue] = set()
        for queue in self._subscribers:
            try:
                queue.put_nowait(event)
            except asyncio.QueueFull:
                stale.add(queue)
        self._subscribers -= stale

    def subscribe(self) -> asyncio.Queue:
        try:
            asyncio.get_event_loop()
        except RuntimeError:
            asyncio.set_event_loop(asyncio.new_event_loop())
        queue: asyncio.Queue = asyncio.Queue(maxsize=100)
        self._subscribers.add(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue) -> None:
        self._subscribers.discard(queue)


class WebSocketServer:
    """WebSocket relay for realtime events."""

    def __init__(
        self,
        event_stream: EventStream,
        port: int = 9475,
        host: str = "127.0.0.1",
        api_keys: set[str] | None = None,
        max_connections: int = 10,
    ):
        self._stream = event_stream
        self._port = int(port)
        self._host = host
        self._api_keys = api_keys
        self._connections = 0
        self._max_connections = int(max_connections)
        self._server = None

    async def start(self) -> bool:
        try:
            import websockets
        except Exception:  # pragma: no cover - optional dependency
            logger.error("websockets required: pip install websockets")
            return False

        self._server = await websockets.serve(
            self._handle,
            self._host,
            self._port,
            max_size=1_000_000,
        )
        logger.info("WebSocket event stream on %s:%d", self._host, self._port)
        return True

    async def run_forever(self) -> None:
        ok = await self.start()
        if not ok:
            return
        assert self._server is not None
        await self._server.wait_closed()

    async def stop(self) -> None:
        if self._server is None:
            return
        self._server.close()
        await self._server.wait_closed()
        self._server = None

    async def _handle(self, ws, path=None):  # pragma: no cover - network I/O
        del path
        if self._connections >= self._max_connections:
            await ws.close(4003, "Max connections reached")
            return

        if self._api_keys:
            try:
                auth_msg = await asyncio.wait_for(ws.recv(), timeout=5)
                auth = json.loads(auth_msg)
                token = str(auth.get("token", ""))
                if not any(hmac.compare_digest(token, candidate) for candidate in self._api_keys):
                    await ws.close(4001, "Unauthorized")
                    return
            except Exception:
                await ws.close(4001, "Auth required")
                return

        self._connections += 1
        queue = self._stream.subscribe()

        try:
            filters = set()
            try:
                filter_msg = await asyncio.wait_for(ws.recv(), timeout=3)
                filter_data = json.loads(filter_msg)
                filters = set(filter_data.get("subscribe", []))
            except (asyncio.TimeoutError, json.JSONDecodeError):
                pass

            while True:
                event = await queue.get()
                if filters and event.get("type") not in filters:
                    continue
                await ws.send(json.dumps(event))

        except Exception as exc:
            logger.debug("WebSocket stream handler closed: %s", exc)
        finally:
            self._stream.unsubscribe(queue)
            self._connections -= 1


EVENT_TYPES = {
    "memory.ingested": "New memory stored",
    "memory.deduplicated": "Duplicate memory skipped",
    "memory.filtered": "Memory filtered by content filter",
    "consolidation.started": "Consolidation cycle began",
    "consolidation.entity_extracted": "New entity extracted",
    "consolidation.relationship_found": "New relationship discovered",
    "consolidation.cluster_formed": "New cluster formed",
    "consolidation.insight_generated": "New insight generated",
    "consolidation.completed": "Consolidation cycle finished",
    "entity.created": "New entity added to graph",
    "entity.updated": "Entity importance/properties updated",
    "relationship.created": "New relationship added",
    "concept.hierarchy_added": "New IS-A/PART-OF link added",
    "prediction.surfaced": "Memory proactively surfaced",
    "prediction.calendar": "Calendar-triggered memory surfacing",
    "peer.discovered": "New peer found via mDNS",
    "peer.connected": "Peer connection established",
    "peer.disconnected": "Peer disconnected",
    "sync.started": "Sync with peer started",
    "sync.completed": "Sync with peer completed",
    "sync.conflict": "Merge conflict detected",
    "sensor.started": "Sensor activated",
    "sensor.stopped": "Sensor deactivated",
    "plugin.loaded": "Plugin loaded",
    "plugin.error": "Plugin encountered error",
    "daemon.started": "Daemon started",
    "daemon.stopped": "Daemon stopped",
}
