from fastapi import WebSocket, WebSocketDisconnect
from utils.logger import session_logger
from services.connection_manager import connection_manager
from services.task_manager import task_manager
from handlers.websocket_auth_handler import authenticate_websocket
from handlers.websocket_message_handler import handle_websocket_messages
from handlers.session_cleanup_handler import cleanup_device_sessions

async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket连接的主入口点
    
    Args:
        websocket: WebSocket连接
    """
    await websocket.accept()
    session_logger.info(
        "unknown", "connection", f"WebSocket 连接已建立: {websocket.client}"
    )
    device_id = None
    
    try:
        authenticated, device_id = await authenticate_websocket(websocket)
        
        if not authenticated:
            session_logger.warning(
                device_id or "unknown", "auth", "未认证的设备尝试连接，断开连接"
            )
            await websocket.send_text('{"status": "error", "message": "Not authenticated"}')
            return
        
        await handle_websocket_messages(websocket, device_id)
        
    except WebSocketDisconnect:
        session_logger.info(
            device_id or "unknown",
            "connection",
            f"WebSocket 连接已关闭: {websocket.client}",
        )
    except Exception as e:
        session_logger.error(
            device_id or "unknown", "error", f"处理 WebSocket 数据时出错: {str(e)}"
        )
    finally:
        if device_id:
            await connection_manager.remove_connection(device_id)
            await cleanup_device_sessions(device_id)
            # 清理设备相关的所有异步任务
            await task_manager.cancel_device_tasks(device_id)
        
        session_logger.info(
            device_id or "unknown",
            "connection",
            "服务端保持WebSocket连接开放，由客户端负责断开连接",
        )