Aniworld/src/server/api/websocket.py

"""WebSocket API endpoints for real-time communication.

This module provides WebSocket endpoints for clients to connect and receive
real-time updates about downloads, queue status, and system events.
"""
from __future__ import annotations

import uuid
from typing import Optional

import structlog
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status
from fastapi.responses import JSONResponse

from src.server.models.websocket import (
    ClientMessage,
    RoomSubscriptionRequest,
    WebSocketMessageType,
)
from src.server.services.websocket_service import (
    WebSocketService,
    get_websocket_service,
)
from src.server.utils.dependencies import get_current_user_optional

logger = structlog.get_logger(__name__)

router = APIRouter(prefix="/ws", tags=["websocket"])


@router.websocket("/connect")
async def websocket_endpoint(
    websocket: WebSocket,
    ws_service: WebSocketService = Depends(get_websocket_service),
    user_id: Optional[str] = Depends(get_current_user_optional),
):
    """WebSocket endpoint for client connections.
    
    Clients connect to this endpoint to receive real-time updates.
    The connection is maintained until the client disconnects or
    an error occurs.
    
    Message flow:
    1. Client connects
    2. Server sends "connected" message
    3. Client can send subscription requests (join/leave rooms)
    4. Server broadcasts updates to subscribed rooms
    5. Client disconnects
    
    Example client subscription:
    ```json
    {
        "action": "join",
        "room": "downloads"
    }
    ```
    
    Server message format:
    ```json
    {
        "type": "download_progress",
        "timestamp": "2025-10-17T10:30:00.000Z",
        "data": {
            "download_id": "abc123",
            "percent": 45.2,
            "speed_mbps": 2.5,
            "eta_seconds": 180
        }
    }
    ```
    """
    connection_id = str(uuid.uuid4())
    
    try:
        # Accept connection and register with service
        await ws_service.connect(websocket, connection_id, user_id=user_id)
        
        # Send connection confirmation
        await ws_service.manager.send_personal_message(
            {
                "type": WebSocketMessageType.CONNECTED,
                "data": {
                    "connection_id": connection_id,
                    "message": "Connected to Aniworld WebSocket",
                },
            },
            connection_id,
        )
        
        logger.info(
            "WebSocket client connected",
            connection_id=connection_id,
            user_id=user_id,
        )
        
        # Handle incoming messages
        while True:
            try:
                # Receive message from client
                data = await websocket.receive_json()
                
                # Parse client message
                try:
                    client_msg = ClientMessage(**data)
                except Exception as e:
                    logger.warning(
                        "Invalid client message format",
                        connection_id=connection_id,
                        error=str(e),
                    )
                    await ws_service.send_error(
                        connection_id,
                        "Invalid message format",
                        "INVALID_MESSAGE",
                    )
                    continue
                
                # Handle room subscription requests
                if client_msg.action in ["join", "leave"]:
                    try:
                        room_req = RoomSubscriptionRequest(
                            action=client_msg.action,
                            room=client_msg.data.get("room", ""),
                        )
                        
                        if room_req.action == "join":
                            await ws_service.manager.join_room(
                                connection_id, room_req.room
                            )
                            await ws_service.manager.send_personal_message(
                                {
                                    "type": WebSocketMessageType.SYSTEM_INFO,
                                    "data": {
                                        "message": (
                                            f"Joined room: {room_req.room}"
                                        )
                                    },
                                },
                                connection_id,
                            )
                        elif room_req.action == "leave":
                            await ws_service.manager.leave_room(
                                connection_id, room_req.room
                            )
                            await ws_service.manager.send_personal_message(
                                {
                                    "type": WebSocketMessageType.SYSTEM_INFO,
                                    "data": {
                                        "message": (
                                            f"Left room: {room_req.room}"
                                        )
                                    },
                                },
                                connection_id,
                            )
                            
                    except Exception as e:
                        logger.warning(
                            "Invalid room subscription request",
                            connection_id=connection_id,
                            error=str(e),
                        )
                        await ws_service.send_error(
                            connection_id,
                            "Invalid room subscription",
                            "INVALID_SUBSCRIPTION",
                        )
                
                # Handle ping/pong for keepalive
                elif client_msg.action == "ping":
                    await ws_service.manager.send_personal_message(
                        {"type": WebSocketMessageType.PONG, "data": {}},
                        connection_id,
                    )
                
                else:
                    logger.debug(
                        "Unknown action from client",
                        connection_id=connection_id,
                        action=client_msg.action,
                    )
                    await ws_service.send_error(
                        connection_id,
                        f"Unknown action: {client_msg.action}",
                        "UNKNOWN_ACTION",
                    )
                    
            except WebSocketDisconnect:
                logger.info(
                    "WebSocket client disconnected",
                    connection_id=connection_id,
                )
                break
            except Exception as e:
                logger.error(
                    "Error handling WebSocket message",
                    connection_id=connection_id,
                    error=str(e),
                )
                await ws_service.send_error(
                    connection_id,
                    "Internal server error",
                    "SERVER_ERROR",
                )
                
    except Exception as e:
        logger.error(
            "WebSocket connection error",
            connection_id=connection_id,
            error=str(e),
        )
    finally:
        # Cleanup connection
        await ws_service.disconnect(connection_id)
        logger.info("WebSocket connection closed", connection_id=connection_id)


@router.get("/status")
async def websocket_status(
    ws_service: WebSocketService = Depends(get_websocket_service),
):
    """Get WebSocket service status and statistics.
    
    Returns information about active connections and rooms.
    Useful for monitoring and debugging.
    """
    connection_count = await ws_service.manager.get_connection_count()
    
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={
            "status": "operational",
            "active_connections": connection_count,
            "supported_message_types": [t.value for t in WebSocketMessageType],
        },
    )