"""
AI 心跳监控服务
实时监控设备健康状态，异常检测与自动恢复，任务队列与心跳周期集成。

核心能力：
  1. 设备健康评分（0-100）
  2. 异常检测（连接断开、心跳丢失、电量低、存储满等）
  3. 自动恢复触发（ADB重连、Agent App重启、反向代理恢复）
  4. 心跳周期内下发待执行 AI 任务
  5. 设备状态仪表盘数据聚合
"""

from __future__ import annotations

import asyncio
import logging
import time
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Deque, Dict, List, Optional

logger = logging.getLogger(__name__)


class AlertLevel(str, Enum):
    INFO = "info"
    WARNING = "warning"
    CRITICAL = "critical"


@dataclass
class DeviceHealthEvent:
    device_id: str
    level: AlertLevel
    category: str
    message: str
    ts: str = field(default_factory=lambda: datetime.now().isoformat())
    meta: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "device_id": self.device_id,
            "level": self.level.value,
            "category": self.category,
            "message": self.message,
            "ts": self.ts,
            **self.meta,
        }


@dataclass
class DeviceHealth:
    device_id: str
    score: int = 100
    last_heartbeat: Optional[str] = None
    heartbeat_age_sec: int = 0
    battery_level: int = -1
    battery_charging: bool = False
    memory_usage_pct: float = 0.0
    storage_free_mb: int = -1
    screen_on: bool = False
    network_type: str = "unknown"
    current_app: str = ""
    wechat_running: bool = False
    ws_connected: bool = False
    adb_connected: bool = False
    hook_available: bool = False
    ai_ready: bool = False
    uptime_sec: int = 0
    consecutive_missed: int = 0

    def to_dict(self) -> dict:
        return {
            "device_id": self.device_id,
            "score": self.score,
            "last_heartbeat": self.last_heartbeat,
            "heartbeat_age_sec": self.heartbeat_age_sec,
            "battery": {"level": self.battery_level, "charging": self.battery_charging},
            "memory_usage_pct": self.memory_usage_pct,
            "storage_free_mb": self.storage_free_mb,
            "screen_on": self.screen_on,
            "network_type": self.network_type,
            "current_app": self.current_app,
            "wechat_running": self.wechat_running,
            "connections": {
                "ws": self.ws_connected,
                "adb": self.adb_connected,
                "hook": self.hook_available,
                "ai_ready": self.ai_ready,
            },
            "uptime_sec": self.uptime_sec,
            "consecutive_missed": self.consecutive_missed,
        }


class AIHeartbeatMonitor:
    """
    AI 心跳监控中枢。
    由 main.py lifespan 启动后台循环，每个心跳周期：
      1. 扫描所有在线设备
      2. 计算健康分
      3. 检测异常 → 生成告警
      4. 触发自动恢复
      5. 下发待执行任务
    """

    MAX_EVENTS = 200
    HEALTH_INTERVAL_SEC = 10

    def __init__(self):
        self._health: Dict[str, DeviceHealth] = {}
        self._events: Deque[DeviceHealthEvent] = deque(maxlen=self.MAX_EVENTS)
        self._task_queue: Dict[str, Deque[dict]] = {}
        self._running = False
        self._recovery_cooldown: Dict[str, float] = {}

    # ------------------------------------------------------------------
    # 生命周期
    # ------------------------------------------------------------------

    async def start(self):
        self._running = True
        logger.info("AI 心跳监控已启动")

    async def stop(self):
        self._running = False
        logger.info("AI 心跳监控已停止")

    async def tick(self):
        """主循环中每次被调用（由 main.py _ai_heartbeat_loop 驱动）"""
        if not self._running:
            return

        from services.ws_hub import ws_hub

        now = datetime.now()

        for device_id, info in list(ws_hub.device_info.items()):
            health = self._health.setdefault(device_id, DeviceHealth(device_id=device_id))

            await self._update_health_from_info(health, info, now)
            self._compute_score(health)
            self._detect_anomalies(health)
            await self._dispatch_pending_tasks(device_id)

        for device_id in list(self._health.keys()):
            if device_id not in ws_hub.device_info and device_id not in ws_hub.connections:
                health = self._health[device_id]
                health.ws_connected = False
                health.score = max(0, health.score - 5)
                if health.consecutive_missed > 3:
                    await self._try_recovery(device_id, "ws_lost")

    # ------------------------------------------------------------------
    # 健康数据更新
    # ------------------------------------------------------------------

    async def _update_health_from_info(self, h: DeviceHealth, info: dict, now: datetime):
        from services.ws_hub import ws_hub

        h.ws_connected = h.device_id in ws_hub.connections

        h.last_heartbeat = info.get("last_heartbeat")
        if h.last_heartbeat:
            try:
                dt = datetime.fromisoformat(h.last_heartbeat)
                h.heartbeat_age_sec = max(0, int((now - dt).total_seconds()))
            except Exception:
                h.heartbeat_age_sec = 9999

        qs = info.get("quick_status") or {}
        h.battery_level = qs.get("battery_level", h.battery_level)
        h.battery_charging = qs.get("battery_charging", h.battery_charging)
        h.memory_usage_pct = qs.get("memory_usage_pct", h.memory_usage_pct)
        h.storage_free_mb = qs.get("storage_free_mb", h.storage_free_mb)
        h.screen_on = qs.get("screen_on", h.screen_on) or info.get("screen_on", h.screen_on)
        h.network_type = qs.get("network_type", h.network_type) or info.get("network_type", h.network_type)
        h.current_app = qs.get("current_app", h.current_app) or info.get("current_app", h.current_app)
        h.wechat_running = qs.get("wechat_running", h.wechat_running) or info.get("wechat_running", h.wechat_running)
        h.uptime_sec = qs.get("uptime_sec", h.uptime_sec)

        h.hook_available = info.get("frida_available", False)
        h.ai_ready = h.ws_connected or h.adb_connected

        from services.connection_priority import connection_priority, ControlMode
        try:
            modes = await connection_priority.async_evaluate(h.device_id)
            for m in modes:
                if m.mode == ControlMode.ADB:
                    h.adb_connected = m.available
        except Exception:
            pass

    def _compute_score(self, h: DeviceHealth):
        """综合健康评分 0-100"""
        score = 100

        if not h.ws_connected:
            score -= 30
        if not h.adb_connected:
            score -= 15

        hb_interval = 10
        if h.heartbeat_age_sec > hb_interval * 3:
            score -= 20
            h.consecutive_missed = h.heartbeat_age_sec // max(hb_interval, 1)
        elif h.heartbeat_age_sec > hb_interval * 2:
            score -= 10
            h.consecutive_missed = max(1, h.heartbeat_age_sec // max(hb_interval, 1) - 1)
        else:
            h.consecutive_missed = 0

        if 0 <= h.battery_level <= 15 and not h.battery_charging:
            score -= 15
        elif 0 <= h.battery_level <= 30 and not h.battery_charging:
            score -= 5

        if h.memory_usage_pct > 90:
            score -= 10
        elif h.memory_usage_pct > 80:
            score -= 5

        if 0 < h.storage_free_mb < 500:
            score -= 10
        elif 0 < h.storage_free_mb < 1000:
            score -= 5

        if h.network_type in ("none", "unknown", ""):
            score -= 10

        h.score = max(0, min(100, score))

    # ------------------------------------------------------------------
    # 异常检测
    # ------------------------------------------------------------------

    def _detect_anomalies(self, h: DeviceHealth):
        did = h.device_id

        if h.consecutive_missed >= 5:
            self._emit(did, AlertLevel.CRITICAL, "heartbeat",
                       f"心跳连续丢失 {h.consecutive_missed} 次")

        if 0 <= h.battery_level <= 10 and not h.battery_charging:
            self._emit(did, AlertLevel.CRITICAL, "battery",
                       f"电量极低 {h.battery_level}%，即将关机")
        elif 0 <= h.battery_level <= 20 and not h.battery_charging:
            self._emit(did, AlertLevel.WARNING, "battery",
                       f"电量低 {h.battery_level}%")

        if 0 < h.storage_free_mb < 200:
            self._emit(did, AlertLevel.CRITICAL, "storage",
                       f"存储空间严重不足 {h.storage_free_mb}MB")
        elif 0 < h.storage_free_mb < 500:
            self._emit(did, AlertLevel.WARNING, "storage",
                       f"存储空间不足 {h.storage_free_mb}MB")

        if h.memory_usage_pct > 95:
            self._emit(did, AlertLevel.CRITICAL, "memory",
                       f"内存使用率 {h.memory_usage_pct:.0f}%，可能 OOM")

        if h.network_type in ("none", ""):
            self._emit(did, AlertLevel.WARNING, "network", "网络不可用")

        if h.score <= 30:
            self._emit(did, AlertLevel.CRITICAL, "health",
                       f"设备健康分 {h.score}，需要紧急关注")
        elif h.score <= 60:
            self._emit(did, AlertLevel.WARNING, "health",
                       f"设备健康分 {h.score}，建议检查")

    def _emit(self, device_id: str, level: AlertLevel, category: str, message: str, **meta):
        recent = [e for e in self._events if e.device_id == device_id and e.category == category]
        if recent:
            last = recent[-1]
            try:
                age = (datetime.now() - datetime.fromisoformat(last.ts)).total_seconds()
                if age < 60:
                    return
            except Exception:
                pass

        evt = DeviceHealthEvent(device_id=device_id, level=level, category=category,
                                message=message, meta=meta)
        self._events.append(evt)
        log_fn = logger.warning if level in (AlertLevel.WARNING, AlertLevel.CRITICAL) else logger.info
        log_fn(f"[AI心跳] [{level.value}] {device_id}: {message}")

    # ------------------------------------------------------------------
    # 自动恢复
    # ------------------------------------------------------------------

    async def _try_recovery(self, device_id: str, reason: str):
        cooldown_key = f"{device_id}:{reason}"
        now = time.time()
        if now - self._recovery_cooldown.get(cooldown_key, 0) < 120:
            return
        self._recovery_cooldown[cooldown_key] = now

        logger.info(f"[AI心跳] 尝试自动恢复 {device_id} ({reason})")
        self._emit(device_id, AlertLevel.INFO, "recovery",
                   f"触发自动恢复: {reason}")

        if reason == "ws_lost":
            await self._recover_ws(device_id)

    async def _recover_ws(self, device_id: str):
        """尝试恢复 WebSocket 连接（通过 ADB 重启 Agent App）"""
        from services.adb_device import adb_manager
        serial = adb_manager.resolve_serial(device_id)
        if not serial:
            return
        dev = adb_manager.get_device(serial)
        if not dev or not dev.is_online():
            return

        try:
            dev._shell("am force-stop com.workphone.agent", timeout=5)
            await asyncio.sleep(2)
            dev._shell(
                "am start -n com.workphone.agent/.MainActivity "
                "-a com.workphone.agent.START",
                timeout=5,
            )
            logger.info(f"[AI心跳] 已通过 ADB 重启 Agent App: {device_id}")
        except Exception as e:
            logger.warning(f"[AI心跳] 恢复失败: {e}")

    # ------------------------------------------------------------------
    # 任务队列
    # ------------------------------------------------------------------

    def enqueue_task(self, device_id: str, task: dict):
        """向设备的心跳任务队列添加待执行任务"""
        q = self._task_queue.setdefault(device_id, deque(maxlen=50))
        task.setdefault("queued_at", datetime.now().isoformat())
        q.append(task)
        logger.info(f"[AI心跳] 任务入队 {device_id}: {task.get('instruction', '')[:50]}")

    async def _dispatch_pending_tasks(self, device_id: str):
        """在心跳周期内下发待执行任务"""
        q = self._task_queue.get(device_id)
        if not q:
            return

        from services.ws_hub import ws_hub
        if not ws_hub.is_online(device_id):
            return

        batch_limit = 3
        dispatched = 0
        while q and dispatched < batch_limit:
            task = q.popleft()
            ok = await ws_hub.send_to_device(device_id, {
                "type": "ai_task",
                "data": task,
            })
            if ok:
                dispatched += 1
                logger.info(f"[AI心跳] 任务下发 {device_id}: {task.get('instruction', '')[:50]}")
            else:
                q.appendleft(task)
                break

    # ------------------------------------------------------------------
    # 查询 API
    # ------------------------------------------------------------------

    def get_device_health(self, device_id: str) -> Optional[dict]:
        h = self._health.get(device_id)
        return h.to_dict() if h else None

    def get_all_health(self) -> List[dict]:
        return [h.to_dict() for h in self._health.values()]

    def get_events(self, device_id: Optional[str] = None, limit: int = 50) -> List[dict]:
        evts = list(self._events)
        if device_id:
            evts = [e for e in evts if e.device_id == device_id]
        return [e.to_dict() for e in evts[-limit:]]

    def get_dashboard(self) -> dict:
        """聚合仪表盘数据"""
        all_h = list(self._health.values())
        total = len(all_h)
        healthy = sum(1 for h in all_h if h.score >= 80)
        warning = sum(1 for h in all_h if 50 <= h.score < 80)
        critical = sum(1 for h in all_h if h.score < 50)
        ws_up = sum(1 for h in all_h if h.ws_connected)
        adb_up = sum(1 for h in all_h if h.adb_connected)

        pending_tasks = sum(len(q) for q in self._task_queue.values())
        recent_alerts = [e.to_dict() for e in list(self._events)[-10:]]

        return {
            "total_devices": total,
            "healthy": healthy,
            "warning": warning,
            "critical": critical,
            "ws_connected": ws_up,
            "adb_connected": adb_up,
            "avg_score": round(sum(h.score for h in all_h) / max(total, 1), 1),
            "pending_tasks": pending_tasks,
            "recent_alerts": recent_alerts,
            "devices": [h.to_dict() for h in sorted(all_h, key=lambda x: x.score)],
        }

    def get_task_queue_status(self, device_id: str) -> dict:
        q = self._task_queue.get(device_id, deque())
        return {
            "device_id": device_id,
            "pending_count": len(q),
            "tasks": [t for t in q],
        }


ai_heartbeat = AIHeartbeatMonitor()
