#!/usr/bin/env python3
"""
UAVMonitor - 无人机状态监控面板
提供美观的终端动态刷新显示
"""
import sys
import time
import threading
from typing import Dict, List, Optional, Callable
from dataclasses import dataclass

# 导入控制器（避免循环导入，在方法中使用）
from .px4_swarm_control_v2 import PX4SwarmController


@dataclass
class Colors:
    """ANSI 颜色代码"""
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BLUE = '\033[94m'
    CYAN = '\033[96m'
    MAGENTA = '\033[95m'
    WHITE = '\033[97m'
    BOLD = '\033[1m'
    DIM = '\033[2m'
    CLEAR = '\033[0m'
    CLEAR_SCREEN = '\033[2J\033[H'
    HIDE_CURSOR = '\033[?25l'
    SHOW_CURSOR = '\033[?25h'


class UAVMonitor:
    """
    无人机集群监控面板
    
    使用示例:
        monitor = UAVMonitor()
        monitor.start()
        
        # 更新状态
        monitor.update_uav(1, data_dict)
        monitor.log("起飞完成")
        
        # 停止监控
        monitor.stop()
    """
    
    def __init__(self, refresh_rate: float = 0.5):
        """
        初始化监控面板
        
        Args:
            refresh_rate: 刷新频率（秒），默认0.5秒
        """
        self.refresh_rate = refresh_rate
        self._uav_data: Dict[int, Dict] = {}
        self._logs: List[str] = []
        self._max_logs = 8
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()
        self._custom_render: Optional[Callable] = None
        
    def set_custom_render(self, func: Callable):
        """设置自定义渲染函数"""
        self._custom_render = func
        
    def start(self):
        """启动监控面板"""
        if self._running:
            return
        self._running = True
        sys.stdout.write(Colors.HIDE_CURSOR)
        sys.stdout.flush()
        self._thread = threading.Thread(target=self._render_loop, daemon=True)
        self._thread.start()
        
    def stop(self):
        """停止监控面板"""
        self._running = False
        if self._thread:
            self._thread.join(timeout=1.0)
        sys.stdout.write(Colors.SHOW_CURSOR)
        sys.stdout.flush()
        
    def update_uav(self, uav_id: int, data: Dict):
        """更新无人机状态数据"""
        with self._lock:
            self._uav_data[uav_id] = data.copy()
            
    def update_uavs(self, data_dict: Dict[int, Dict]):
        """批量更新多架无人机状态"""
        with self._lock:
            self._uav_data = {k: v.copy() for k, v in data_dict.items()}
            
    def log(self, message: str):
        """添加日志消息"""
        timestamp = time.strftime('%H:%M:%S')
        with self._lock:
            self._logs.append(f"[{timestamp}] {message}")
            if len(self._logs) > self._max_logs:
                self._logs.pop(0)
                
    def clear_logs(self):
        """清空日志"""
        with self._lock:
            self._logs.clear()
            
    def _clear_screen(self):
        """清屏"""
        sys.stdout.write(Colors.CLEAR_SCREEN)
        sys.stdout.flush()
        
    def _render_loop(self):
        """渲染循环（在后台线程运行）"""
        while self._running:
            self._render()
            time.sleep(self.refresh_rate)
            
    def _render(self):
        """渲染面板"""
        with self._lock:
            uav_data = self._uav_data.copy()
            logs = self._logs.copy()
            
        self._clear_screen()
        
        # 标题
        print(f"{Colors.BOLD}{'='*70}{Colors.CLEAR}")
        print(f"{Colors.CYAN}{Colors.BOLD}{'🚁 无人机集群监控面板':^70}{Colors.CLEAR}")
        print(f"{Colors.DIM}{f'刷新频率: {self.refresh_rate*1000:.0f}ms | 按 Ctrl+C 停止':^70}{Colors.CLEAR}")
        print(f"{Colors.BOLD}{'='*70}{Colors.CLEAR}")
        
        # 自定义渲染区域
        if self._custom_render:
            try:
                self._custom_render()
            except Exception as e:
                print(f"  {Colors.RED}渲染错误: {e}{Colors.CLEAR}")
                
        # 无人机状态
        if uav_data:
            for uav_id, data in sorted(uav_data.items()):
                self._render_uav(uav_id, data)
        else:
            print(f"\n  {Colors.YELLOW}⚠️ 暂无无人机数据{Colors.CLEAR}")
            
        # 日志区域
        print(f"\n{Colors.BOLD}{'─'*70}{Colors.CLEAR}")
        print(f"{Colors.BOLD}📋 任务日志:{Colors.CLEAR}")
        if logs:
            for log in logs:
                print(f"  {log}")
        else:
            print(f"  {Colors.DIM}暂无日志...{Colors.CLEAR}")
        print(f"{Colors.BOLD}{'='*70}{Colors.CLEAR}")
        
    def _render_uav(self, uav_id: int, data: Dict):
        """渲染单个无人机状态"""
        connected = data.get('connected', False)
        
        if not connected:
            status = f"{Colors.RED}● 离线{Colors.CLEAR}"
            print(f"\n  UAV {uav_id}: {status}")
            print(f"    {Colors.DIM}{'─'*60}{Colors.CLEAR}")
            return
            
        # 在线状态
        status = f"{Colors.GREEN}● 在线{Colors.CLEAR}"
        
        # 飞行模式（兼容新旧字段名）
        mode = data.get('flight_mode', data.get('flight_mode_str', 'Unknown'))
        if mode == 'Offboard':
            mode_str = f"{Colors.CYAN}{mode}{Colors.CLEAR}"
        elif mode in ['Takeoff', 'Land']:
            mode_str = f"{Colors.YELLOW}{mode}{Colors.CLEAR}"
        elif mode in ['Hold', 'Return']:
            mode_str = f"{Colors.MAGENTA}{mode}{Colors.CLEAR}"
        else:
            mode_str = f"{Colors.WHITE}{mode}{Colors.CLEAR}"
            
        # 解锁状态
        armed = data.get('is_armed', False)
        if armed:
            armed_str = f"{Colors.GREEN}已解锁 🔓{Colors.CLEAR}"
        else:
            armed_str = f"{Colors.RED}锁定 🔒{Colors.CLEAR}"
            
        # 位置数据
        x = data.get('local_x', 0)
        y = data.get('local_y', 0)
        z = data.get('local_z', 0)
        height = -z  # NED 转高度
        
        # 电池（兼容新旧字段名）
        voltage = data.get('battery_voltage', data.get('battery', 0))
        if voltage > 15:
            bat_str = f"{Colors.GREEN}{voltage:.2f}V 🔋{Colors.CLEAR}"
        elif voltage > 14:
            bat_str = f"{Colors.YELLOW}{voltage:.2f}V 🔋{Colors.CLEAR}"
        else:
            bat_str = f"{Colors.RED}{voltage:.2f}V ⚠️{Colors.CLEAR}"
            
        # 打印状态卡片
        print(f"\n  UAV {uav_id}: {status}  |  模式: {mode_str}  |  {armed_str}")
        print(f"    {Colors.DIM}┌{'─'*58}┐{Colors.CLEAR}")
        print(f"    {Colors.DIM}│{Colors.CLEAR} 位置: X={x:7.2f}m  Y={y:7.2f}m  高度={height:6.2f}m{' '*10}{Colors.DIM}│{Colors.CLEAR}")
        print(f"    {Colors.DIM}│{Colors.CLEAR} 电池: {bat_str}{' '*43}{Colors.DIM}│{Colors.CLEAR}")
        print(f"    {Colors.DIM}└{'─'*58}┘{Colors.CLEAR}")


class SimpleSwarmDemo:
    """
    简化的无人机集群演示基类
    用户只需继承此类并实现 run() 方法
    """
    
    def __init__(self, uav_ids=None, refresh_rate=0.5):
        """
        初始化演示
        
        Args:
            uav_ids: 无人机ID列表，如 [1, 2, 3]
            refresh_rate: 监控刷新频率
        """
        self.uav_ids = uav_ids or [1]
        self.monitor = UAVMonitor(refresh_rate=refresh_rate)
        self.controller = None
        self._running = False
        
    def log(self, message: str):
        """添加日志"""
        self.monitor.log(message)
        
    def sleep(self, seconds: float):
        """睡眠（非阻塞，监控继续刷新）"""
        time.sleep(seconds)
        
    def wait_for_connection(self, uav_id: int, timeout: float = 15.0) -> bool:
        """等待指定无人机连接"""
        self.log(f"等待 UAV {uav_id} 连接...")
        start = time.time()
        while time.time() - start < timeout:
            if self.controller:
                data = self.controller.get_uav_data(uav_id)
                if data.get('connected'):
                    self.log(f"✅ UAV {uav_id} 已连接")
                    return True
            self.sleep(0.1)
        self.log(f"❌ UAV {uav_id} 连接超时")
        return False
        
    def wait_for_height(self, uav_id: int, target_height: float, 
                       tolerance: float = 0.3, timeout: float = 30.0) -> bool:
        """等待无人机到达目标高度"""
        self.log(f"等待到达高度 {target_height}m...")
        start = time.time()
        while time.time() - start < timeout:
            data = self.controller.get_uav_data(uav_id)
            current_height = -data.get('local_z', 0)
            if abs(current_height - target_height) < tolerance:
                self.log(f"✅ 已到达目标高度")
                return True
            self.sleep(0.2)
        self.log(f"⚠️ 高度等待超时")
        return False
        
    def run(self):
        """
        用户需要重写此方法实现任务逻辑
        
        示例:
            def run(self):
                self.takeoff(1, 3.0)
                self.sleep(5)
                self.goto_enu(1, 2, 0, 3)
                self.sleep(5)
                self.land(1)
        """
        raise NotImplementedError("请继承此类并实现 run() 方法")
        
    def takeoff(self, uav_id: int, height: float, blocking: bool = True, timeout: float = None):
        """起飞"""
        self.log(f"🚀 UAV {uav_id} 起飞到 {height}m")
        success, result, message = self.controller.takeoff(uav_id, height, blocking=blocking, timeout=timeout)
        if success:
            self.log(f"✅ UAV {uav_id} 起飞完成")
        else:
            self.log(f"❌ UAV {uav_id} 起飞失败: {message}")
        return success
        
    def land(self, uav_id: int):
        """降落"""
        self.log(f"🛬 UAV {uav_id} 降落")
        self.controller.land(uav_id)
        
    def goto_enu(self, uav_id: int, x: float, y: float, z: float):
        """飞往 ENU 坐标"""
        self.log(f"➡️ UAV {uav_id} 飞往 ({x}, {y}, {z})")
        self.controller.goto_enu(uav_id, x, y, z)
        
    def goto_ned(self, uav_id: int, x: float, y: float, z: float):
        """飞往 NED 坐标"""
        self.log(f"➡️ UAV {uav_id} 飞往 NED ({x}, {y}, {z})")
        self.controller.goto_ned(uav_id, x, y, z)
        
    def goto_body(self, uav_id: int, x: float, y: float, z: float, yaw: float = None):
        """Body坐标系增量控制 (X=前, Y=右, Z=下)"""
        yaw_info = f", Yaw={yaw:.1f}°" if yaw is not None else ""
        self.log(f"➡️ UAV {uav_id} Body移动 (Δx={x}, Δy={y}, Δz={z}{yaw_info})")
        self.controller.goto_body(uav_id, x, y, z, yaw)
        
    def hover(self, uav_id: int):
        """悬停"""
        self.log(f"⏸️ UAV {uav_id} 悬停")
        self.controller.hover(uav_id)
        
    def _update_monitor(self):
        """后台更新监控数据"""
        while self._running:
            if self.controller:
                data = self.controller.get_all_uav_data()
                self.monitor.update_uavs(data)
            self.sleep(0.1)
            
    def start(self):
        """启动演示"""
        import rclpy
        from rclpy.executors import MultiThreadedExecutor
        
        rclpy.init()
        
        # 创建控制器
        self.controller = PX4SwarmController()
        for uav_id in self.uav_ids:
            self.controller.setup_subscribers(uav_id, {})
            
        # 启动 ROS
        executor = MultiThreadedExecutor()
        executor.add_node(self.controller)
        spin_thread = threading.Thread(target=executor.spin, daemon=True)
        spin_thread.start()
        
        # 启动监控
        self.monitor.start()
        self._running = True
        
        # 启动数据更新线程
        update_thread = threading.Thread(target=self._update_monitor, daemon=True)
        update_thread.start()
        
        try:
            # 等待连接
            for uav_id in self.uav_ids:
                if not self.wait_for_connection(uav_id):
                    self.log("连接失败，退出")
                    return
                    
            # 执行用户任务
            self.sleep(1)
            self.run()
            self.sleep(2)
            self.log("✅ 任务完成")
            
        except KeyboardInterrupt:
            self.log("⚠️ 用户中断")
        finally:
            self.sleep(2)
            self._running = False
            self.monitor.stop()
            self.controller.destroy_node()
            rclpy.shutdown()
