"""WebSocket connection manager for real-time notifications."""
import asyncio
import json
from datetime import datetime, timezone
from typing import Set
from fastapi import WebSocket


class ConnectionManager:
    """Manages WebSocket connections and broadcasts."""

    def __init__(self):
        self._connections: Set[WebSocket] = set()
        self._lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket) -> None:
        """Accept and track a new WebSocket connection."""
        await websocket.accept()
        async with self._lock:
            self._connections.add(websocket)

        # Send initial connected event
        await self._send_to_client(websocket, {
            "type": "event",
            "event": "connected",
            "data": {"message": "Connected to Dangerous Pi event stream"},
            "timestamp": datetime.now(timezone.utc).isoformat()
        })

    async def disconnect(self, websocket: WebSocket) -> None:
        """Remove a disconnected WebSocket."""
        async with self._lock:
            self._connections.discard(websocket)

    async def broadcast(self, event_type: str, data: dict) -> None:
        """Broadcast an event to all connected clients."""
        message = {
            "type": "event",
            "event": event_type,
            "data": data,
            "timestamp": datetime.now(timezone.utc).isoformat()
        }

        async with self._lock:
            dead_connections: Set[WebSocket] = set()
            for connection in self._connections:
                try:
                    await asyncio.wait_for(
                        self._send_to_client(connection, message),
                        timeout=5.0
                    )
                except asyncio.TimeoutError:
                    dead_connections.add(connection)
                except Exception:
                    dead_connections.add(connection)

            # Remove dead connections
            self._connections -= dead_connections

    async def _send_to_client(self, websocket: WebSocket, message: dict) -> None:
        """Send a message to a specific client."""
        await websocket.send_json(message)

    @property
    def connection_count(self) -> int:
        """Return number of active connections."""
        return len(self._connections)


# Global singleton instance
ws_manager = ConnectionManager()