"""
工作手机SDK v3.0 - WebSocket Hub
管理所有设备的WebSocket连接
"""

from fastapi import WebSocket
from typing import Dict, Optional, Callable, Any
import asyncio
import json
import uuid
import logging
from datetime import datetime
from config import settings

logger = logging.getLogger(__name__)

# 延迟导入避免循环依赖，在 handle_message 内使用
def _get_device_manager():
    from services.device_manager import device_manager
    return device_manager


async def _record_device_event(device_id: str, event_type: str, payload: dict) -> None:
    """把真实 WS 设备事件写入 Hook 事件总线，供存客宝/超管/AI 数字员工实时订阅。"""
    try:
        from services.hook_module_service import hook_module_service

        await hook_module_service.add_event({
            "source": "ws_agent",
            "platform": "android",
            "device_id": device_id,
            "event_type": event_type,
            "payload": payload,
        })
    except Exception as e:
        logger.debug(f"设备事件落总线失败 [{device_id}] {event_type}: {e}")


class WebSocketHub:
    """WebSocket连接管理中心"""
    
    def __init__(self):
        # 设备连接: device_id -> WebSocket
        self.connections: Dict[str, WebSocket] = {}
        
        # 设备信息: device_id -> device_info
        self.device_info: Dict[str, dict] = {}
        
        # 项目设备映射: project_id -> [device_id, ...]
        self.project_devices: Dict[str, list] = {}
        
        # 命令响应: command_id -> asyncio.Future
        self.pending_commands: Dict[str, asyncio.Future] = {}
        # 设备在途命令计数（sweep 时跳过误踢）
        self.pending_device_commands: Dict[str, int] = {}
        
        # 命令结果回调: command_id -> result
        self.command_results: Dict[str, dict] = {}
        
        # 消息处理器
        self.message_handlers: Dict[str, Callable] = {}
        # 设备心跳配置: device_id -> interval_seconds
        self.heartbeat_config: Dict[str, int] = {}
    
    async def connect(self, websocket: WebSocket, device_id: str):
        """设备连接"""
        await websocket.accept()
        self.connections[device_id] = websocket
        logger.info(f"✅ 设备连接: {device_id}")
    
    async def disconnect(self, device_id: str):
        """设备断开"""
        old_info = dict(self.device_info.get(device_id) or {})
        # 从项目中移除
        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]
        
        if device_id in self.connections:
            del self.connections[device_id]
        
        logger.info(f"❌ 设备断开: {device_id}")
        await _record_device_event(device_id, "device_disconnected", {
            "project_id": old_info.get("project_id"),
            "model": old_info.get("model"),
            "last_heartbeat": old_info.get("last_heartbeat"),
        })
    
    def is_online(self, device_id: str) -> bool:
        """检查设备是否在线"""
        return device_id in self.connections
    
    async def handle_message(self, device_id: str, data: dict):
        """处理设备消息"""
        msg_type = data.get("type")
        
        if msg_type == "register":
            # 设备注册（支持APP直接发送的格式）
            # APP发送: {type: "register", device_id: xxx, project_id: xxx, ...}
            # Agent发送: {type: "register", data: {...}}
            
            device_data = data.get("data") or data
            project_id = device_data.get("project_id", "")
            heartbeat_interval = int(device_data.get("heartbeat_interval_seconds") or settings.WS_HEARTBEAT_INTERVAL)
            heartbeat_interval = max(5, min(120, heartbeat_interval))
            
            now_iso = datetime.now().isoformat()
            self.device_info[device_id] = {
                **device_data,
                "device_id": device_id,
                "project_id": project_id,
                "status": "online",
                "connected_at": now_iso,
                "last_heartbeat": now_iso,
                "heartbeat_interval_seconds": heartbeat_interval,
            }
            self.heartbeat_config[device_id] = heartbeat_interval

            # 添加到项目设备映射
            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}")
            else:
                logger.info(f"📱 设备注册: {device_id} (无项目绑定)")
            
            # 先发确认（不被 MongoDB 阻塞）
            await self.send_to_device(device_id, {
                "type": "registered", 
                "success": True,
                "device_id": device_id,
                "project_id": project_id,
                "heartbeat_interval_seconds": heartbeat_interval,
            })

            # 注册落库 fire-and-forget
            _reg_data = {
                "device_id": device_id, "project_id": project_id,
                "status": "online", "connected_at": now_iso, "last_heartbeat": now_iso,
                **{k: v for k, v in device_data.items() if k not in ("device_id",)}
            }
            async def _reg_persist():
                try:
                    dm = _get_device_manager()
                    await dm.register_device(_reg_data)
                except Exception as e:
                    logger.warning(f"设备落库失败: {e}")
            asyncio.create_task(_reg_persist())
            asyncio.create_task(_record_device_event(device_id, "device_registered", {
                "project_id": device_data.get("project_id"),
                "model": device_data.get("model"),
                "agent_version": device_data.get("agent_version"),
                "capabilities": device_data.get("capabilities") or [],
                "frida_available": device_data.get("frida_available"),
                "ai_brain_enabled": device_data.get("ai_brain_enabled"),
            }))
        
        elif msg_type == "heartbeat":
            now_iso = datetime.now().isoformat()
            if device_id in self.device_info:
                self.device_info[device_id]["last_heartbeat"] = now_iso
                hb_status = data.get("status") or {}
                if hb_status:
                    self.device_info[device_id]["quick_status"] = hb_status
                    _sync_fields = [
                        "wechat_running", "wechat_foreground", "network_type",
                        "screen_on", "current_app", "battery_level",
                        "battery_charging", "memory_usage_pct", "storage_free_mb",
                        "uptime_sec", "cpu_usage_pct",
                        "connect_stage",  # BIND-07 寻服阶段实时同步
                    ]
                    for _f in _sync_fields:
                        if hb_status.get(_f) is not None:
                            self.device_info[device_id][_f] = hb_status[_f]
                    # 兼容旧 Agent（≤3.x/4.0 心跳用 "battery" 键）：统一落 battery_level
                    if hb_status.get("battery") is not None and hb_status.get("battery_level") is None:
                        self.device_info[device_id]["battery_level"] = hb_status["battery"]
                self.device_info[device_id]["heartbeat_count"] = (
                    self.device_info[device_id].get("heartbeat_count", 0) + 1
                )
                guard_data = data.get("guard")
                if guard_data:
                    self.device_info[device_id]["guard"] = guard_data
                    self._store_guard_events(device_id, guard_data)
            # 先回 pong（不被 MongoDB 阻塞）；双 type 兼容旧 Agent
            pong_payload: dict = {
                "type": "pong",
                "server_time": now_iso,
                "heartbeat_ack": True,
            }
            from services.ai_heartbeat import ai_heartbeat
            q = ai_heartbeat.get_task_queue_status(device_id)
            if q.get("pending_count", 0) > 0:
                pong_payload["pending_tasks"] = q["pending_count"]
            await self.send_to_device(device_id, pong_payload)

            # 心跳落库改为 fire-and-forget（MongoDB 挂了不影响 ACK）
            async def _hb_persist():
                try:
                    dm = _get_device_manager()
                    await dm.update_heartbeat(device_id)
                except Exception as e:
                    logger.debug(f"心跳落库: {e}")
            asyncio.create_task(_hb_persist())
        
        elif msg_type in ("response", "result"):
            # 命令响应
            # - Python Agent: {type:"response", code, message, data}
            # - Android App: {type:"result", success, message}
            command_id = data.get("command_id")
            if not command_id:
                return
            if msg_type == "result":
                success = bool(data.get("success"))
                msg = data.get("message") or ("success" if success else "failed")
                raw_data = data.get("data") or {}
                if not raw_data and msg:
                    raw_data = {"output": msg}
                normalized = {
                    "type": "response",
                    "command_id": command_id,
                    "device_id": device_id,
                    "code": 200 if success else 500,
                    "message": msg,
                    "data": raw_data,
                    "timestamp": data.get("timestamp"),
                }
                data = normalized
            if command_id in self.pending_commands:
                future = self.pending_commands.pop(command_id)
                cnt = self.pending_device_commands.get(device_id, 0) - 1
                if cnt <= 0:
                    self.pending_device_commands.pop(device_id, None)
                else:
                    self.pending_device_commands[device_id] = cnt
                if device_id in self.device_info:
                    self.device_info[device_id].pop("command_grace_seconds", None)
                if not future.done():
                    future.set_result(data)

        elif msg_type == "status_report":
            # 详细状态上报（设备端定期上报）
            status_payload = data.get("data", data)
            if device_id in self.device_info:
                self.device_info[device_id]["last_status"] = status_payload
                self.device_info[device_id]["last_heartbeat"] = datetime.now().isoformat()
                # 与心跳同源：把电量/网络等关键状态提升到顶层 device_info，
                # 让 fleet/status、设备详情等聚合端「各端统一」读到真实电量（否则只在 last_status 里）。
                if isinstance(status_payload, dict):
                    _sync_fields = [
                        "wechat_running", "wechat_foreground", "network_type",
                        "screen_on", "current_app", "battery_level",
                        "battery_charging", "battery_status", "charging",
                        "memory_usage_pct", "storage_free_mb",
                        "uptime_sec", "cpu_usage_pct", "connect_stage",
                    ]
                    for _f in _sync_fields:
                        if status_payload.get(_f) is not None:
                            self.device_info[device_id][_f] = status_payload[_f]
                    # 旧 Agent 用 "battery" 键：统一落 battery_level
                    if status_payload.get("battery") is not None and status_payload.get("battery_level") is None:
                        self.device_info[device_id]["battery_level"] = status_payload["battery"]
            logger.debug(f"设备状态上报: {device_id}")
            asyncio.create_task(_record_device_event(device_id, "device_status_report", status_payload))

        elif msg_type == "event":
            event_name = data.get("event", "unknown")
            event_data = data.get("data", {})
            logger.info(f"设备事件 [{device_id}] {event_name}: {event_data}")
            asyncio.create_task(_record_device_event(device_id, event_name, event_data))

            if event_name == "ai_brain_acted":
                await self._handle_ai_brain_acted(device_id, event_data)
            elif event_name == "offline_buffer_upload":
                await self._handle_offline_buffer(device_id, event_data)
            elif event_name == "agent_started":
                if device_id in self.device_info:
                    self.device_info[device_id]["frida_available"] = event_data.get("frida_available", False)
                    self.device_info[device_id]["ai_brain_enabled"] = event_data.get("ai_brain_enabled", False)

            handler = self.message_handlers.get("event")
            if handler:
                await handler(device_id, data)

        elif msg_type == "device_request":
            # 设备端请求服务端执行操作（拉配置、落库、通知业务等）
            request_id = data.get("request_id", "")
            action = data.get("action", "")
            params = data.get("params", {})
            logger.info(f"设备请求 [{device_id}] {action} request_id={request_id}")
            ack = await self._handle_device_request(device_id, request_id, action, params)
            if ack:
                await self.send_to_device(device_id, ack)
            handler = self.message_handlers.get("device_request")
            if handler:
                await handler(device_id, data)
        
        else:
            # 自定义处理器
            handler = self.message_handlers.get(msg_type)
            if handler:
                await handler(device_id, data)
    
    async def _handle_device_request(
        self, device_id: str, request_id: str, action: str, params: dict
    ) -> Optional[dict]:
        """处理设备端 device_request，返回要下发给设备的 ack（可选）"""
        if action == "get_config":
            interval = self.heartbeat_config.get(device_id, settings.WS_HEARTBEAT_INTERVAL)
            config_data = {
                "heartbeat_interval": interval,
                "heartbeat_interval_seconds": interval,
                "ws_timeout_seconds": settings.WS_TIMEOUT,
            }
            ai_cfg = getattr(settings, "AI_BRAIN_CONFIG", None)
            if ai_cfg:
                config_data["ai_brain"] = ai_cfg
            return {
                "type": "device_request_ack",
                "request_id": request_id,
                "success": True,
                "data": config_data,
            }
        if action == "log_result":
            # 仅记录日志，可选落库由 message_handlers 扩展
            logger.info(f"设备请求落库 [{device_id}] {params}")
            return {
                "type": "device_request_ack",
                "request_id": request_id,
                "success": True,
            }
        logger.warning(f"未知 device_request action: {action}")
        return {
            "type": "device_request_ack",
            "request_id": request_id,
            "success": False,
            "error": f"unknown action: {action}",
        }

    # ================================================================
    # AI Brain 服务端支持
    # ================================================================

    async def _handle_ai_brain_acted(self, device_id: str, event_data: dict):
        """处理 AI Brain 自主执行事件"""
        results_count = event_data.get("results_count", 0)
        reason = event_data.get("reason", "")
        logger.info(f"🧠 AI Brain [{device_id}] 执行了 {results_count} 个操作: {reason}")
        if device_id in self.device_info:
            ai_stats = self.device_info[device_id].setdefault("ai_brain_stats", {
                "total_acts": 0, "last_act": None,
            })
            ai_stats["total_acts"] += results_count
            ai_stats["last_act"] = datetime.now().isoformat()
            ai_stats["last_reason"] = reason

    async def _handle_offline_buffer(self, device_id: str, event_data: dict):
        """处理设备离线期间缓冲的执行结果"""
        count = event_data.get("count", 0)
        results = event_data.get("results", [])
        logger.info(f"📥 离线缓冲上传 [{device_id}]: {count} 条结果")
        if device_id in self.device_info:
            self.device_info[device_id].setdefault("offline_uploads", []).append({
                "uploaded_at": datetime.now().isoformat(),
                "count": count,
            })
        try:
            dm = _get_device_manager()
            await dm.log_offline_results(device_id, results)
        except Exception as e:
            logger.debug(f"离线缓冲落库: {e}")

    async def push_ai_task(self, device_id: str, instruction: str, priority: int = 5) -> bool:
        """向设备推送 AI 任务"""
        return await self.send_to_device(device_id, {
            "type": "ai_task",
            "data": {"instruction": instruction, "priority": priority},
        })

    async def push_standing_order(self, device_id: str, order: str) -> bool:
        """向设备推送常驻指令"""
        return await self.send_to_device(device_id, {
            "type": "standing_order",
            "data": {"order": order},
        })

    def get_ai_brain_status(self, device_id: str) -> Optional[dict]:
        """获取设备 AI Brain 状态"""
        info = self.device_info.get(device_id)
        if not info:
            return None
        return {
            "device_id": device_id,
            "ai_brain_enabled": info.get("ai_brain_enabled", False),
            "frida_available": info.get("frida_available", False),
            "ai_brain_stats": info.get("ai_brain_stats"),
            "offline_uploads": info.get("offline_uploads", [])[-5:],
            "last_status": info.get("last_status", {}),
        }

    def _store_guard_events(self, device_id: str, guard_data: dict):
        """将心跳中的守护数据存入事件存储（供 API 查询）"""
        try:
            from routers.devices import _guard_event_store, _MAX_EVENTS_PER_DEVICE
            import time
            if device_id not in _guard_event_store:
                _guard_event_store[device_id] = []
            store = _guard_event_store[device_id]

            popups = guard_data.get("popups")
            if popups and popups.get("count", 0) > 0:
                store.append({
                    "type": "popup_dismissed",
                    "dismissed": popups.get("dismissed", []),
                    "count": popups["count"],
                    "server_received_at": time.time(),
                })

            full_cycle = guard_data.get("full_cycle")
            if full_cycle:
                store.append({
                    "type": "guard_cycle",
                    "all_ok": full_cycle.get("all_ok", True),
                    "u2_alive": full_cycle.get("u2", {}).get("alive"),
                    "network_ok": full_cycle.get("network", {}).get("network"),
                    "screen_on": full_cycle.get("screen", {}).get("screen_on"),
                    "server_received_at": time.time(),
                })
                if not full_cycle.get("all_ok"):
                    store.append({
                        "type": "guard_alert",
                        "issues": full_cycle,
                        "server_received_at": time.time(),
                    })

            if len(store) > _MAX_EVENTS_PER_DEVICE:
                _guard_event_store[device_id] = store[-_MAX_EVENTS_PER_DEVICE:]
        except Exception as e:
            logger.debug(f"守护事件存储异常: {e}")

    async def send_to_device(self, device_id: str, data: dict) -> bool:
        """发送消息到设备（失败重试，避免单次抖动即踢下线）"""
        ws = self.connections.get(device_id)
        if not ws:
            return False

        last_err = None
        for attempt in range(3):
            try:
                await ws.send_json(data)
                return True
            except Exception as e:
                last_err = e
                logger.warning(f"发送失败 [{device_id}] attempt={attempt + 1}/3: {e}")
                if attempt < 2:
                    await asyncio.sleep(0.4 * (attempt + 1))

        logger.error(f"发送失败 [{device_id}] 已重试 3 次: {last_err}")
        await self.disconnect(device_id)
        return False

    def get_heartbeat_status(self, device_id: str) -> Optional[dict]:
        info = self.device_info.get(device_id)
        if not info:
            return None
        last = info.get("last_heartbeat")
        if not last:
            return None
        try:
            dt = datetime.fromisoformat(last)
            age = max(0, int((datetime.now() - dt).total_seconds()))
        except Exception:
            age = -1
        interval = int(info.get("heartbeat_interval_seconds") or self.heartbeat_config.get(device_id, settings.WS_HEARTBEAT_INTERVAL))
        return {
            "device_id": device_id,
            "online": self.is_online(device_id),
            "last_heartbeat": last,
            "heartbeat_interval_seconds": interval,
            "heartbeat_age_seconds": age,
            "stale": age >= (interval * 3) if age >= 0 else True,
        }

    async def set_heartbeat_interval(self, device_id: str, interval_seconds: int) -> bool:
        if interval_seconds < 5 or interval_seconds > 120:
            return False
        self.heartbeat_config[device_id] = interval_seconds
        if device_id in self.device_info:
            self.device_info[device_id]["heartbeat_interval_seconds"] = interval_seconds
        if self.is_online(device_id):
            await self.send_to_device(device_id, {
                "type": "config",
                "heartbeat_interval_seconds": interval_seconds,
            })
        return True

    async def sweep_stale_devices(self, timeout_seconds: int):
        # 主动清理超时连接，避免设备长时间假在线
        to_drop = []
        now = datetime.now()
        for device_id, info in list(self.device_info.items()):
            if self.pending_device_commands.get(device_id, 0) > 0:
                logger.debug(f"跳过超时下线（设备有在途命令）: {device_id}")
                continue
            last = info.get("last_heartbeat")
            if not last:
                continue
            try:
                age = (now - datetime.fromisoformat(last)).total_seconds()
            except Exception:
                age = timeout_seconds + 1
            grace = int(info.get("command_grace_seconds") or 0)
            effective_timeout = timeout_seconds + max(0, grace)
            if age > effective_timeout:
                to_drop.append(device_id)
        for device_id in to_drop:
            logger.warning(f"⏱️ 心跳超时下线: {device_id}")
            await self.disconnect(device_id)
    
    async def send_command(
        self,
        device_id: str,
        command: dict,
        timeout: int = 30
    ) -> dict:
        """发送命令并等待响应"""
        
        if not self.is_online(device_id):
            return {"code": 503, "message": "设备不在线"}
        
        # 生成命令ID
        command_id = str(uuid.uuid4())
        command["command_id"] = command_id
        command["timestamp"] = int(datetime.now().timestamp())

        if device_id in self.device_info:
            self.device_info[device_id]["last_command_at"] = datetime.now().isoformat()
            # 长任务（微信登录等）期间延长心跳宽限期，避免误踢
            timeout_hint = int(command.get("timeout") or command.get("data", {}).get("timeout") or 0)
            grace = max(120, min(timeout_hint, 600))
            self.device_info[device_id]["command_grace_seconds"] = grace

        # Android App 协议兼容：execute 需要扁平 action/params
        try:
            info = self.device_info.get(device_id) or {}
            platform = (info.get("platform") or "").lower()
            if platform == "android" and command.get("type") == "execute":
                payload = command.get("data") or {}
                if isinstance(payload, dict):
                    action = payload.get("action")
                    params = payload.get("params") or {}
                    if action:
                        script = payload.get("script")
                        hook_only = payload.get("hook_only")
                        command.pop("data", None)
                        command["action"] = action
                        command["params"] = params
                        if script is not None:
                            command["script"] = script
                        if hook_only is not None:
                            command["hook_only"] = hook_only
        except Exception:
            pass
        
        # 创建Future等待响应
        future = asyncio.get_running_loop().create_future()
        self.pending_commands[command_id] = future
        self.pending_device_commands[device_id] = self.pending_device_commands.get(device_id, 0) + 1

        try:
            # 发送命令（保留 pending_commands 直至收到 response 或超时）
            success = await self.send_to_device(device_id, command)
            if not success:
                self.pending_commands.pop(command_id, None)
                cnt = self.pending_device_commands.get(device_id, 0) - 1
                if cnt <= 0:
                    self.pending_device_commands.pop(device_id, None)
                else:
                    self.pending_device_commands[device_id] = cnt
                if not future.done():
                    future.cancel()
                return {"code": 503, "message": "发送失败"}

            result = await asyncio.wait_for(future, timeout=timeout)
            return result
        
        except asyncio.TimeoutError:
            self.pending_commands.pop(command_id, None)
            cnt = self.pending_device_commands.get(device_id, 0) - 1
            if cnt <= 0:
                self.pending_device_commands.pop(device_id, None)
            else:
                self.pending_device_commands[device_id] = cnt
            if device_id in self.device_info:
                self.device_info[device_id].pop("command_grace_seconds", None)
            logger.warning(f"[ws_hub] 设备响应超时 device_id={device_id} command_id={command_id} timeout={timeout}s")
            return {"code": 408, "message": "设备响应超时"}
        
        except Exception as e:
            self.pending_commands.pop(command_id, None)
            cnt = self.pending_device_commands.get(device_id, 0) - 1
            if cnt <= 0:
                self.pending_device_commands.pop(device_id, None)
            else:
                self.pending_device_commands[device_id] = cnt
            return {"code": 500, "message": str(e)}
    
    def get_online_devices(self) -> list:
        """获取所有在线设备"""
        return list(self.device_info.values())
    
    def get_device_info(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_projects(self) -> dict:
        """获取所有项目及其设备数量"""
        result = {}
        for project_id, device_ids in self.project_devices.items():
            online_count = sum(1 for did in device_ids if did in self.device_info)
            result[project_id] = {
                "project_id": project_id,
                "total_devices": len(device_ids),
                "online_devices": online_count
            }
        return result
    
    async def broadcast_to_project(self, project_id: str, command: dict) -> dict:
        """向项目下所有设备广播命令"""
        device_ids = self.project_devices.get(project_id, [])
        results = {"success": 0, "failed": 0, "devices": []}
        
        for device_id in device_ids:
            if device_id in self.connections:
                success = await self.send_to_device(device_id, command)
                if success:
                    results["success"] += 1
                else:
                    results["failed"] += 1
                results["devices"].append({
                    "device_id": device_id,
                    "success": success
                })
        
        return results
    
    async def execute_on_project(
        self, 
        project_id: str, 
        command: dict,
        timeout: int = 30
    ) -> list:
        """在项目下所有设备执行命令并等待结果"""
        device_ids = self.project_devices.get(project_id, [])
        tasks = []
        
        for device_id in device_ids:
            if device_id in self.connections:
                task = self.send_command(device_id, command.copy(), timeout)
                tasks.append((device_id, task))
        
        results = []
        for device_id, task in tasks:
            try:
                result = await task
                results.append({
                    "device_id": device_id,
                    "result": result
                })
            except Exception as e:
                results.append({
                    "device_id": device_id,
                    "error": str(e)
                })
        
        return results


# 全局实例
ws_hub = WebSocketHub()
