#!/usr/bin/env python3
# 多无人机offboard控制程序 - 支持原地起飞和状态监控

# 导入必要的ROS 2库
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
# 导入PX4消息类型
from px4_msgs.msg import OffboardControlMode, TrajectorySetpoint, VehicleCommand, VehicleLocalPosition, VehicleStatus
import time  # 用于添加延迟


class UAVController:
    """单个无人机的控制器类"""
    
    def __init__(self, node: Node, uav_id: int):
        """初始化单个无人机控制器"""
        self.node = node  # ROS 2节点
        self.uav_id = uav_id  # 无人机ID
        self.namespace = f'/uav{uav_id}'  # 话题命名空间
        
        # 配置QoS配置文件，用于发布和订阅
        qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,  # 尽力而为的可靠性
            durability=DurabilityPolicy.TRANSIENT_LOCAL,  # 瞬态本地持久性
            history=HistoryPolicy.KEEP_LAST,  # 只保留最后一条消息
            depth=1  # 队列深度为1
        )

        # 创建发布器，使用无人机特定的话题
        self.offboard_control_mode_publisher = self.node.create_publisher(
            OffboardControlMode, f'{self.namespace}/fmu/in/offboard_control_mode', qos_profile)
        self.trajectory_setpoint_publisher = self.node.create_publisher(
            TrajectorySetpoint, f'{self.namespace}/fmu/in/trajectory_setpoint', qos_profile)
        self.vehicle_command_publisher = self.node.create_publisher(
            VehicleCommand, f'{self.namespace}/fmu/in/vehicle_command', qos_profile)

        # 创建订阅器，使用无人机特定的话题
        self.vehicle_local_position_subscriber = self.node.create_subscription(
            VehicleLocalPosition, f'{self.namespace}/fmu/out/vehicle_local_position_v1', 
            self.vehicle_local_position_callback, qos_profile)
        self.vehicle_status_subscriber = self.node.create_subscription(
            VehicleStatus, f'{self.namespace}/fmu/out/vehicle_status_v1', 
            self.vehicle_status_callback, qos_profile)

        # 初始化变量
        self.offboard_setpoint_counter = 0  # offboard设定点计数器
        self.vehicle_local_position = VehicleLocalPosition()  # 车辆本地位置
        self.vehicle_status = VehicleStatus()  # 车辆状态
        self.takeoff_height = -3.0  # 起飞高度（负值表示NED坐标系）
        self.has_taken_off = False  # 是否已起飞标志
        self.is_armed = False  # 是否已解锁标志
        self.is_in_offboard = False  # 是否在offboard模式标志
        self.initial_position_set = False  # 初始位置是否已设置标志
        self.initial_x = 0.0  # 初始X坐标
        self.initial_y = 0.0  # 初始Y坐标
        self.arm_attempts = 0  # 解锁尝试次数
        self.max_arm_attempts = 100  # 最大解锁尝试次数
        self.offboard_attempts = 0  # offboard模式尝试次数
        self.max_offboard_attempts = 50  # 最大offboard模式尝试次数
        self.offboard_mode_engaged = False  # offboard模式是否已启动标志
        self.arming_started = False  # 解锁过程是否已开始标志

    def vehicle_local_position_callback(self, vehicle_local_position):
        """车辆本地位置话题订阅回调函数"""
        self.vehicle_local_position = vehicle_local_position  # 更新本地位置
        
        # 记录初始位置（只在第一次收到位置信息时）
        if not self.initial_position_set and vehicle_local_position.xy_valid:
            self.initial_x = vehicle_local_position.x  # 记录初始X坐标
            self.initial_y = vehicle_local_position.y  # 记录初始Y坐标
            self.initial_position_set = True  # 标记初始位置已设置
            # 可选：打印初始位置信息（已注释以减少日志输出）
            # self.node.get_logger().info(f"UAV{self.uav_id}: Initial position recorded - x: {self.initial_x:.2f}, y: {self.initial_y:.2f}")

    def vehicle_status_callback(self, vehicle_status):
        """车辆状态话题订阅回调函数"""
        self.vehicle_status = vehicle_status  # 更新车辆状态
        # 检查是否已解锁
        self.is_armed = (vehicle_status.arming_state == VehicleStatus.ARMING_STATE_ARMED)
        # 检查是否在offboard模式
        self.is_in_offboard = (vehicle_status.nav_state == VehicleStatus.NAVIGATION_STATE_OFFBOARD)

    def arm(self):
        """发送解锁命令到车辆"""
        # 发送解锁命令，同时设置目标系统为无人机ID
        self.publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM, param1=1.0)
        # 可选：打印解锁命令信息（已注释以减少日志输出）
        # self.node.get_logger().info(f'UAV{self.uav_id}: Arm command sent (target_system={self.uav_id})')

    def disarm(self):
        """发送上锁命令到车辆"""
        self.publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM, param1=0.0)
        # 可选：打印上锁命令信息（已注释以减少日志输出）
        # self.node.get_logger().info(f'UAV{self.uav_id}: Disarm command sent')

    def engage_offboard_mode(self):
        """切换到offboard模式"""
        self.publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_DO_SET_MODE, param1=1.0, param2=6.0)
        # 可选：打印模式切换信息（已注释以减少日志输出）
        # self.node.get_logger().info(f"UAV{self.uav_id}: Switching to offboard mode")

    def land(self):
        """切换到降落模式"""
        self.publish_vehicle_command(VehicleCommand.VEHICLE_CMD_NAV_LAND)
        # 可选：打印降落模式信息（已注释以减少日志输出）
        # self.node.get_logger().info(f"UAV{self.uav_id}: Switching to land mode")

    def publish_offboard_control_heartbeat_signal(self):
        """发布offboard控制模式心跳信号"""
        msg = OffboardControlMode()  # 创建offboard控制模式消息
        msg.position = True  # 启用位置控制
        msg.velocity = False  # 禁用速度控制
        msg.acceleration = False  # 禁用加速度控制
        msg.attitude = False  # 禁用姿态控制
        msg.body_rate = False  # 禁用角速度控制
        msg.timestamp = int(self.node.get_clock().now().nanoseconds / 1000)  # 设置时间戳
        self.offboard_control_mode_publisher.publish(msg)  # 发布消息

    def publish_position_setpoint(self, x: float, y: float, z: float):
        """发布轨迹设定点"""
        msg = TrajectorySetpoint()  # 创建轨迹设定点消息
        msg.position = [x, y, z]  # 设置位置坐标
        msg.yaw = 1.57079  # 设置偏航角（90度）
        msg.timestamp = int(self.node.get_clock().now().nanoseconds / 1000)  # 设置时间戳
        self.trajectory_setpoint_publisher.publish(msg)  # 发布消息
        # 可选：打印位置设定点信息（已注释以减少日志输出）
        # self.node.get_logger().info(f"UAV{self.uav_id}: Publishing position setpoints {[x, y, z]}")

    def publish_vehicle_command(self, command, **params) -> None:
        """发布车辆命令"""
        msg = VehicleCommand()  # 创建车辆命令消息
        msg.command = command  # 设置命令类型
        msg.param1 = params.get("param1", 0.0)  # 设置参数1
        msg.param2 = params.get("param2", 0.0)  # 设置参数2
        msg.param3 = params.get("param3", 0.0)  # 设置参数3
        msg.param4 = params.get("param4", 0.0)  # 设置参数4
        msg.param5 = params.get("param5", 0.0)  # 设置参数5
        msg.param6 = params.get("param6", 0.0)  # 设置参数6
        msg.param7 = params.get("param7", 0.0)  # 设置参数7
        msg.target_system = self.uav_id  # 使用无人机ID作为目标系统
        msg.target_component = 1  # 设置目标组件
        msg.source_system = 1  # 设置源系统
        msg.source_component = 1  # 设置源组件
        msg.from_external = True  # 标记为外部命令
        msg.timestamp = int(self.node.get_clock().now().nanoseconds / 1000)  # 设置时间戳
        self.vehicle_command_publisher.publish(msg)  # 发布消息

    def update(self):
        """更新无人机状态和控制逻辑"""
        self.publish_offboard_control_heartbeat_signal()  # 发布offboard心跳信号

        # 阶段1: 发送offboard模式切换命令
        if not self.offboard_mode_engaged and self.offboard_attempts < self.max_offboard_attempts:
            if self.offboard_setpoint_counter >= 5:  # 发送几个心跳后再切换模式
                self.engage_offboard_mode()  # 切换到offboard模式
                self.offboard_attempts += 1  # 增加尝试次数
                if self.offboard_attempts % 5 == 0:  # 每5次重试打印一次日志
                    self.node.get_logger().info(f"UAV{self.uav_id}: Offboard mode attempt {self.offboard_attempts}/{self.max_offboard_attempts}")

        # 阶段2: 确认进入offboard模式后开始解锁
        if self.is_in_offboard and not self.offboard_mode_engaged:
            self.offboard_mode_engaged = True  # 标记offboard模式已启动
            # 可选：打印成功进入offboard模式信息（已注释以减少日志输出）
            # self.node.get_logger().info(f"UAV{self.uav_id}: Successfully entered offboard mode")
            self.arming_started = True  # 标记解锁过程已开始

        # 阶段3: 发送解锁命令
        if self.offboard_mode_engaged and not self.is_armed:
            # 发送解锁命令（每次循环都发送，但减少日志输出）
            self.arm()  # 发送解锁命令
            self.arm_attempts += 1  # 增加解锁尝试次数
            
            # 只在重要时刻打印日志
            if self.arm_attempts == 1:
                self.node.get_logger().info(f"UAV{self.uav_id}: Started arming process")
            elif self.arm_attempts == 10:
                self.node.get_logger().info(f"UAV{self.uav_id}: Still arming... ({self.arm_attempts} attempts)")

        # 阶段4: 在offboard模式下发送位置设定点（无论是否解锁）
        if self.is_in_offboard:
            # 原地起飞：保持当前位置，只改变高度
            if self.initial_position_set:
                # 使用记录的初始位置
                target_x = self.initial_x
                target_y = self.initial_y
            else:
                # 如果还没有收到位置信息，使用当前位置
                target_x = self.vehicle_local_position.x
                target_y = self.vehicle_local_position.y
            
            self.publish_position_setpoint(target_x, target_y, self.takeoff_height)  # 发布位置设定点
            
            # 基于状态检测起飞情况
            if self.is_armed and not self.has_taken_off:
                # 检查是否到达目标高度
                if self.vehicle_local_position.z <= self.takeoff_height:
                    self.has_taken_off = True  # 标记已起飞
                    self.node.get_logger().info(f"UAV{self.uav_id}: ✅ 成功起飞到目标高度 {self.vehicle_local_position.z:.2f}m")
                else:
                    # 如果已解锁但未起飞，打印起飞状态
                    if self.offboard_setpoint_counter % 30 == 0:  # 每3秒打印一次
                        self.node.get_logger().info(f"UAV{self.uav_id}: 🚀 正在起飞 - 当前高度: {self.vehicle_local_position.z:.2f}m, 目标: {self.takeoff_height}m")
            elif not self.is_armed and self.offboard_mode_engaged:
                # 如果已进入offboard但未解锁，打印等待解锁状态
                if self.offboard_setpoint_counter % 50 == 0:  # 每5秒打印一次
                    self.node.get_logger().info(f"UAV{self.uav_id}: ⏳ 等待解锁 - offboard模式已就绪")

        # 阶段5: 状态监控（已通过阶段4的状态检测处理）
        pass

        if self.offboard_setpoint_counter < 100:  # 增加计数器上限
            self.offboard_setpoint_counter += 1  # 增加计数器

    def has_reached_target(self):
        """检查无人机是否到达目标高度"""
        return self.has_taken_off  # 返回是否已起飞标志


class MultiUAVOffboardControl(Node):
    """多无人机offboard控制节点"""

    def __init__(self, num_uavs: int = 4) -> None:
        """初始化多无人机控制节点"""
        super().__init__('multi_uav_offboard_control')  # 调用父类构造函数
        
        self.num_uavs = num_uavs  # 无人机数量
        self.uav_controllers = []  # 无人机控制器列表
        self.timer_count = 0  # 定时器计数器
        
        # 创建所有无人机的控制器
        for i in range(1, num_uavs + 1):
            uav_controller = UAVController(self, i)  # 创建无人机控制器
            self.uav_controllers.append(uav_controller)  # 添加到控制器列表
            self.get_logger().info(f"Initialized UAV{i} controller")  # 打印初始化信息

        # 创建定时器来发布所有无人机的控制命令
        self.timer = self.create_timer(0.1, self.timer_callback)  # 0.1秒周期
        
        # 可选：打印初始化完成信息（已注释以减少日志输出）
        # self.get_logger().info(f"Multi-UAV offboard control initialized with {num_uavs} UAVs")

    def timer_callback(self) -> None:
        """定时器回调函数，更新所有无人机"""
        all_reached_target = True  # 所有无人机是否到达目标标志
        active_uavs = 0  # 已起飞无人机数量
        armed_uavs = 0  # 已解锁无人机数量
        offboard_uavs = 0  # 在offboard模式无人机数量
        stuck_uavs = []  # 记录卡住的无人机ID列表
        
        for i, uav_controller in enumerate(self.uav_controllers):
            # 为每个无人机添加微小延迟，避免同时发送命令
            if i > 0:
                time.sleep(0.05)  # 50ms延迟
            
            uav_controller.update()  # 更新无人机状态和控制
            
            # 统计状态
            if uav_controller.has_reached_target():
                active_uavs += 1  # 增加已起飞计数
            else:
                all_reached_target = False  # 标记未全部到达
                
                # 检测卡住的无人机（在offboard模式但未解锁超过20秒）
                if uav_controller.is_in_offboard and not uav_controller.is_armed and uav_controller.arm_attempts > 200:
                    stuck_uavs.append(uav_controller.uav_id)  # 添加到卡住列表
            
            if uav_controller.is_armed:
                armed_uavs += 1  # 增加已解锁计数
            if uav_controller.is_in_offboard:
                offboard_uavs += 1  # 增加offboard模式计数
        
        # 定期打印状态汇总（减少频率）
        if self.timer_count % 100 == 0:  # 每10秒打印一次
            self.get_logger().info(f"📊 无人机状态汇总: {active_uavs}/{len(self.uav_controllers)} 已起飞, {armed_uavs}/{len(self.uav_controllers)} 已解锁, {offboard_uavs}/{len(self.uav_controllers)} 在offboard模式")
            
            # 如果有卡住的无人机，报告具体编号
            if stuck_uavs:
                self.get_logger().warning(f"🚨 以下无人机卡在offboard模式但未解锁: {stuck_uavs}")
        
        # 如果所有无人机都到达目标高度，可以执行下一步操作
        if all_reached_target:
            self.get_logger().info("🎉 所有无人机都已成功起飞到目标高度!")
            # 这里可以添加后续的编队飞行或其他任务
            # 例如：self.execute_formation_flight()
        
        self.timer_count += 1  # 增加定时器计数器


def main(args=None) -> None:
    """主函数"""
    print('Starting multi-UAV offboard control node...')  # 打印启动信息
    rclpy.init(args=args)  # 初始化ROS 2
    
    # 可以通过命令行参数指定无人机数量，默认为2架
    import sys  # 导入系统模块
    num_uavs = 4  # 默认无人机数量
    if len(sys.argv) > 1:  # 检查是否有命令行参数
        try:
            num_uavs = int(sys.argv[1])  # 尝试将参数转换为整数
        except ValueError:
            print(f"Invalid number of UAVs: {sys.argv[1]}, using default: 2")  # 参数无效时使用默认值
    
    multi_uav_control = MultiUAVOffboardControl(num_uavs=num_uavs)  # 创建多无人机控制节点
    
    try:
        rclpy.spin(multi_uav_control)  # 运行ROS 2节点
    except KeyboardInterrupt:
        print("\n" + "="*50)  # 打印分隔线
        print("多无人机控制程序被用户中断")  # 打印中断信息
        print("="*50)  # 打印分隔线
        
        # 生成最终状态报告
        success_uavs = []  # 成功起飞无人机列表
        stuck_uavs = []  # 卡住未解锁无人机列表
        failed_uavs = []  # 其他失败无人机列表
        
        for uav_controller in multi_uav_control.uav_controllers:
            if uav_controller.has_reached_target():  # 检查是否成功起飞
                success_uavs.append(uav_controller.uav_id)  # 添加到成功列表
            elif uav_controller.is_in_offboard and not uav_controller.is_armed:  # 检查是否卡住
                stuck_uavs.append(uav_controller.uav_id)  # 添加到卡住列表
            else:
                failed_uavs.append(uav_controller.uav_id)  # 添加到失败列表
        
        print(f"✅ 成功起飞: {len(success_uavs)}架 - UAV{success_uavs}")  # 打印成功起飞信息
        print(f"🚨 卡住未解锁: {len(stuck_uavs)}架 - UAV{stuck_uavs}")  # 打印卡住信息
        print(f"❌ 其他失败: {len(failed_uavs)}架 - UAV{failed_uavs}")  # 打印失败信息
        print("="*50)  # 打印分隔线
    finally:
        multi_uav_control.destroy_node()  # 销毁节点
        rclpy.shutdown()  # 关闭ROS 2


if __name__ == '__main__':
    try:
        main()  # 调用主函数
    except Exception as e:
        print(e)  # 打印异常信息
           