"""WebSocket API endpoints"""

import asyncio
from datetime import UTC, datetime
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from jose import jwt, JWTError

from app.core.config import settings
from app.core.logging import get_logger
from app.core.time import to_iso8601_utc
from app.core.websocket.manager import manager

logger = get_logger(__name__, service="api")
router = APIRouter()


async def authenticate_token(token: str) -> Optional[dict]:
    """Authenticate WebSocket connection via token"""
    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        if payload.get("type") != "access":
            logger.warning_event(
                "WebSocket auth failed: wrong token type",
                event="auth.websocket.invalid_token_type",
            )
            return None
        return payload
    except JWTError as e:
        logger.warning_event(
            "WebSocket auth failed",
            event="auth.websocket.decode_failed",
            context={"error": str(e)},
        )
        return None


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str | None = Query(None),
):
    """WebSocket endpoint for real-time data"""
    logger.info_event(
        "WebSocket connection attempt",
        event="auth.websocket.connection_attempt",
        context={"token_preview": f"{token[:8]}..." if token else "anonymous"},
    )
    payload = await authenticate_token(token) if token else None
    if token and payload is None:
        logger.warning_event(
            "WebSocket authentication failed, closing connection",
            event="auth.websocket.connection_rejected",
        )
        await websocket.close(code=4001)
        return

    is_anonymous = payload is None
    user_id = str(payload.get("sub")) if payload else f"anonymous:{id(websocket)}"
    supported_channels = ["vessels", "earth_news"] if is_anonymous else [
        "gpu_clusters",
        "submarine_cables",
        "ixp_nodes",
        "alerts",
        "dashboard",
        "datasource_tasks",
        "vessels",
        "earth_news",
    ]
    await manager.connect(websocket, user_id)

    try:
        await websocket.send_json(
            {
                "type": "connection_established",
                "data": {
                    "connection_id": f"conn_{user_id}",
                    "server_version": settings.VERSION,
                    "heartbeat_interval": 30,
                    "supported_channels": supported_channels,
                },
            }
        )

        while True:
            try:
                data = await asyncio.wait_for(websocket.receive_json(), timeout=30)

                if data.get("type") == "heartbeat":
                    await websocket.send_json(
                        {
                            "type": "heartbeat",
                            "data": {"action": "pong", "timestamp": to_iso8601_utc(datetime.now(UTC))},
                        }
                    )
                elif data.get("type") == "subscribe":
                    payload_data = data.get("data", {})
                    if not isinstance(payload_data, dict):
                        payload_data = {}
                    channels = payload_data.get("channels", [])
                    if isinstance(channels, str):
                        channels = [channels]
                    elif not isinstance(channels, list):
                        channels = []
                    channel = payload_data.get("channel")
                    if channel and channel not in channels:
                        channels = [*channels, channel]
                    if is_anonymous:
                        channels = [channel for channel in channels if channel in supported_channels]
                    vessel_subscription = None
                    if "vessels" in channels and "bbox" in payload_data:
                        try:
                            vessel_subscription = manager.subscribe_vessels(websocket, payload_data)
                        except ValueError as exc:
                            await websocket.send_json(
                                {
                                    "type": "subscription_error",
                                    "data": {"channel": "vessels", "detail": str(exc)},
                                }
                            )
                            continue
                        channels = [channel for channel in channels if channel != "vessels"]
                    manager.subscribe(websocket, channels)
                    await websocket.send_json(
                        {
                            "type": "subscription_confirmed",
                            "data": {
                                "action": "subscribe",
                                "channels": [
                                    *channels,
                                    *(["vessels"] if vessel_subscription else []),
                                ],
                                "vessels": vessel_subscription,
                            },
                        }
                    )
                elif data.get("type") == "unsubscribe":
                    channels = data.get("data", {}).get("channels", [])
                    manager.unsubscribe(websocket, channels)
                    await websocket.send_json(
                        {
                            "type": "subscription_confirmed",
                            "data": {"action": "unsubscribe", "channels": channels},
                        }
                    )
                elif data.get("type") == "control_frame":
                    await websocket.send_json(
                        {"type": "control_acknowledged", "data": {"received": True}}
                    )
                else:
                    await websocket.send_json({"type": "ack", "data": {"received": True}})

            except asyncio.TimeoutError:
                await websocket.send_json({"type": "heartbeat", "data": {"action": "ping"}})

    except WebSocketDisconnect:
        pass
    finally:
        manager.disconnect(websocket, user_id)