#!/usr/bin/env python3
"""
PX4 无人机集群控制器（优化重构版）
主要改进:
- 修复时间戳格式，使用系统启动后的微秒时间
- 优化 QoS 配置，控制消息使用 RELIABLE
- 优化锁粒度，避免长时间持有锁
- 添加常量定义，消除魔法数字
- 完善错误处理和日志
- 添加类型注解
"""
from __future__ import annotations

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
from rclpy.callback_groups import ReentrantCallbackGroup
from px4_msgs.msg import (
    VehicleStatus, BatteryStatus, VehicleLocalPosition, VehicleGlobalPosition,
    OffboardControlMode, TrajectorySetpoint, VehicleCommand
)
from sensor_msgs.msg import Joy
import threading
import time
import math
from typing import Dict, Set, Callable, Any, Optional, List, Tuple, NamedTuple
from copy import deepcopy
from dataclasses import dataclass
from enum import Enum, auto


# ============================================================================
# 常量定义
# ============================================================================

class Constants:
    """系统常量"""
    # 控制频率
    SETPOINT_RATE_HZ = 20          # 设定点发布频率 (Hz)
    SETPOINT_PERIOD = 1.0 / SETPOINT_RATE_HZ  # 设定点发布周期 (s)
    OFFBOARD_RATE_HZ = 10          # Offboard 定时器频率 (Hz)
    OFFBOARD_PERIOD = 1.0 / OFFBOARD_RATE_HZ  # Offboard 定时器周期 (s)
    
    # 默认参数
    DEFAULT_YAW = math.pi / 2      # 默认偏航角 (90度，东)
    DEFAULT_TAKEOFF_HEIGHT = 3.0   # 默认起飞高度 (m)
    
    # 容差和阈值
    HEIGHT_TOLERANCE = 0.3         # 高度控制容差 (m)
    POSITION_TOLERANCE = 0.3       # 位置控制容差 (m)
    BATTERY_LOW_THRESHOLD = 14.0   # 低电量阈值 (V)
    BATTERY_CRITICAL_THRESHOLD = 13.5  # 严重低电量阈值 (V)
    
    # 超时设置
    CONNECTION_TIMEOUT = 15.0      # 连接超时 (s)
    MODE_SWITCH_TIMEOUT = 5.0      # 模式切换超时 (s)
    ARM_TIMEOUT = 5.0              # 解锁超时 (s)
    TAKEOFF_TIMEOUT = 30.0         # 起飞超时 (s)
    LAND_TIMEOUT = 30.0            # 降落超时 (s)
    
    # 重试间隔
    COMMAND_RETRY_INTERVAL = 0.2   # 命令重试间隔 (s)
    STATUS_CHECK_INTERVAL = 0.5    # 状态检查间隔 (s)
    
    # 计数器循环
    COUNTER_MAX = 1000             # 计数器最大值（循环）
    LOG_PRINT_INTERVAL = 50        # 日志打印间隔（Offboard 回调次数）
    COMMAND_SEND_INTERVAL = 20     # 命令发送间隔（Offboard 回调次数）
    
    # 最大无人机数量
    MAX_UAV_COUNT = 50


# ============================================================================
# 数据类定义
# ============================================================================

@dataclass
class Setpoint:
    """目标设定点"""
    x: float
    y: float
    z: float
    yaw: float = Constants.DEFAULT_YAW
    
    def to_tuple(self) -> Tuple[float, float, float, float]:
        """转换为元组"""
        return (self.x, self.y, self.z, self.yaw)


class TakeoffResult(Enum):
    """起飞结果"""
    SUCCESS = auto()
    OFFBOARD_FAILED = auto()
    ARM_FAILED = auto()
    HEIGHT_TIMEOUT = auto()
    NOT_CONNECTED = auto()


class FlightMode(Enum):
    """PX4 飞行模式映射"""
    MANUAL = 0
    ALTITUDE = 1
    POSITION = 2
    AUTO_MISSION = 3
    AUTO_LOITER = 4
    AUTO_RTL = 5
    POSITION_SLOW = 6
    ACRO = 10
    DESCEND = 12
    TERMINATION = 13
    OFFBOARD = 14
    STABILIZED = 15
    AUTO_TAKEOFF = 17
    AUTO_LAND = 18
    AUTO_FOLLOW_TARGET = 19
    AUTO_PRECLAND = 20
    ORBIT = 21
    AUTO_VTOL_TAKEOFF = 22
    EXTERNAL1 = 23
    EXTERNAL2 = 24
    EXTERNAL3 = 25
    EXTERNAL4 = 26
    
    @classmethod
    def to_string(cls, nav_state: int) -> str:
        """转换为字符串"""
        mode_map = {
            0: 'Manual',
            1: 'Altitude',
            2: 'Position',
            3: 'Auto Mission',
            4: 'Auto Loiter',
            5: 'Auto RTL',
            6: 'Position Slow',
            7: 'Free5',
            8: 'Free4',
            9: 'Free3',
            10: 'Acro',
            11: 'Free2',
            12: 'Descend',
            13: 'Termination',
            14: 'Offboard',
            15: 'Stabilized',
            16: 'Free1',
            17: 'Auto Takeoff',
            18: 'Auto Land',
            19: 'Auto Follow Target',
            20: 'Auto PrecLand',
            21: 'Orbit',
            22: 'Auto VTOL Takeoff',
            23: 'External1',
            24: 'External2',
            25: 'External3',
            26: 'External4',
        }
        return mode_map.get(nav_state, f'Unknown({nav_state})')


class UAVStatus(NamedTuple):
    """无人机状态（不可变）"""
    connected: bool = False
    armed: bool = False
    flight_mode: str = 'Unknown'
    battery_voltage: float = 0.0
    local_x: float = 0.0
    local_y: float = 0.0
    local_z: float = 0.0
    yaw: float = 0.0
    global_lat: float = 0.0
    global_lon: float = 0.0
    global_alt: float = 0.0


# ============================================================================
# 坐标转换工具
# ============================================================================

def ned_to_enu(x: float, y: float, z: float) -> Tuple[float, float, float]:
    """NED坐标系转ENU坐标系"""
    return y, x, -z


def enu_to_ned(x: float, y: float, z: float) -> Tuple[float, float, float]:
    """ENU坐标系转NED坐标系"""
    return y, x, -z


def body_to_ned(bx: float, by: float, bz: float, yaw: float) -> Tuple[float, float, float]:
    """体轴系坐标转NED坐标系"""
    nx = bx * math.cos(yaw) - by * math.sin(yaw)
    ny = bx * math.sin(yaw) + by * math.cos(yaw)
    nz = bz
    return nx, ny, nz


# ============================================================================
# 主控制器
# ============================================================================

class PX4SwarmController(Node):
    """PX4 无人机集群控制器"""
    
    def __init__(self, node_name: str = 'px4_swarm_controller'):
        super().__init__(node_name)
        
        # 线程锁
        self._lock = threading.RLock()
        
        # 状态存储
        self._uav_data: Dict[int, Dict[str, Any]] = {}
        self._target_sp: Dict[int, Setpoint] = {}
        self._offboard_active: Set[int] = set()
        self._hovering: Set[int] = set()
        self._landing: Set[int] = set()
        self._returning: Set[int] = set()
        self._selected_uavs: Set[int] = set()
        
        # 回调存储
        self._external_callbacks: Dict[int, Dict[str, Callable]] = {}
        
        # 定时器
        self._offboard_timers: Dict[int, Any] = {}
        self._counters: Dict[int, int] = {}
        
        # 摇杆数据
        self._joy_data = {
            'axes': [0.0] * 8,
            'buttons': [0] * 12,
            'connected': False,
            'last_update': None
        }
        
        # 系统启动时间基准（用于 PX4 时间戳）
        self._startup_time = self.get_clock().now()
        
        # QoS 配置
        self._init_qos_profiles()
        
        # 回调组
        self._cb_group = ReentrantCallbackGroup()
        
        # 预创建发布者（1-50号UAV）
        self._pubs = {
            i: self._create_publishers(i) 
            for i in range(1, Constants.MAX_UAV_COUNT + 1)
        }
        
        # 订阅摇杆
        self._setup_joy_subscription()
        
        self.get_logger().info('✅ PX4SwarmController 初始化完成')
    
    def _init_qos_profiles(self):
        """初始化 QoS 配置"""
        # 控制消息：使用 RELIABLE 确保送达
        self._control_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.VOLATILE,
            history=HistoryPolicy.KEEP_LAST,
            depth=1
        )
        
        # 状态订阅：使用 BEST_EFFORT 减少开销
        self._status_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            history=HistoryPolicy.KEEP_LAST,
            depth=1
        )
        
        # 摇杆订阅：使用 RELIABLE
        self._joy_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.VOLATILE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10
        )
    
    def _get_px4_timestamp(self) -> int:
        """获取 PX4 格式的时间戳（系统启动后的微秒）"""
        now = self.get_clock().now()
        elapsed = now - self._startup_time
        return int(elapsed.nanoseconds / 1000)
    
    def _create_publishers(self, uav_id: int) -> Dict[str, Any]:
        """为指定 UAV 创建发布者"""
        ns = f'/uav{uav_id}'
        return {
            'offboard': self.create_publisher(
                OffboardControlMode, 
                f'{ns}/fmu/in/offboard_control_mode', 
                self._control_qos
            ),
            'trajectory': self.create_publisher(
                TrajectorySetpoint, 
                f'{ns}/fmu/in/trajectory_setpoint', 
                self._control_qos
            ),
            'vehicle_command': self.create_publisher(
                VehicleCommand, 
                f'{ns}/fmu/in/vehicle_command', 
                self._control_qos
            )
        }
    
    def _setup_joy_subscription(self):
        """设置摇杆订阅"""
        self.create_subscription(
            Joy, '/joy', self._joy_callback,
            self._joy_qos,
            callback_group=self._cb_group
        )
    
    # ---------------------------------------------------------------------
    # 订阅管理
    # ---------------------------------------------------------------------
    
    def setup_subscribers(self, uav_id: int, callbacks: Dict[str, Callable]):
        """设置指定 UAV 的订阅"""
        ns = f'/uav{uav_id}'
        
        with self._lock:
            self._external_callbacks[uav_id] = callbacks
            
            # 初始化 UAV 数据
            if uav_id not in self._uav_data:
                self._uav_data[uav_id] = {
                    'connected': False,
                    'armed': False,
                    'flight_mode': 'Unknown',
                    'battery_voltage': 0.0,
                    'local_x': 0.0,
                    'local_y': 0.0,
                    'local_z': 0.0,
                    'yaw': 0.0,
                    'global_lat': 0.0,
                    'global_lon': 0.0,
                    'global_alt': 0.0
                }
        
        # 创建订阅（在锁外）
        self.create_subscription(
            VehicleStatus, f'{ns}/fmu/out/vehicle_status',
            self._make_callback(self._vehicle_status_callback, uav_id),
            self._status_qos,
            callback_group=self._cb_group
        )
        self.create_subscription(
            BatteryStatus, f'{ns}/fmu/out/battery_status',
            self._make_callback(self._battery_callback, uav_id),
            self._status_qos,
            callback_group=self._cb_group
        )
        self.create_subscription(
            VehicleLocalPosition, f'{ns}/fmu/out/vehicle_local_position',
            self._make_callback(self._local_position_callback, uav_id),
            self._status_qos,
            callback_group=self._cb_group
        )
        self.create_subscription(
            VehicleGlobalPosition, f'{ns}/fmu/out/vehicle_global_position',
            self._make_callback(self._global_position_callback, uav_id),
            self._status_qos,
            callback_group=self._cb_group
        )
    
    def _make_callback(self, cb_func: Callable, uid: int) -> Callable:
        """创建闭包安全的回调"""
        def callback(msg):
            return cb_func(msg, uid)
        return callback
    
    # ---------------------------------------------------------------------
    # 回调函数
    # ---------------------------------------------------------------------
    
    def _vehicle_status_callback(self, msg: VehicleStatus, uid: int):
        """车辆状态回调"""
        with self._lock:
            if uid not in self._uav_data:
                return
            
            d = self._uav_data[uid]
            d['connected'] = True
            d['armed'] = (msg.arming_state == VehicleStatus.ARMING_STATE_ARMED)
            d['flight_mode'] = FlightMode.to_string(msg.nav_state)
            
            self.get_logger().debug(
                f'UAV{uid}: 状态更新 - 模式={d["flight_mode"]}, 解锁={d["armed"]}'
            )
        
        # 调用外部回调（在锁外）
        self._call_external_callback(uid, 'vehicle_status', msg)
    
    def _battery_callback(self, msg: BatteryStatus, uid: int):
        """电池状态回调"""
        voltage = getattr(msg, 'voltage_v',
            getattr(msg, 'voltage_filtered_v',
                getattr(msg, 'voltage', 0.0)))
        
        with self._lock:
            if uid not in self._uav_data:
                return
            self._uav_data[uid]['battery_voltage'] = round(voltage, 2)
        
        self._call_external_callback(uid, 'battery_status', msg)
    
    def _local_position_callback(self, msg: VehicleLocalPosition, uid: int):
        """本地位置回调"""
        with self._lock:
            if uid not in self._uav_data:
                return
            
            d = self._uav_data[uid]
            d['local_x'] = msg.x
            d['local_y'] = msg.y
            d['local_z'] = msg.z
            d['yaw'] = getattr(msg, 'heading', 0.0)
            
            self.get_logger().debug(
                f'UAV{uid}: 位置更新 - X={msg.x:.2f}, Y={msg.y:.2f}, Z={msg.z:.2f}'
            )
        
        self._call_external_callback(uid, 'local_position', msg)
    
    def _global_position_callback(self, msg: VehicleGlobalPosition, uid: int):
        """全局位置回调"""
        with self._lock:
            if uid not in self._uav_data:
                return
            
            d = self._uav_data[uid]
            d['global_lat'] = msg.lat
            d['global_lon'] = msg.lon
            d['global_alt'] = msg.alt
        
        self._call_external_callback(uid, 'global_position', msg)
    
    def _joy_callback(self, msg: Joy):
        """摇杆回调"""
        with self._lock:
            self._joy_data.update({
                'axes': list(msg.axes),
                'buttons': list(msg.buttons),
                'connected': True,
                'last_update': time.time()
            })
    
    def _call_external_callback(self, uid: int, event: str, msg):
        """调用外部回调"""
        callbacks = self._external_callbacks.get(uid, {})
        if event in callbacks:
            try:
                callbacks[event](msg, uid)
            except Exception as e:
                self.get_logger().error(f'UAV{uid}: 外部回调错误 - {e}')
    
    # ---------------------------------------------------------------------
    # 状态查询
    # ---------------------------------------------------------------------
    
    def get_uav_status(self, uid: int) -> UAVStatus:
        """获取 UAV 状态（线程安全）"""
        with self._lock:
            d = self._uav_data.get(uid, {})
            return UAVStatus(
                connected=d.get('connected', False),
                armed=d.get('armed', False),
                flight_mode=d.get('flight_mode', 'Unknown'),
                battery_voltage=d.get('battery_voltage', 0.0),
                local_x=d.get('local_x', 0.0),
                local_y=d.get('local_y', 0.0),
                local_z=d.get('local_z', 0.0),
                yaw=d.get('yaw', 0.0),
                global_lat=d.get('global_lat', 0.0),
                global_lon=d.get('global_lon', 0.0),
                global_alt=d.get('global_alt', 0.0)
            )
    
    def get_uav_data(self, uid: int) -> Dict[str, Any]:
        """获取 UAV 数据（兼容旧 API）"""
        return self.get_uav_status(uid)._asdict()
    
    def get_all_uav_data(self) -> Dict[int, Dict[str, Any]]:
        """获取所有 UAV 数据"""
        result = {}
        with self._lock:
            for uid in self._uav_data:
                result[uid] = self.get_uav_data(uid)
        return result
    
    def is_connected(self, uid: int) -> bool:
        """检查 UAV 是否已连接"""
        return self.get_uav_status(uid).connected
    
    def is_armed(self, uid: int) -> bool:
        """检查 UAV 是否已解锁"""
        return self.get_uav_status(uid).armed
    
    # ---------------------------------------------------------------------
    # 高级控制 API
    # ---------------------------------------------------------------------
    
    def takeoff(self, uid: int, height: Optional[float] = None, 
                blocking: bool = True, timeout: float = None) -> Tuple[bool, TakeoffResult, str]:
        """
        起飞到指定高度
        
        Args:
            uid: 无人机ID
            height: 目标高度（米），默认 3.0
            blocking: 是否阻塞等待
            timeout: 超时时间（秒），默认 30
            
        Returns:
            (成功, 结果码, 详细信息)
        """
        if timeout is None:
            timeout = Constants.TAKEOFF_TIMEOUT
        
        if height is None:
            height = Constants.DEFAULT_TAKEOFF_HEIGHT
        
        target_z = -height  # NED 坐标
        
        # 检查连接
        if not self.is_connected(uid):
            return False, TakeoffResult.NOT_CONNECTED, f'UAV{uid}: 未连接'
        
        # 获取当前位置
        status = self.get_uav_status(uid)
        current_x, current_y = status.local_x, status.local_y
        
        self.get_logger().info(f'UAV{uid}: 🚀 准备起飞到 {height}m...')
        
        # 步骤1: 发送设定点流
        self.get_logger().info(f'UAV{uid}: 步骤1/4: 发送设定点流...')
        start_time = time.time()
        while time.time() - start_time < 2.0:
            self._publish_offboard_mode(uid)
            self._publish_position(uid, current_x, current_y, target_z, Constants.DEFAULT_YAW)
            time.sleep(Constants.SETPOINT_PERIOD)
        
        # 步骤2: 启动 Offboard 定时器
        with self._lock:
            self._target_sp[uid] = Setpoint(current_x, current_y, target_z)
            self._offboard_active.add(uid)
            self._hovering.discard(uid)
        self._start_offboard_timer(uid)
        
        # 步骤3: 切换到 Offboard 模式
        self.get_logger().info(f'UAV{uid}: 步骤2/4: 切换 Offboard 模式...')
        if not self._wait_for_mode(uid, 'Offboard', Constants.MODE_SWITCH_TIMEOUT):
            self._leave_offboard(uid)
            return False, TakeoffResult.OFFBOARD_FAILED, f'UAV{uid}: Offboard 模式切换失败'
        self.get_logger().info(f'UAV{uid}: ✅ Offboard 模式已激活')
        
        # 步骤4: 解锁
        self.get_logger().info(f'UAV{uid}: 步骤3/4: 解锁...')
        if not self._arm(uid, timeout=Constants.ARM_TIMEOUT):
            self._leave_offboard(uid)
            return False, TakeoffResult.ARM_FAILED, f'UAV{uid}: 解锁失败'
        self.get_logger().info(f'UAV{uid}: ✅ 已解锁')
        
        # 步骤5: 等待到达目标高度
        if blocking:
            self.get_logger().info(f'UAV{uid}: 步骤4/4: 等待到达目标高度...')
            if not self._wait_for_height(uid, height, timeout=timeout):
                return False, TakeoffResult.HEIGHT_TIMEOUT, f'UAV{uid}: 高度等待超时'
            self.get_logger().info(f'UAV{uid}: ✅ 已到达目标高度 {height}m')
        
        return True, TakeoffResult.SUCCESS, f'UAV{uid}: 起飞成功'
    
    def land(self, uid: int, blocking: bool = False, timeout: float = None) -> bool:
        """降落"""
        if timeout is None:
            timeout = Constants.LAND_TIMEOUT
        
        self._leave_offboard(uid)
        
        with self._lock:
            self._landing.add(uid)
        
        self._publish_vehicle_command(uid, VehicleCommand.VEHICLE_CMD_NAV_LAND)
        self.get_logger().info(f'UAV{uid}: 降落')
        
        if blocking:
            return self._wait_for_disarm(uid, timeout)
        return True
    
    def hover(self, uid: int):
        """悬停"""
        status = self.get_uav_status(uid)
        if not status.connected:
            self.get_logger().error(f'UAV{uid}: 未连接，无法悬停')
            return
        
        with self._lock:
            self._target_sp[uid] = Setpoint(status.local_x, status.local_y, status.local_z)
            self._offboard_active.add(uid)
            self._hovering.add(uid)
        
        self._start_offboard_timer(uid)
        self.get_logger().info(f'UAV{uid}: 悬停')
    
    def goto_enu(self, uid: int, x: float, y: float, z: float, yaw: float = None):
        """飞往 ENU 坐标"""
        if yaw is None:
            yaw = Constants.DEFAULT_YAW
        
        ned_x, ned_y, ned_z = enu_to_ned(x, y, z)
        
        with self._lock:
            self._target_sp[uid] = Setpoint(ned_x, ned_y, ned_z, yaw)
            self._offboard_active.add(uid)
        
        self._start_offboard_timer(uid)
        self.get_logger().info(f'UAV{uid}: 飞往 ENU ({x:.1f}, {y:.1f}, {z:.1f})')
    
    def goto_ned(self, uid: int, x: float, y: float, z: float, yaw: float = None):
        """飞往 NED 坐标"""
        if yaw is None:
            yaw = Constants.DEFAULT_YAW
        
        with self._lock:
            self._target_sp[uid] = Setpoint(x, y, z, yaw)
            self._offboard_active.add(uid)
        
        self._start_offboard_timer(uid)
        self.get_logger().info(f'UAV{uid}: 飞往 NED ({x:.1f}, {y:.1f}, {z:.1f})')
    
    def goto_body(self, uid: int, x: float, y: float, z: float, yaw: float = None):
        """体轴系移动"""
        if yaw is None:
            yaw = Constants.DEFAULT_YAW
        
        status = self.get_uav_status(uid)
        if not status.connected:
            self.get_logger().error(f'UAV{uid}: 未连接')
            return
        
        # 体轴系转 NED
        delta_x, delta_y, delta_z = body_to_ned(x, y, z, status.yaw)
        target_x = status.local_x + delta_x
        target_y = status.local_y + delta_y
        target_z = status.local_z + delta_z
        
        with self._lock:
            self._target_sp[uid] = Setpoint(target_x, target_y, target_z, yaw)
            self._offboard_active.add(uid)
        
        self._start_offboard_timer(uid)
        self.get_logger().info(f'UAV{uid}: 体轴系移动 ({x:.1f}, {y:.1f}, {z:.1f})')
    
    def return_home(self, uid: int):
        """返航"""
        self._leave_offboard(uid)
        with self._lock:
            self._returning.add(uid)
        self._publish_vehicle_command(uid, VehicleCommand.VEHICLE_CMD_NAV_RETURN_TO_LAUNCH)
        self.get_logger().info(f'UAV{uid}: 返航')
    
    def disarm(self, uid: int):
        """上锁"""
        self._leave_offboard(uid)
        self._publish_vehicle_command(uid, VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM, param1=0.0)
        self.get_logger().info(f'UAV{uid}: 上锁')
    
    def emergency_stop(self, uid: int):
        """紧急停止"""
        self._leave_offboard(uid)
        self._publish_vehicle_command(uid, VehicleCommand.VEHICLE_CMD_DO_FLIGHTTERMINATION, param1=1.0)
        self.get_logger().error(f'UAV{uid}: 紧急停止！')
    
    def reboot(self, uid: int):
        """重启飞控"""
        self._publish_vehicle_command(
            uid, VehicleCommand.VEHICLE_CMD_PREFLIGHT_REBOOT_SHUTDOWN,
            param1=1.0, param2=1.0
        )
        self.get_logger().info(f'UAV{uid}: 重启飞控')
    
    # ---------------------------------------------------------------------
    # 等待方法
    # ---------------------------------------------------------------------
    
    def _wait_for_mode(self, uid: int, target_mode: str, timeout: float) -> bool:
        """等待切换到目标模式"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            status = self.get_uav_status(uid)
            if status.flight_mode == target_mode:
                return True
            
            # 发送模式切换命令
            self._publish_vehicle_command(
                uid, VehicleCommand.VEHICLE_CMD_DO_SET_MODE,
                param1=1.0, param2=6.0
            )
            time.sleep(Constants.COMMAND_RETRY_INTERVAL)
        
        return False
    
    def _arm(self, uid: int, timeout: float) -> bool:
        """解锁 UAV"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            if self.is_armed(uid):
                return True
            
            self._publish_vehicle_command(
                uid, VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM,
                param1=1.0
            )
            time.sleep(Constants.COMMAND_RETRY_INTERVAL)
        
        return False
    
    def _wait_for_height(self, uid: int, target_height: float, 
                         tolerance: float = None, timeout: float = None) -> bool:
        """等待到达目标高度"""
        if tolerance is None:
            tolerance = Constants.HEIGHT_TOLERANCE
        if timeout is None:
            timeout = Constants.TAKEOFF_TIMEOUT
        
        start_time = time.time()
        while time.time() - start_time < timeout:
            status = self.get_uav_status(uid)
            current_height = -status.local_z
            
            if abs(current_height - target_height) < tolerance:
                return True
            
            time.sleep(Constants.STATUS_CHECK_INTERVAL)
        
        return False
    
    def _wait_for_disarm(self, uid: int, timeout: float) -> bool:
        """等待上锁"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            if not self.is_armed(uid):
                return True
            time.sleep(Constants.STATUS_CHECK_INTERVAL)
        return False
    
    # ---------------------------------------------------------------------
    # 内部辅助
    # ---------------------------------------------------------------------
    
    def _leave_offboard(self, uid: int):
        """退出 Offboard 模式"""
        # 先标记为不活跃
        with self._lock:
            self._offboard_active.discard(uid)
            self._hovering.discard(uid)
            self._target_sp.pop(uid, None)
        
        # 在锁外取消定时器
        timer = self._offboard_timers.pop(uid, None)
        if timer:
            timer.cancel()
    
    def _start_offboard_timer(self, uid: int):
        """启动 Offboard 定时器"""
        if uid in self._offboard_timers:
            return
        
        timer = self.create_timer(
            Constants.OFFBOARD_PERIOD,
            lambda: self._offboard_timer_callback(uid),
            callback_group=self._cb_group
        )
        self._offboard_timers[uid] = timer
    
    def _offboard_timer_callback(self, uid: int):
        """Offboard 定时器回调"""
        with self._lock:
            if uid not in self._offboard_active:
                return
            
            target = self._target_sp.get(uid, Setpoint(0, 0, -Constants.DEFAULT_TAKEOFF_HEIGHT))
            status = self._uav_data.get(uid, {})
            
            if not status:
                return
            
            # 更新计数器（循环）
            self._counters[uid] = (self._counters.get(uid, 0) + 1) % Constants.COUNTER_MAX
            counter = self._counters[uid]
        
        # 在锁外发布消息
        self._publish_offboard_mode(uid)
        self._publish_position(uid, target.x, target.y, target.z, target.yaw)
        
        # 定期检查和发送命令
        if counter % Constants.COMMAND_SEND_INTERVAL == 0:
            status_obj = self.get_uav_status(uid)
            
            # 定期打印状态
            if counter % Constants.LOG_PRINT_INTERVAL == 0:
                self.get_logger().info(
                    f'UAV{uid}: Offboard 运行中 - '
                    f'目标=({target.x:.1f}, {target.y:.1f}, {target.z:.1f}), '
                    f'模式={status_obj.flight_mode}, 解锁={status_obj.armed}, '
                    f'高度={-status_obj.local_z:.1f}m'
                )
            
            # 如果需要，重新发送模式切换命令
            if status_obj.flight_mode != 'Offboard':
                self._publish_vehicle_command(
                    uid, VehicleCommand.VEHICLE_CMD_DO_SET_MODE,
                    param1=1.0, param2=6.0
                )
            
            # 如果需要，重新发送解锁命令
            if not status_obj.armed:
                self._publish_vehicle_command(
                    uid, VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM,
                    param1=1.0
                )
    
    # ---------------------------------------------------------------------
    # 消息发布
    # ---------------------------------------------------------------------
    
    def _publish_offboard_mode(self, uid: int) -> bool:
        """发布 Offboard 控制模式"""
        try:
            msg = OffboardControlMode()
            msg.position = True
            msg.velocity = False
            msg.acceleration = False
            msg.attitude = False
            msg.body_rate = False
            msg.thrust_and_torque = False
            msg.direct_actuator = False
            msg.timestamp = self._get_px4_timestamp()
            
            self._pubs[uid]['offboard'].publish(msg)
            return True
        except Exception as e:
            self.get_logger().error(f'UAV{uid}: 发布 Offboard 模式失败 - {e}')
            return False
    
    def _publish_position(self, uid: int, x: float, y: float, z: float, yaw: float) -> bool:
        """发布位置设定点"""
        try:
            msg = TrajectorySetpoint()
            msg.position = [float(x), float(y), float(z)]
            msg.velocity = [float('nan')] * 3
            msg.acceleration = [float('nan')] * 3
            msg.jerk = [float('nan')] * 3
            msg.yaw = float(yaw)
            msg.yawspeed = float('nan')
            msg.timestamp = self._get_px4_timestamp()
            
            self._pubs[uid]['trajectory'].publish(msg)
            return True
        except Exception as e:
            self.get_logger().error(f'UAV{uid}: 发布位置设定点失败 - {e}')
            return False
    
    def _publish_vehicle_command(self, uid: int, cmd: int, **params) -> bool:
        """发布车辆命令"""
        try:
            msg = VehicleCommand()
            msg.command = cmd
            
            for i in range(1, 8):
                setattr(msg, f'param{i}', float(params.get(f'param{i}', 0.0)))
            
            msg.target_system = uid
            msg.target_component = 1
            msg.source_system = 1
            msg.source_component = 1
            msg.from_external = True
            msg.timestamp = self._get_px4_timestamp()
            
            self._pubs[uid]['vehicle_command'].publish(msg)
            return True
        except Exception as e:
            self.get_logger().error(f'UAV{uid}: 发布命令失败 - {e}')
            return False
    
    # ---------------------------------------------------------------------
    # 选择管理（兼容旧 API）
    # ---------------------------------------------------------------------
    
    def select_uav(self, uid: int):
        """选中 UAV"""
        with self._lock:
            self._selected_uavs.add(uid)
    
    def deselect_uav(self, uid: int):
        """取消选中 UAV"""
        self._leave_offboard(uid)
        with self._lock:
            self._selected_uavs.discard(uid)
    
    def get_selected_uavs(self) -> Set[int]:
        """获取选中的 UAV 列表"""
        with self._lock:
            return self._selected_uavs.copy()
    
    @property
    def target_sp(self) -> Dict[int, Setpoint]:
        """目标设定点（兼容旧 API）"""
        with self._lock:
            return self._target_sp.copy()
    
    # ---------------------------------------------------------------------
    # 兼容性属性（供旧 UI 使用）
    # ---------------------------------------------------------------------
    
    @property
    def uav_data(self) -> Dict[int, Dict[str, Any]]:
        """UAV 数据（兼容旧 API）"""
        return self.get_all_uav_data()
    
    @property
    def offboard_active(self) -> Set[int]:
        """活跃的 Offboard（兼容旧 API）"""
        with self._lock:
            return self._offboard_active.copy()
    
    @property
    def landing_uavs(self) -> Set[int]:
        """正在降落的 UAV（兼容旧 API）"""
        with self._lock:
            return self._landing.copy()
    
    @property
    def hovering_uavs(self) -> Set[int]:
        """正在悬停的 UAV（兼容旧 API）"""
        with self._lock:
            return self._hovering.copy()
    
    @property
    def returning_uavs(self) -> Set[int]:
        """正在返航的 UAV（兼容旧 API）"""
        with self._lock:
            return self._returning.copy()
    
    @property
    def selected_uavs(self) -> Set[int]:
        """选中的 UAV（兼容旧 API）"""
        return self.get_selected_uavs()
    
    @selected_uavs.setter
    def selected_uavs(self, value: Set[int]):
        """设置选中的 UAV"""
        with self._lock:
            self._selected_uavs = value.copy()
    
    @property
    def offboard_timers(self) -> Dict[int, Any]:
        """Offboard 定时器（兼容旧 API）"""
        with self._lock:
            return self._offboard_timers.copy()
    
    @property
    def _global_offboard_flag(self) -> bool:
        """全局 Offboard 标志（兼容旧 API）"""
        return len(self._offboard_active) > 0
    
    @property
    def global_offboard_running(self) -> bool:
        """全局 Offboard 运行状态（兼容旧 API）"""
        return self._global_offboard_flag
    
    @property
    def selected_uavs_offboard_running(self) -> Dict[int, bool]:
        """选中 UAV 的 Offboard 运行状态（兼容旧 API）"""
        with self._lock:
            return {uid: uid in self._offboard_timers for uid in self._selected_uavs}
    
    def get_joy_data(self) -> Dict[str, Any]:
        """获取摇杆数据（兼容旧 API）"""
        with self._lock:
            d = self._joy_data.copy()
            if d['last_update'] and time.time() - d['last_update'] > 2.0:
                d['connected'] = False
            return d
    
    # ---------------------------------------------------------------------
    # 兼容性发布方法（供旧 UI 使用）
    # ---------------------------------------------------------------------
    
    def publish_offboard_control_mode(self, uid: int):
        """发布 Offboard 控制模式（兼容旧 API）"""
        self._publish_offboard_mode(uid)
    
    def publish_position_setpoint(self, uid: int, x: float, y: float, z: float, yaw: float = None):
        """发布位置设定点（兼容旧 API）"""
        if yaw is None:
            yaw = Constants.DEFAULT_YAW
        self._publish_position(uid, x, y, z, yaw)
    
    def publish_vehicle_command(self, uid: int, cmd: int, **params):
        """发布车辆命令（兼容旧 API）"""
        self._publish_vehicle_command(uid, cmd, **params)
    
    # ---------------------------------------------------------------------
    # 工具方法
    # ---------------------------------------------------------------------
    
    def print_status(self):
        """打印状态"""
        self.get_logger().info('=' * 60)
        self.get_logger().info('PX4SwarmController 状态')
        self.get_logger().info(f'已连接: {sum(1 for s in [self.get_uav_status(i) for i in range(1, Constants.MAX_UAV_COUNT + 1)] if s.connected)}')
        
        for uid in sorted(self._uav_data.keys()):
            status = self.get_uav_status(uid)
            if status.connected:
                self.get_logger().info(
                    f'  UAV{uid}: 模式={status.flight_mode}, '
                    f'电池={status.battery_voltage:.1f}V, '
                    f'位置=({status.local_x:+.1f}, {status.local_y:+.1f}, {status.local_z:+.1f})'
                )
        self.get_logger().info('=' * 60)


# ============================================================================
# 模块自测
# ============================================================================

def main():
    import sys
    rclpy.init(args=sys.argv)
    node = PX4SwarmController()
    
    try:
        node.print_status()
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info('Ctrl-C 退出')
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
