"""
设备WebSocket路由
接收Android Agent APP的连接和命令
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from typing import Dict, Optional
import json
import asyncio
import logging
from datetime import datetime

router = APIRouter(tags=["WebSocket"])
logger = logging.getLogger(__name__)


class DeviceConnectionManager:
    """
    设备连接管理器
    管理所有连接的Android设备
    """
    
    def __init__(self):
        # device_id -> WebSocket连接
        self.active_connections: Dict[str, WebSocket] = {}
        # device_id -> 设备信息
        self.device_info: Dict[str, dict] = {}
        # project_id -> [device_id, ...]
        self.project_devices: Dict[str, list] = {}
    
    async def connect(self, device_id: str, websocket: WebSocket):
        """接受新连接"""
        await websocket.accept()
        self.active_connections[device_id] = websocket
        logger.info(f"设备连接: {device_id}")
    
    def disconnect(self, device_id: str):
        """断开连接"""
        if device_id in self.active_connections:
            del self.active_connections[device_id]
        if device_id in self.device_info:
            project_id = self.device_info[device_id].get("project_id")
            if project_id and project_id in self.project_devices:
                if device_id in self.project_devices[project_id]:
                    self.project_devices[project_id].remove(device_id)
            del self.device_info[device_id]
        logger.info(f"设备断开: {device_id}")
    
    def register_device(self, device_id: str, info: dict):
        """注册设备信息"""
        self.device_info[device_id] = {
            **info,
            "connected_at": datetime.now().isoformat(),
            "last_heartbeat": datetime.now().isoformat()
        }
        
        # 添加到项目
        project_id = info.get("project_id")
        if project_id:
            if project_id not in self.project_devices:
                self.project_devices[project_id] = []
            if device_id not in self.project_devices[project_id]:
                self.project_devices[project_id].append(device_id)
        
        logger.info(f"设备注册: {device_id}, 项目: {project_id}")
    
    def update_heartbeat(self, device_id: str):
        """更新心跳时间"""
        if device_id in self.device_info:
            self.device_info[device_id]["last_heartbeat"] = datetime.now().isoformat()
    
    async def send_to_device(self, device_id: str, message: dict) -> bool:
        """发送消息到指定设备"""
        if device_id in self.active_connections:
            try:
                await self.active_connections[device_id].send_json(message)
                return True
            except Exception as e:
                logger.error(f"发送失败: {device_id}, {e}")
                return False
        return False
    
    async def broadcast_to_project(self, project_id: str, message: dict):
        """广播消息到项目下所有设备"""
        if project_id in self.project_devices:
            for device_id in self.project_devices[project_id]:
                await self.send_to_device(device_id, message)
    
    def get_device(self, device_id: str) -> Optional[dict]:
        """获取设备信息"""
        return self.device_info.get(device_id)
    
    def get_project_devices(self, project_id: str) -> list:
        """获取项目下所有设备"""
        device_ids = self.project_devices.get(project_id, [])
        return [self.device_info[did] for did in device_ids if did in self.device_info]
    
    def get_all_devices(self) -> list:
        """获取所有设备"""
        return list(self.device_info.values())
    
    def is_connected(self, device_id: str) -> bool:
        """检查设备是否在线"""
        return device_id in self.active_connections


# 全局连接管理器
manager = DeviceConnectionManager()


@router.websocket("/ws/device/{device_id}")
async def websocket_endpoint(websocket: WebSocket, device_id: str):
    """
    设备WebSocket连接端点
    
    消息类型:
    - register: 设备注册
    - heartbeat: 心跳
    - result: 命令执行结果
    """
    await manager.connect(device_id, websocket)
    
    try:
        while True:
            data = await websocket.receive_json()
            msg_type = data.get("type")
            
            if msg_type == "register":
                # 设备注册
                manager.register_device(device_id, data)
                await websocket.send_json({
                    "type": "registered",
                    "device_id": device_id,
                    "message": "注册成功"
                })
            
            elif msg_type == "heartbeat":
                # 心跳
                manager.update_heartbeat(device_id)
                await websocket.send_json({
                    "type": "pong",
                    "timestamp": datetime.now().timestamp() * 1000
                })
            
            elif msg_type == "result":
                # 命令执行结果
                command_id = data.get("command_id")
                logger.info(f"命令结果: {command_id}, 成功: {data.get('success')}")
                # 这里可以存储结果或通知其他系统
            
            else:
                logger.warning(f"未知消息类型: {msg_type}")
    
    except WebSocketDisconnect:
        manager.disconnect(device_id)
    except Exception as e:
        logger.error(f"WebSocket错误: {device_id}, {e}")
        manager.disconnect(device_id)


# ========== REST API接口 ==========

@router.get("/api/v3/devices")
async def list_devices():
    """获取所有在线设备"""
    return {
        "success": True,
        "devices": manager.get_all_devices(),
        "total": len(manager.get_all_devices())
    }


@router.get("/api/v3/devices/{device_id}")
async def get_device(device_id: str):
    """获取指定设备信息"""
    device = manager.get_device(device_id)
    if device:
        return {
            "success": True,
            "device": device,
            "online": manager.is_connected(device_id)
        }
    return {"success": False, "message": "设备不存在"}


@router.get("/api/v3/projects/{project_id}/devices")
async def get_project_devices(project_id: str):
    """获取项目下所有设备"""
    devices = manager.get_project_devices(project_id)
    return {
        "success": True,
        "project_id": project_id,
        "devices": devices,
        "total": len(devices)
    }


@router.post("/api/v3/devices/{device_id}/execute")
async def execute_command(device_id: str, command: dict):
    """
    向设备发送执行命令
    
    请求体:
    {
        "action": "open_app",
        "params": {"package": "com.tencent.mm"}
    }
    """
    if not manager.is_connected(device_id):
        return {"success": False, "message": "设备离线"}
    
    command_id = f"cmd_{datetime.now().timestamp()}"
    message = {
        "type": "execute",
        "command_id": command_id,
        "action": command.get("action"),
        "params": command.get("params", {})
    }
    
    success = await manager.send_to_device(device_id, message)
    
    return {
        "success": success,
        "command_id": command_id,
        "message": "命令已发送" if success else "发送失败"
    }


@router.post("/api/v3/projects/{project_id}/broadcast")
async def broadcast_to_project(project_id: str, command: dict):
    """向项目下所有设备广播命令"""
    command_id = f"cmd_{datetime.now().timestamp()}"
    message = {
        "type": "execute",
        "command_id": command_id,
        "action": command.get("action"),
        "params": command.get("params", {})
    }
    
    await manager.broadcast_to_project(project_id, message)
    
    devices = manager.get_project_devices(project_id)
    return {
        "success": True,
        "command_id": command_id,
        "devices_count": len(devices),
        "message": f"已广播到 {len(devices)} 台设备"
    }