import json
import asyncio
import time
from fastapi import WebSocket
from services.device_auth_manager import device_auth_manager
from services.connection_manager import connection_manager
from utils.logger import session_logger

async def authenticate_websocket(websocket: WebSocket):
    """
    处理设备认证流程
    
    Args:
        websocket: WebSocket连接
        
    Returns:
        tuple: (认证状态, 设备ID) - (是否认证成功, 设备ID)
    """
    device_id = None
    authenticated = False
    
    auth_timeout = 10  # 10秒认证超时
    auth_start_time = time.time()
    
    while not authenticated and time.time() - auth_start_time < auth_timeout:
        try:
            message = await asyncio.wait_for(
                websocket.receive(), 
                timeout=auth_timeout - (time.time() - auth_start_time)
            )
            
            if message["type"] == "websocket.receive" and "text" in message:
                try:
                    auth_data = json.loads(message["text"])
                    if "device_id" in auth_data and "serial_number" in auth_data:
                        device_id = auth_data["device_id"]
                        serial_number = auth_data["serial_number"]
                        authenticated = await device_auth_manager.authenticate_device(
                            device_id, serial_number
                        )
                        if authenticated:
                            await connection_manager.add_connection(device_id, websocket)
                            session_logger.info(
                                device_id, "auth", f"设备 {device_id} 认证成功"
                            )
                            await websocket.send_text(
                                json.dumps({"status": "authenticated"})
                            )
                        else:
                            session_logger.warning(
                                device_id, "auth", f"设备 {device_id} 认证失败"
                            )
                            await websocket.send_text(
                                json.dumps(
                                    {"status": "error", "message": "Authentication failed"}
                                )
                            )
                except json.JSONDecodeError:
                    session_logger.warning("unknown", "auth", "收到的不是有效的JSON认证消息")
            elif message["type"] == "websocket.receive" and "bytes" in message:
                pass
        except asyncio.TimeoutError:
            session_logger.warning("unknown", "auth", "WebSocket认证超时")
            break
    
    return authenticated, device_id