"""WebSocket endpoint routes."""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from .manager import ws_manager
from ..api.token_store import validate_token
from .. import config

router = APIRouter()


@router.websocket("/events")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time event streaming.

    When AUTH_ENABLED=true, a valid token must be provided as a query
    parameter: ws://host/ws/events?token=<token>
    Tokens are obtained via POST /api/auth/token with Basic Auth.
    Unauthorized connections are closed with code 4401.

    Events include:
    - connected: Initial connection confirmation
    - system_stats: CPU, memory, temperature updates (every 5s)
    - pm3_devices: Device list on connect/disconnect
    - pm3_status: PM3 connection status changes
    - ups_battery: Battery level updates
    - ups_warning: Low battery warning
    - ups_critical: Critical battery level
    - ups_shutdown: Shutdown initiated
    - ups_config_required: UPS needs configuration
    - update_available: New version available
    - update_downloading: Download progress
    """
    # Validate auth token when authentication is enabled
    if config.AUTH_ENABLED:
        token = websocket.query_params.get("token")
        if not token or not validate_token(token):
            # Must accept before sending close code so client receives it
            await websocket.accept()
            await websocket.close(code=4401, reason="Authentication required")
            return

    await ws_manager.connect(websocket)
    try:
        while True:
            # Keep connection alive by waiting for messages or disconnect
            # Client doesn't send data, but we need to detect disconnection
            try:
                await websocket.receive_text()
            except WebSocketDisconnect:
                break
    finally:
        await ws_manager.disconnect(websocket)