#!/usr/bin/env python3
"""
PX4无人机控制器（ENU坐标系）
提供简单的API接口控制无人机：解锁、起飞、降落、位置控制等
"""

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
from px4_msgs.msg import (
    OffboardControlMode, TrajectorySetpoint, VehicleCommand, 
    VehicleLocalPosition, VehicleStatus
)
import time
import math
import threading


def enu_to_ned(x_enu: float, y_enu: float, z_enu: float):
    """ENU坐标系转NED坐标系
    ENU: X=东, Y=北, Z=上
    NED: X=北, Y=东, Z=下
    """
    x_ned = y_enu      # NED X = ENU Y (北 = 北)
    y_ned = x_enu      # NED Y = ENU X (东 = 东)
    z_ned = -z_enu     # NED Z = -ENU Z (下 = -上)
    return x_ned, y_ned, z_ned


def ned_to_enu(x_ned: float, y_ned: float, z_ned: float):
    """NED坐标系转ENU坐标系
    NED: X=北, Y=东, Z=下
    ENU: X=东, Y=北, Z=上
    """
    x_enu = y_ned      # ENU X = NED Y (东 = 东)
    y_enu = x_ned      # ENU Y = NED X (北 = 北)
    z_enu = -z_ned     # ENU Z = -NED Z (上 = -下)
    return x_enu, y_enu, z_enu


class PX4UAVController(Node):
    """PX4无人机控制器类（ENU坐标系）"""
    
    def __init__(self, uav_id: int = 1, node_name: str = 'px4_uav_controller'):
        """
        初始化无人机控制器
        
        Args:
            uav_id: 无人机ID（默认1）
            node_name: ROS2节点名称
        """
        super().__init__(node_name)
        
        self.uav_id = uav_id
        
        # 配置QoS
        qos_profile = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            history=HistoryPolicy.KEEP_LAST,
            depth=1
        )
        
        # 创建发布者
        ns = f'/uav{uav_id}'
        self.offboard_control_publisher = self.create_publisher(
            OffboardControlMode, f'{ns}/fmu/in/offboard_control_mode', qos_profile)
        self.trajectory_setpoint_publisher = self.create_publisher(
            TrajectorySetpoint, f'{ns}/fmu/in/trajectory_setpoint', qos_profile)
        self.vehicle_command_publisher = self.create_publisher(
            VehicleCommand, f'{ns}/fmu/in/vehicle_command', qos_profile)
        
        # 创建订阅者
        self.vehicle_local_position_subscriber = self.create_subscription(
            VehicleLocalPosition, f'{ns}/fmu/out/vehicle_local_position',
            self._vehicle_local_position_callback, qos_profile)
        self.vehicle_status_subscriber = self.create_subscription(
            VehicleStatus, f'{ns}/fmu/out/vehicle_status',
            self._vehicle_status_callback, qos_profile)
        
        # 初始化变量
        self._vehicle_local_position = VehicleLocalPosition()
        self._vehicle_status = VehicleStatus()
        self._current_target_enu = (0.0, 0.0, 0.0)  # 当前ENU目标位置
        self._current_yaw = 90.0  # 当前偏航角（度）
        self._offboard_mode_active = False
        
        # 创建定时器（10Hz）
        self._timer = self.create_timer(0.1, self._timer_callback)
        
        self.get_logger().info(f'UAV{uav_id} 控制器已初始化')
        
    # ==================== 内部回调 ====================
    def _vehicle_local_position_callback(self, msg):
        """本地位置回调（NED坐标系）"""
        self._vehicle_local_position = msg
        
    def _vehicle_status_callback(self, msg):
        """车辆状态回调"""
        self._vehicle_status = msg
        
    def _timer_callback(self):
        """定时器回调（10Hz）"""
        self._publish_offboard_control_mode()
        
        # 如果offboard模式激活，持续发送目标位置
        if self._offboard_mode_active:
            x_ned, y_ned, z_ned = enu_to_ned(
                self._current_target_enu[0],
                self._current_target_enu[1],
                self._current_target_enu[2]
            )
            self._publish_position_setpoint(
                x_ned, y_ned, z_ned,
                math.radians(self._current_yaw)
            )
            
    # ==================== 内部发布方法 ====================
    def _publish_offboard_control_mode(self):
        """发布offboard控制模式"""
        msg = OffboardControlMode()
        msg.position = True
        msg.velocity = False
        msg.acceleration = False
        msg.attitude = False
        msg.body_rate = False
        msg.timestamp = int(self.get_clock().now().nanoseconds / 1000)
        self.offboard_control_publisher.publish(msg)
        
    def _publish_position_setpoint(self, x_ned: float, y_ned: float, z_ned: float, yaw_rad: float):
        """发布位置设定点（NED坐标系）"""
        msg = TrajectorySetpoint()
        msg.position = [x_ned, y_ned, z_ned]
        msg.yaw = yaw_rad
        msg.timestamp = int(self.get_clock().now().nanoseconds / 1000)
        self.trajectory_setpoint_publisher.publish(msg)
        
    def _publish_vehicle_command(self, command: int, **params):
        """发布车辆命令"""
        msg = VehicleCommand()
        msg.command = command
        msg.param1 = params.get("param1", 0.0)
        msg.param2 = params.get("param2", 0.0)
        msg.param3 = params.get("param3", 0.0)
        msg.param4 = params.get("param4", 0.0)
        msg.param5 = params.get("param5", 0.0)
        msg.param6 = params.get("param6", 0.0)
        msg.param7 = params.get("param7", 0.0)
        msg.target_system = self.uav_id
        msg.target_component = 1
        msg.source_system = 1
        msg.source_component = 1
        msg.from_external = True
        msg.timestamp = int(self.get_clock().now().nanoseconds / 1000)
        self.vehicle_command_publisher.publish(msg)
        
    # ==================== 公共API方法 ====================
    def arm(self):
        """解锁无人机"""
        self._publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM, param1=1.0)
        self.get_logger().info(f'UAV{self.uav_id}: 发送解锁命令')
        
    def disarm(self):
        """上锁无人机"""
        self._publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_COMPONENT_ARM_DISARM, param1=0.0)
        self.get_logger().info(f'UAV{self.uav_id}: 发送上锁命令')
        
    def engage_offboard_mode(self):
        """切换到offboard模式"""
        self._publish_vehicle_command(
            VehicleCommand.VEHICLE_CMD_DO_SET_MODE, param1=1.0, param2=6.0)
        self.get_logger().info(f'UAV{self.uav_id}: 切换到OFFBOARD模式')
        
    def land(self):
        """切换到降落模式"""
        self._offboard_mode_active = False
        self._publish_vehicle_command(VehicleCommand.VEHICLE_CMD_NAV_LAND)
        self.get_logger().info(f'UAV{self.uav_id}: 切换到降落模式')
        
    def takeoff(self, height: float = 3.0, timeout: float = 30.0):
        """
        起飞到指定高度
        
        Args:
            height: 目标高度（米，ENU坐标系，正值）
            timeout: 超时时间（秒）
        """
        self.get_logger().info(f'UAV{self.uav_id}: 准备起飞到 {height} 米高度')
        
        # 检查当前状态
        # self.get_logger().info(f'UAV{self.uav_id}: 当前飞行模式: {self._vehicle_status.nav_state}')
        # self.get_logger().info(f'UAV{self.uav_id}: 当前是否解锁: {self.is_armed()}')
        # self.get_logger().info(f'UAV{self.uav_id}: 当前高度: {-self._vehicle_local_position.z:.2f}米')
        
        # 获取当前位置
        current_x_ned = self._vehicle_local_position.x
        current_y_ned = self._vehicle_local_position.y
        
        # NED转ENU
        current_x_enu = current_y_ned
        current_y_enu = current_x_ned
        
        # 设置目标位置
        self._current_target_enu = (current_x_enu, current_y_enu, height)
        self._current_yaw = 90.0
        
        # 激活offboard模式
        self._offboard_mode_active = True
        
        # 准备阶段：持续发送setpoint一段时间，让PX4识别offboard控制源
        self.get_logger().info(f'UAV{self.uav_id}: 发送offboard控制信号（准备阶段）...')
        prepare_time = 3.0  # 准备3秒
        prepare_start = time.time()
        while (time.time() - prepare_start) < prepare_time:
            rclpy.spin_once(self, timeout_sec=0.1)
            time.sleep(0.1)
        
        self.get_logger().info(f'UAV{self.uav_id}: 准备完成，检查状态...')
        
        # 等待无人机真正解锁和进入offboard模式
        self.get_logger().info(f'UAV{self.uav_id}: 等待无人机解锁和进入offboard模式...')
        wait_start = time.time()
        max_wait = 10.0  # 最多等待10秒
        
        while (time.time() - wait_start) < max_wait:
            rclpy.spin_once(self, timeout_sec=0.1)
            
            is_armed = self.is_armed()
            is_offboard = (self._vehicle_status.nav_state == VehicleStatus.NAVIGATION_STATE_OFFBOARD)
            
            if is_armed and is_offboard:
                self.get_logger().info(f'UAV{self.uav_id}: 无人机已解锁且在offboard模式！')
                break
            elif int(time.time() - wait_start) % 2 == 0:
                self.get_logger().info(f'UAV{self.uav_id}: 等待中... 解锁: {is_armed}, Offboard: {is_offboard}')
            
            time.sleep(0.1)
        else:
            self.get_logger().warn(f'UAV{self.uav_id}: 等待解锁和offboard模式超时！')
            self.get_logger().warn(f'UAV{self.uav_id}: 当前状态 - 解锁: {self.is_armed()}, 飞行模式: {self._vehicle_status.nav_state}')
            return False
        
        # 记录起飞前的高度
        start_height = -self._vehicle_local_position.z
        self.get_logger().info(f'UAV{self.uav_id}: 起飞前高度: {start_height:.2f}米')
        self.get_logger().info(f'UAV{self.uav_id}: 开始起飞...')
        
        # 等待到达目标高度
        start_time = time.time()
        while (time.time() - start_time) < timeout:
            # 处理ROS回调
            rclpy.spin_once(self, timeout_sec=0.1)
            
            current_height_enu = -self._vehicle_local_position.z  # NED Z为负，转换后为正值
            
            # 检查高度是否有变化（至少上升0.3米）
            if current_height_enu >= start_height + 0.3:
                self.get_logger().info(f'UAV{self.uav_id}: 检测到起飞，当前高度: {current_height_enu:.2f}米')
            
            # 检查是否达到目标高度
            if current_height_enu >= height * 0.9:
                self.get_logger().info(f'UAV{self.uav_id}: 起飞完成！当前高度: {current_height_enu:.2f}米')
                return True
                
            time.sleep(0.1)
            
        self.get_logger().warn(f'UAV{self.uav_id}: 起飞超时')
        self.get_logger().warn(f'UAV{self.uav_id}: 最终高度: {-self._vehicle_local_position.z:.2f}米')
        return False
        
    def goto_enu(self, x: float, y: float, z: float, yaw: float = 90.0, timeout: float = 30.0):
        """
        飞到ENU坐标系目标位置
        
        Args:
            x: 东向位置（米）
            y: 北向位置（米）
            z: 高度（米，正值）
            yaw: 偏航角（度，0=北，90=东，180=南，270=西）
            timeout: 超时时间（秒）
        """
        self.get_logger().info(f'UAV{self.uav_id}: 飞往ENU坐标 [{x:.2f}, {y:.2f}, {z:.2f}], 偏航角: {yaw:.1f}°')
        
        # 设置目标位置
        self._current_target_enu = (x, y, z)
        self._current_yaw = yaw
        
        # 激活offboard模式
        self._offboard_mode_active = True
        
        # 等待到达目标位置
        start_time = time.time()
        last_progress_time = 0
        while (time.time() - start_time) < timeout:
            # 处理ROS回调
            rclpy.spin_once(self, timeout_sec=0.1)
            
            # 当前位置（NED）
            current_x_ned = self._vehicle_local_position.x
            current_y_ned = self._vehicle_local_position.y
            current_z_ned = self._vehicle_local_position.z
            
            # 目标位置（NED）
            target_x_ned, target_y_ned, target_z_ned = enu_to_ned(x, y, z)
            
            # 计算距离
            distance = math.sqrt(
                (current_x_ned - target_x_ned)**2 +
                (current_y_ned - target_y_ned)**2 +
                (current_z_ned - target_z_ned)**2
            )
            
            # 每2秒打印一次进度
            current_time = time.time()
            if current_time - last_progress_time >= 2.0:
                # 当前ENU位置
                current_x_enu, current_y_enu, current_z_enu = self.get_position_enu()
                self.get_logger().info(f'UAV{self.uav_id}: 当前位置: ENU[{current_x_enu:.2f}, {current_y_enu:.2f}, {current_z_enu:.2f}], 距离目标: {distance:.2f}米')
                last_progress_time = current_time
            
            if distance < 0.5:  # 到达目标点（50cm内）
                self.get_logger().info(f'UAV{self.uav_id}: 到达目标位置！')
                return True
                
            time.sleep(0.1)
            
        self.get_logger().warn(f'UAV{self.uav_id}: 到达目标位置超时')
        # 最终位置信息
        current_x_enu, current_y_enu, current_z_enu = self.get_position_enu()
        self.get_logger().warn(f'UAV{self.uav_id}: 最终位置: ENU[{current_x_enu:.2f}, {current_y_enu:.2f}, {current_z_enu:.2f}]')
        return False
        
    # ==================== 机体坐标系控制方法 ====================
    def move_body_frame(self, x: float = 0.0, y: float = 0.0, z: float = 0.0, yaw: float = 0.0, timeout: float = 10.0):
        """
        机体坐标系下移动无人机 - 简洁API
        
        Args:
            x: 前进距离（米，正值向前）
            y: 右移距离（米，正值向右）
            z: 上升距离（米，正值向上）
            yaw: 偏航角变化（度，正值向右转）
            timeout: 超时时间（秒，默认10秒）
        """
        return self._move_body_frame_internal(forward=x, right=y, up=z, yaw_deg=yaw, timeout=timeout)
        
    def move_body_frame_xyz(self, x: float = 0.0, y: float = 0.0, z: float = 0.0, yaw: float = 0.0, timeout: float = 10.0):
        """
        机体坐标系下移动无人机 - XYZ坐标API
        
        Args:
            x: 前进距离（米，正值向前）
            y: 右移距离（米，正值向右）
            z: 上升距离（米，正值向上）
            yaw: 偏航角变化（度，正值向右转，可选）
            timeout: 超时时间（秒，默认10秒，可选）
        """
        return self._move_body_frame_internal(forward=x, right=y, up=z, yaw_deg=yaw, timeout=timeout)
        
    def _move_body_frame_internal(self, forward: float = 0.0, right: float = 0.0, up: float = 0.0, yaw_deg: float = 0.0, timeout: float = 10.0):
        """
        机体坐标系下移动无人机 - 内部实现
        
        Args:
            forward: 前进距离（米，正值向前）
            right: 右移距离（米，正值向右）
            up: 上升距离（米，正值向上）
            yaw_deg: 偏航角变化（度，正值向右转）
            timeout: 超时时间（秒）
        """
        # 获取当前位置
        current_x_enu, current_y_enu, current_z_enu = self.get_position_enu()
        current_yaw = self._current_yaw
        
        # 计算机体坐标系下的目标位置
        # 机体坐标系：X=前，Y=右，Z=上
        # 需要根据当前偏航角转换到世界坐标系
        
        # 将机体坐标系位移转换到世界坐标系
        yaw_rad = math.radians(current_yaw)
        
        # 机体坐标系下的位移向量
        body_dx = forward
        body_dy = right
        
        # 转换到世界坐标系（ENU）
        world_dx = body_dx * math.cos(yaw_rad) - body_dy * math.sin(yaw_rad)
        world_dy = body_dx * math.sin(yaw_rad) + body_dy * math.cos(yaw_rad)
        
        # 计算目标位置
        target_x = current_x_enu + world_dx
        target_y = current_y_enu + world_dy
        target_z = current_z_enu + up
        target_yaw = current_yaw + yaw_deg
        
        self.get_logger().info(f'UAV{self.uav_id}: 机体坐标系移动 - 前:{forward}m, 右:{right}m, 上:{up}m, 偏航:{yaw_deg}°')
        self.get_logger().info(f'UAV{self.uav_id}: 从ENU[{current_x_enu:.2f}, {current_y_enu:.2f}, {current_z_enu:.2f}] -> ENU[{target_x:.2f}, {target_y:.2f}, {target_z:.2f}]')
        
        # 设置目标位置
        self._current_target_enu = (target_x, target_y, target_z)
        self._current_yaw = target_yaw
        self._offboard_mode_active = True
        
        # 等待到达目标位置
        start_time = time.time()
        while (time.time() - start_time) < timeout:
            # 处理ROS回调
            rclpy.spin_once(self, timeout_sec=0.1)
            
            # 获取当前位置
            current_x, current_y, current_z = self.get_position_enu()
            
            # 计算到目标位置的距离
            distance = math.sqrt(
                (current_x - target_x)**2 +
                (current_y - target_y)**2 +
                (current_z - target_z)**2
            )
            
            if distance < 0.3:  # 30cm精度
                self.get_logger().info(f'UAV{self.uav_id}: 机体坐标系移动完成！')
                return True
                
            time.sleep(0.1)
        
        self.get_logger().warn(f'UAV{self.uav_id}: 机体坐标系移动超时')
        return False
        
    def move_forward(self, distance: float, timeout: float = 10.0):
        """向前移动指定距离"""
        return self._move_body_frame_internal(forward=distance, timeout=timeout)
        
    def move_backward(self, distance: float, timeout: float = 10.0):
        """向后移动指定距离"""
        return self._move_body_frame_internal(forward=-distance, timeout=timeout)
        
    def move_right(self, distance: float, timeout: float = 10.0):
        """向右移动指定距离"""
        return self._move_body_frame_internal(right=distance, timeout=timeout)
        
    def move_left(self, distance: float, timeout: float = 10.0):
        """向左移动指定距离"""
        return self._move_body_frame_internal(right=-distance, timeout=timeout)
        
    def move_up(self, distance: float, timeout: float = 10.0):
        """向上移动指定距离"""
        return self._move_body_frame_internal(up=distance, timeout=timeout)
        
    def move_down(self, distance: float, timeout: float = 10.0):
        """向下移动指定距离"""
        return self._move_body_frame_internal(up=-distance, timeout=timeout)
        
    def turn_right(self, angle_deg: float, timeout: float = 5.0):
        """向右转指定角度"""
        return self._move_body_frame_internal(yaw_deg=angle_deg, timeout=timeout)
        
    def turn_left(self, angle_deg: float, timeout: float = 5.0):
        """向左转指定角度"""
        return self._move_body_frame_internal(yaw_deg=-angle_deg, timeout=timeout)
        
    def set_yaw(self, yaw_deg: float):
        """
        设置偏航角
        
        Args:
            yaw_deg: 偏航角（度，0=北，90=东，180=南，270=西）
        """
        self._current_yaw = yaw_deg
        self.get_logger().info(f'UAV{self.uav_id}: 设置偏航角 {yaw_deg:.1f}°')
        
    def get_position_enu(self):
        """
        获取当前ENU位置
        
        Returns:
            tuple: (x, y, z) ENU坐标
        """
        x_ned = self._vehicle_local_position.x
        y_ned = self._vehicle_local_position.y
        z_ned = self._vehicle_local_position.z
        
        return ned_to_enu(x_ned, y_ned, z_ned)
        
    def get_status(self):
        """
        获取无人机状态
        
        Returns:
            dict: 包含armed, flight_mode, position等信息
        """
        is_armed = self._vehicle_status.arming_state == VehicleStatus.ARMING_STATE_ARMED
        flight_mode = self._vehicle_status.nav_state
        
        x_enu, y_enu, z_enu = self.get_position_enu()
        
        return {
            'armed': is_armed,
            'flight_mode': flight_mode,
            'position_enu': (x_enu, y_enu, z_enu),
            'height_enu': -self._vehicle_local_position.z,  # NED Z为负，转换后为正值
            'yaw': self._current_yaw
        }
        
    def is_armed(self):
        """检查是否已解锁"""
        return self._vehicle_status.arming_state == VehicleStatus.ARMING_STATE_ARMED
        
    def get_flight_mode(self):
        """获取当前飞行模式"""
        return self._vehicle_status.nav_state
        
    def wait_for_connection(self, timeout: float = 30.0):
        """
        等待无人机连接
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            bool: 是否成功连接
        """
        self.get_logger().info(f'UAV{self.uav_id}: 等待无人机连接...')
        self.get_logger().info(f'UAV{self.uav_id}: 订阅话题: /uav{self.uav_id}/fmu/out/vehicle_status')
        
        start_time = time.time()
        while (time.time() - start_time) < timeout:
            # 处理ROS回调
            rclpy.spin_once(self, timeout_sec=0.1)
            
            # 检查是否有数据更新
            if self._vehicle_status.nav_state != 0:  # 0是未定义状态
                self.get_logger().info(f'UAV{self.uav_id}: 无人机已连接')
                self.get_logger().info(f'UAV{self.uav_id}: 当前飞行模式: {self._vehicle_status.nav_state}')
                return True
                
            # 每5秒打印一次状态
            if int(time.time() - start_time) % 5 == 0 and int(time.time() - start_time) > 0:
                self.get_logger().info(f'UAV{self.uav_id}: 等待中... {int(time.time() - start_time)}s')
                self.get_logger().info(f'UAV{self.uav_id}: 当前nav_state: {self._vehicle_status.nav_state}')
            
            time.sleep(0.1)
            
        self.get_logger().warn(f'UAV{self.uav_id}: 等待连接超时')
        self.get_logger().warn(f'UAV{self.uav_id}: 请检查:')
        self.get_logger().warn(f'UAV{self.uav_id}: 1. 无人机是否已启动')
        self.get_logger().warn(f'UAV{self.uav_id}: 2. 话题 /uav{self.uav_id}/fmu/out/vehicle_status 是否有数据')
        self.get_logger().warn(f'UAV{self.uav_id}: 3. 运行: ros2 topic echo /uav{self.uav_id}/fmu/out/vehicle_status')
        return False


# ==================== 使用示例 ====================
