"""WebSocket API endpoints"""

import asyncio
import logging
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.time import to_iso8601_utc
from app.core.websocket.manager import manager
from app.core.websocket.ue_scene import expand_scene_payload_for_transport, ue_scene_state_store
from app.db.session import async_session_factory

logger = logging.getLogger(__name__)
router = APIRouter()
SUPPORTED_CHANNELS = [
    "gpu_clusters",
    "submarine_cables",
    "ixp_nodes",
    "alerts",
    "dashboard",
    "datasource_tasks",
    "ue_scene",
]


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",
                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                "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",
                            "timestamp": to_iso8601_utc(datetime.now(UTC)),
                            "data": {"action": "pong", "timestamp": to_iso8601_utc(datetime.now(UTC))},
                        }
                    )
                elif data.get("type") == "subscribe":
                    requested_channels = data.get("data", {}).get("channels", [])
                    channels = manager.subscribe(websocket, requested_channels)
                    await websocket.send_json(
                        {
                            "type": "subscription_confirmed",
                            "timestamp": to_iso8601_utc(datetime.now(UTC)),
                            "data": {"action": "subscribe", "channels": channels},
                        }
                    )
                elif data.get("type") == "sync_request":
                    sync_data = data.get("data", {})
                    channel = sync_data.get("channel")
                    if channel == "ue_scene":
                        async with async_session_factory() as session:
                            scene_payloads = await ue_scene_state_store.get_sync_payloads(
                                session,
                                last_sequence=sync_data.get("last_sequence"),
                                reason=sync_data.get("reason"),
                            )
                        for scene_payload in scene_payloads:
                            for transport_payload in expand_scene_payload_for_transport(scene_payload):
                                await websocket.send_json(
                                    {
                                        "type": "data_frame",
                                        "channel": "ue_scene",
                                        "timestamp": to_iso8601_utc(datetime.now(UTC)),
                                        "data": transport_payload,
                                    }
                                )
                    else:
                        await websocket.send_json(
                            {
                                "type": "error",
                                "timestamp": to_iso8601_utc(datetime.now(UTC)),
                                "data": {
                                    "message": f"Unsupported sync channel: {channel}",
                                },
                            }
                        )
                elif data.get("type") == "control_frame":
                    await websocket.send_json(
                        {
                            "type": "control_acknowledged",
                            "timestamp": to_iso8601_utc(datetime.now(UTC)),
                            "data": {
                                "target": data.get("data", {}).get("target"),
                                "command": data.get("data", {}).get("command"),
                                "accepted": True,
                            },
                        }
                    )
                else:
                    await websocket.send_json(
                        {
                            "type": "ack",
                            "timestamp": to_iso8601_utc(datetime.now(UTC)),
                            "data": {"received": True},
                        }
                    )

            except asyncio.TimeoutError:
                await websocket.send_json(
                    {
                        "type": "heartbeat",
                        "timestamp": to_iso8601_utc(datetime.now(UTC)),
                        "data": {"action": "ping"},
                    }
                )

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