"""WebSocket API endpoints"""

import asyncio
import json
import logging
from datetime import 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.websocket.manager import manager

logger = logging.getLogger(__name__)
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(f"WebSocket auth failed: wrong token type")
            return None
        return payload
    except JWTError as e:
        logger.warning(f"WebSocket auth failed: {e}")
        return None


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Query(...),
):
    """WebSocket endpoint for real-time data"""
    logger.info(f"WebSocket connection attempt with token: {token[:20]}...")
    payload = await authenticate_token(token)
    if payload is None:
        logger.warning("WebSocket authentication failed, closing connection")
        await websocket.close(code=4001)
        return

    user_id = str(payload.get("sub"))
    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": [
                        "gpu_clusters",
                        "submarine_cables",
                        "ixp_nodes",
                        "alerts",
                        "dashboard",
                    ],
                },
            }
        )

        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": datetime.utcnow().isoformat()},
                        }
                    )
                elif data.get("type") == "subscribe":
                    channels = data.get("data", {}).get("channels", [])
                    await websocket.send_json(
                        {
                            "type": "subscription_confirmed",
                            "data": {"action": "subscribe", "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)