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

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query

from dashboard.dependencies import get_ws_manager, get_device_manager
from dashboard.services.websocket_manager import WebSocketManager
from dashboard.services.device_manager import DeviceManager

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


@router.websocket("")
async def websocket_endpoint(
    websocket: WebSocket,
    client_id: str | None = Query(None),
    ws_manager: WebSocketManager = Depends(get_ws_manager),
    device_manager: DeviceManager = Depends(get_device_manager),
):
    """WebSocket endpoint for real-time updates.

    This endpoint provides real-time updates for:
    - Device connection/disconnection
    - Task execution progress
    - Screenshot updates
    - Task completion

    Query parameters:
        client_id: Optional client ID (auto-generated if not provided)

    Message types (client -> server):
        - {"type": "subscribe", "device_id": "device_id"} - Subscribe to device updates
        - {"type": "unsubscribe", "device_id": "device_id"} - Unsubscribe from device updates
        - {"type": "ping"} - Ping server

    Message types (server -> client):
        - {"type": "device_update", "data": {...}} - Device status update
        - {"type": "task_started", "data": {...}} - Task started
        - {"type": "task_step", "data": {...}} - Task step update
        - {"type": "task_completed", "data": {...}} - Task completed
        - {"type": "task_failed", "data": {...}} - Task failed
        - {"type": "task_stopped", "data": {...}} - Task stopped
        - {"type": "screenshot", "data": {...}} - Screenshot update
        - {"type": "error", "data": {...}} - Error occurred
        - {"type": "pong"} - Pong response
    """
    # Accept connection
    client_id = await ws_manager.connect(websocket, client_id)

    try:
        # Send initial devices
        devices = await device_manager.refresh_devices()
        for device in devices:
            await websocket.send_json({
                "type": "device_update",
                "data": {
                    "device_id": device.device_id,
                    "status": device.status.value,
                    "device_type": device.device_type.value,
                    "model": device.model,
                    "android_version": device.android_version,
                    "current_app": device.current_app,
                    "is_connected": device.is_connected,
                }
            })

        # Message loop
        while True:
            data = await websocket.receive_json()
            msg_type = data.get("type")

            if msg_type == "subscribe":
                # Subscribe to device updates
                device_id = data.get("device_id", "*")
                ws_manager.subscribe_to_device(client_id, device_id)

            elif msg_type == "unsubscribe":
                # Unsubscribe from device updates
                device_id = data.get("device_id")
                if device_id:
                    ws_manager.unsubscribe_from_device(client_id, device_id)

            elif msg_type == "ping":
                # Respond to ping
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        ws_manager.disconnect(client_id)
    except Exception:
        ws_manager.disconnect(client_id)


@router.websocket("/device/{device_id}")
async def device_websocket(
    device_id: str,
    websocket: WebSocket,
    ws_manager: WebSocketManager = Depends(get_ws_manager),
):
    """Device-specific WebSocket endpoint for real-time updates.

    This endpoint provides real-time updates for a specific device.
    Automatically subscribes to the device's updates.

    Args:
        device_id: Device identifier
    """
    # Accept connection and auto-subscribe to device
    client_id = await ws_manager.connect(websocket)
    ws_manager.subscribe_to_device(client_id, device_id)

    try:
        # Keep connection alive
        while True:
            data = await websocket.receive_json()

            # Handle client messages
            if data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        ws_manager.disconnect(client_id)
    except Exception:
        ws_manager.disconnect(client_id)