"""
WebSocket endpoint.

URL: /ws
     /ws?server_id=1
     /ws?server_id=1&server_id=2

Authentication: JWT passed as a query parameter `token` because
browser WebSocket API does not support custom headers.
If the token is missing or invalid, the connection is closed with code 4001.

After authentication, the client receives:
- A "connected" welcome message with the list of subscribed server IDs
- All events for subscribed servers pushed by BroadcastThread
"""
from __future__ import annotations

import logging
from typing import Optional

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from core.auth.utils import decode_access_token

logger = logging.getLogger(__name__)

router = APIRouter(tags=["websocket"])


@router.websocket("/ws")
async def websocket_endpoint(
    ws: WebSocket,
    token: Optional[str] = Query(default=None),
    server_id: Optional[list[int]] = Query(default=None),
) -> None:
    """
    WebSocket endpoint for real-time server events.

    Query parameters:
        token: JWT access token (required)
        server_id: One or more server IDs to subscribe to (optional, default=all)
    """
    # Authenticate before accepting
    if not token:
        await ws.close(code=4001, reason="Missing token")
        return

    try:
        user = decode_access_token(token)
    except Exception as exc:
        logger.warning("WebSocket: token decode failed: %s", exc)
        user = None

    if user is None:
        await ws.close(code=4001, reason="Invalid or expired token")
        return

    # Get WebSocketManager from app state
    ws_manager = ws.app.state.ws_manager

    await ws_manager.connect(ws, server_ids=server_id)
    logger.info(
        "WebSocket: user '%s' connected, subscribed to servers=%s",
        user.get("sub"),
        server_id,
    )

    try:
        # Send welcome message
        await ws_manager.send_to_connection(ws, {
            "type": "connected",
            "data": {
                "user": user.get("sub"),
                "subscriptions": server_id or "all",
            },
        })

        # Keep connection alive — wait for client to disconnect
        while True:
            data = await ws.receive_text()

    except WebSocketDisconnect:
        logger.info(
            "WebSocket: user '%s' disconnected",
            user.get("sub"),
        )
    except Exception as exc:
        logger.error("WebSocket: unexpected error: %s", exc)
    finally:
        await ws_manager.disconnect(ws)