"""
BLINK.py - MAVLink飞控通信与控制模块
功能说明：
1. 建立与飞控的MAVLink通信连接（串口/TCP）
2. 实时接收并解析飞控状态数据（姿态、GPS、电池状态等）
3. 支持航点任务上传与自动飞行控制
4. 提供飞行模式切换和参数配置接口
5. 多线程处理数据接收，保证主线程控制流畅性

关键特性：
- 采用MAVLink v2.0协议，支持ArduPilot/PX4
- 消息频率动态配置优化带宽
- 线程安全的状态数据访问机制
- 航点任务支持跳跃指令循环
- 完备的错误处理和重连机制
作者: YZY & Shiqi
日期: 2025-07-21
版本: 0.1.5
"""

import time
from pymavlink import mavutil
import pymavlink.dialects.v20.all as dialect
import threading
import dronekit


class BLINK:
    """飞控通信主类"""

    def __init__(self, serial_port="COM5", baud_rate=57600):
        """初始化飞控连接参数
        Args:
            serial_port (str): 串口名称，默认为COM5
            baud_rate (int): 波特率，默认为57600
        """
        self.serial_port = serial_port
        self.baud_rate = baud_rate
        self.master = None
        self.lock = threading.Lock()  # 线程锁
        # 默认命令消息模板
        self.default_msg = {
            'target_system': 1,
            'target_component': 0,
            'command': None,
            'confirmation': 0,
            'param1': 0,
            'param2': 0,
            'param3': 0,
            'param4': 0,
            'param5': 0,
            'param6': 0,
            'param7': 0
        }

        self.latest_state = {}  # 存储飞控最新状态
        self.receiver_thread = None  # 接收线程
        self.mission_seq_request = -1
        self.mission_ack = [-1, -1]

        # ArduPilot飞行模式映射表
        self.ID_TO_MODES = {
            0: "MANUAL",  # 未解锁
            1: "CIRCLE",
            2: "STABILIZE",
            3: "TRAINING",
            4: "ACRO",
            5: "FBWA",
            6: "FBWB",
            7: "CRUISE",
            8: "AUTOTUNE",
            10: "AUTO",  # 自动
            11: "RTL",
            12: "LOITER",
            13: "TAKEOFF",
            14: "AVOID_ADSB",
            15: "GUIDED",
            16: "INITIALIZING",
            17: "QSTABILIZE",
            18: "QHOVER",
            19: "QLOITER",
            20: "QLAND",
            21: "QRTL",
            22: "QAUTOTUNE",
            23: "QACRO",
            24: "THERMAL",
            25: "LOITER_ALT_QLAND",
            26: "AUTOLAND"
        }

    def get_para(self, para_name):
        """获取飞控参数最新值
        Args:
            para_name (str): 参数名称
        Returns:
            参数值，若不存在返回None
        """
        return self.latest_state[para_name]

    def wait_mission_ack_and_reset(self, timeout):
        mission_ack = []
        for i in range(timeout):
            if self.mission_ack != [-1, -1]:
                mission_ack = self.mission_ack
                self.mission_ack = [-1, -1]
                return mission_ack
            time.sleep(0.01)

        return "wait mission_ack timeout"

    # =========================================================================
    # ------------------------- 消息频率设置方法 ------------------------------
    # =========================================================================

    def set_message_send_frequency(self, msg_id, interval):
        """设置消息发送频率（通用方法）
        Args:
            msg_id (int): MAVLink消息ID
            interval (int): 发送间隔（Hz正数或-1表示禁用）
        """
        # 构建设置消息频率的命令
        msg = dialect.MAVLink_command_long_message(target_system=self.master.target_system,
                                                   target_component=self.master.target_component,
                                                   command=dialect.MAV_CMD_SET_MESSAGE_INTERVAL,
                                                   confirmation=0,
                                                   param1=msg_id,
                                                   param2=interval,
                                                   param3=0,
                                                   param4=0,
                                                   param5=0,
                                                   param6=0,
                                                   param7=0)
        self.master.mav.send(msg)

    def set_gps_data_rate(self, rate):
        """设置GPS数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_global_position_int_message.id, rate)

    def set_attitude_data_rate(self, rate):
        """设置姿态数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_attitude_message.id, rate)

    def set_sys_status_data_rate(self, rate):
        """设置系统状态数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_sys_status_message.id, rate)

    def set_mission_current_data_rate(self, rate):
        """设置任务当前数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_mission_current_message.id, rate)

    def set_rc_channels_data_rate(self, rate):
        """设置遥控通道数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_rc_channels_message.id, rate)

    def set_system_time_data_rate(self, rate):
        """设置系统时间数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_system_time_message.id, rate)

    def set_ahrs_data_rate(self, rate):
        """设置AHRS数据发送频率"""
        self.set_message_send_frequency(dialect.MAVLink_ahrs_message.id, rate)

    def message_frequency_init(self):
        """初始化消息频率，禁用不必要的数据流

        该函数禁用多种不常用消息以减少通信负载
        """
        # 禁用各种不必要消息的数据流
        # 需要配置频率的MAVLink消息ID列表
        message_ids = [
            dialect.MAVLink_vfr_hud_message.id,
            dialect.MAVLink_aoa_ssa_message.id,
            dialect.MAVLink_power_status_message.id,
            dialect.MAVLink_meminfo_message.id,
            dialect.MAVLink_servo_output_raw_message.id,
            dialect.MAVLink_raw_imu_message.id,
            dialect.MAVLink_scaled_imu2_message.id,
            dialect.MAVLink_scaled_imu3_message.id,
            dialect.MAVLink_scaled_pressure_message.id,
            dialect.MAVLink_scaled_pressure2_message.id,
            dialect.MAVLink_gps_raw_int_message.id,
            dialect.MAVLink_wind_message.id,
            dialect.MAVLink_ekf_status_report_message.id,
            dialect.MAVLink_vibration_message.id,
            dialect.MAVLink_battery_status_message.id,
            dialect.MAVLink_mcu_status_message.id
        ]

        for msg_id in message_ids:
            self.set_message_send_frequency(msg_id, -1)
            # 阻塞等待命令确认消息
            self.master.recv_match(type=dialect.MAVLink_command_ack_message.msgname, blocking=True)

    # =========================================================================
    # --------------------- 连接管理与基础通信方法 -----------------------------
    # =========================================================================

    def connect(self):
        """连接飞控并等待心跳包

        尝试建立连接直到成功为止
        """
        while True:
            print(f"[MAVLink] Connecting to flight controller on {self.serial_port} at {self.baud_rate} baud...")
            try:
                # 创建MAVLink连接
                self.master = mavutil.mavlink_connection(self.serial_port, baud=self.baud_rate, source_system=1)
                # 等待心跳包确认连接
                self.master.wait_heartbeat(timeout=5)
                print("[MAVLink] Connected to flight controller")
                break
            except Exception as e:
                print(f"[Error] Connection failed: {e}")
                time.sleep(5)
        # 更新默认消息的目标系统参数
        self.default_msg['target_system'] = self.master.target_system
        self.default_msg['target_component'] = self.master.target_component

    def make_msg(self, **kwargs):
        """生成 Mavlink 消息
        Args:
            **kwargs: 需要修改的参数，其他参数使用默认值
            - target_system: 目标系统ID
            - target_component: 目标组件ID
            - command: 命令ID
            - confirmation: 确认标志
            - param1, param2, param3, param4, param5, param6, param7: 命令参数

        Returns:
            dict: Mavlink 消息字典，使用**解包可以直接发送
        """
        # 使用默认参数创建消息副本
        msg = self.default_msg.copy()
        # 更新用户提供的参数
        msg.update(kwargs)
        return msg

    def send_text(self, text, severity=6):
        byte_text = text.encode('utf-8')
        msg = dialect.MAVLink_statustext_message(severity, byte_text)
        self.master.mav.send(msg)
        # print("send text successful")

    '''def request_data_from_fc(self, msg_id, rate_hz):
        """请求指定消息类型的数据，注：若未使用关闭请求，则飞控下次会一直发送请求
        Args:
            msg_id (int): 消息类型ID, 参见 MAVLink 文档
            rate_hz (int): 正数表示请求的频率，否则表示取消请求
        """
        # 构建请求数据命令
        msg = self.make_msg(command=mavutil.mavlink.MAV_CMD_SET_MESSAGE_INTERVAL,
                            param1=msg_id,
                            param2=int(1e6 / rate_hz) if rate_hz > 0 else -1)
        # 发送请求命令
        self.master.mav.command_long_send(**msg)'''

    # =========================================================================
    # ------------------------ 参数管理方法 -----------------------------------
    # =========================================================================
    def get_param_from_fc(self, param_id):
        """读取一次指定飞控参数
        Args:
            param_id (str): 参数名称(不多于16个字符)
        """
        self.master.mav.param_request_read_send(self.default_msg['target_system'],
                                                self.default_msg['target_component'],
                                                param_id.encode('utf-8'), -1)

    def set_param_into_fc(self, param_id, param_value):
        """
        设置飞控参数
        Args:
            param_id (str): 参数名称(不多于16个字符)
            param_value (float): 参数值
        """
        # 确保参数已存在本地缓存中
        if param_id not in self.latest_state:
            # 请求一次参数，等待其回传
            self.get_param_from_fc(param_id)

            # 最多等待1秒，每100ms检查一次是否收到参数
            for _ in range(10):
                time.sleep(0.1)
                if param_id in self.latest_state:
                    break
            else:
                print(f"[Warning] 参数 {param_id} 未收到回传，无法设置")
                return

        param_type = self.latest_state[param_id]['type']
        # 发送设置参数命令
        self.master.mav.param_set_send(
            self.default_msg['target_system'],
            self.default_msg['target_component'],
            param_id.encode('utf-8'),
            param_value,
            param_type
        )

    # =========================================================================
    # ------------------------ 状态监控方法 -----------------------------------
    # =========================================================================
    def receiver(self):
        """持续接收飞控消息，更新最新状态字典
        持续接收飞控消息并更新状态
        在独立线程中运行，解析各类消息并更新latest_state字典
        """
        while True:
            # 接收一条消息, 超时时间设为0.2秒
            msg = self.master.recv_match(blocking=True, timeout=0.2)
            if msg:
                msg_type = msg.get_type()
                # print(msg_type)
                # 使用线程锁保护状态字典
                with self.lock:
                    if msg_type == 'ATTITUDE':
                        # 姿态数据更新
                        self.latest_state['roll'] = msg.roll
                        self.latest_state['pitch'] = msg.pitch
                        self.latest_state['yaw'] = msg.yaw
                    elif msg_type == 'GLOBAL_POSITION_INT':
                        # GPS数据更新
                        self.latest_state['latitude'] = msg.lat / 1e7
                        self.latest_state['longitude'] = msg.lon / 1e7
                        self.latest_state['altitude'] = msg.relative_alt
                    elif msg_type == 'SYSTEM_TIME':
                        # 系统时间数据更新
                        self.latest_state['sys_time'] = msg.time_unix_usec
                    elif msg_type == 'PARAM_VALUE':
                        # 参数数据更新
                        self.latest_state[msg.param_id] = {'value': msg.param_value, 'type': msg.param_type}
                    elif msg_type == 'HEARTBEAT' and msg.type == 1:
                        # 心跳包（飞行模式和系统状态）数据更新
                        self.latest_state['flight_mode'] = self.ID_TO_MODES.get(msg.custom_mode,
                                                                                f"未知模式({msg.custom_mode})")
                        self.latest_state['system_status'] = msg.system_status
                    elif msg_type == dialect.MAVLink_mission_request_message.msgname:
                        msg = msg.to_dict()
                        print(msg)
                        self.mission_seq_request = msg["seq"]
                    elif msg_type == dialect.MAVLink_mission_ack_message.msgname:
                        msg = msg.to_dict()
                        print(msg)
                        self.mission_ack[0] = msg["type"]
                        self.mission_ack[1] = msg["mission_type"]

                    else:
                        pass
                        # print(f"[MAVLink] Received message: {msg_type} - {msg.to_dict()}")

    def start_receiver(self):
        """启动接收线程"""
        self.receiver_thread = threading.Thread(target=self.receiver, daemon=True)
        self.receiver_thread.start()

    # =========================================================================
    # ------------------------ 任务管理方法 -----------------------------------
    # =========================================================================

    def upload_mission(self, waypoints):
        """上传航点任务
            Args:
                waypoints (list): 航点列表，每个航点格式为(命令, param1, param2, param3, param4, lat, lon, alt)
        """
        # 清除现有任务
        msg = self.master.mav.mission_clear_all_send(target_system=self.master.target_system,
                                                     target_component=self.master.target_component, )
        # 等待清除确认
        # while True:
        #    ack = self.master.recv_match(type='MISSION_ACK', blocking=True, timeout=5)
        #    if ack and ack.type == mavutil.mavlink.MAV_MISSION_ACCEPTED:
        #        print("现有任务已清除")
        #        break
        mission_count_msg = dialect.MAVLink_mission_count_message(target_system=self.master.target_system,
                                                                  target_component=self.master.target_component,
                                                                  count=len(waypoints) + 1,
                                                                  mission_type=dialect.MAV_MISSION_TYPE_MISSION)
        self.master.mav.send(mission_count_msg)
        print('send mission count')

        self.wait_mission_ack_and_reset(300)
        print(self.mission_ack)

        last_seq = -1
        # 上传航点
        while True:

            # 处理航点请求
            if self.mission_seq_request != last_seq:
                seq = self.mission_seq_request
                last_seq = self.mission_seq_request

                if seq == 0:
                    # 上传第一个航点HOME点
                    upload_msg = dialect.MAVLink_mission_item_int_message(
                        target_system=self.master.target_system,
                        target_component=self.master.target_component,
                        seq=seq,
                        frame=dialect.MAV_FRAME_GLOBAL,
                        command=dialect.MAV_CMD_NAV_WAYPOINT,
                        current=0,
                        autocontinue=0,
                        param1=0,
                        param2=0,
                        param3=0,
                        param4=0,
                        x=0,
                        y=0,
                        z=0,
                        mission_type=dialect.MAV_MISSION_TYPE_MISSION
                    )
                else:
                    # 普通航点
                    upload_msg = dialect.MAVLink_mission_item_int_message(
                        target_system=self.master.target_system,
                        target_component=self.master.target_component,
                        seq=seq,
                        frame=dialect.MAV_FRAME_GLOBAL_RELATIVE_ALT,
                        command=waypoints[seq - 1][0],
                        current=0,
                        autocontinue=0,
                        param1=waypoints[seq - 1][1],
                        param2=waypoints[seq - 1][2],
                        param3=waypoints[seq - 1][3],
                        param4=waypoints[seq - 1][4],
                        x=int(waypoints[seq - 1][5] * 1e7),
                        y=int(waypoints[seq - 1][6] * 1e7),
                        z=waypoints[seq - 1][7],
                        mission_type=dialect.MAV_MISSION_TYPE_MISSION
                    )

                    print(f"上传航点 {seq}/{len(waypoints)}")
                    # 发送上传航点命令
                self.master.mav.send(upload_msg)
            # 处理航点回传

            if self.mission_ack[0] == dialect.MAV_MISSION_ACCEPTED and self.mission_ack[
                1] == dialect.MAV_MISSION_TYPE_MISSION:
                self.wait_mission_ack_and_reset(200)
                print('mission upload successful')
                break

    # =========================================================================
    # ------------------------ 飞行控制方法 -----------------------------------
    # =========================================================================

    def flight_mode_set(self, mode):
        """设置飞行模式
            Args:
                mode (int): 飞行模式代码（参考ArduPilot文档）
        """
        # 构建设置模式命令
        set_mode_msg = dialect.MAVLink_command_long_message(target_system=self.master.target_system,
                                                            target_component=self.master.target_component,
                                                            command=dialect.MAV_CMD_DO_SET_MODE,
                                                            confirmation=0,
                                                            param1=dialect.MAV_MODE_FLAG_CUSTOM_MODE_ENABLED,
                                                            param2=mode,
                                                            param3=0,
                                                            param4=0,
                                                            param5=0,
                                                            param6=0,
                                                            param7=0)
        self.master.mav.send(set_mode_msg)
        # 等待命令确认


if __name__ == "__main__":
    # uav = BLINK(serial_port="COM11", baud_rate=57600)  # 连接飞控,端口自己从电脑设备管理器读取,注意不要跟MissionPlanner占一个端口
    uav = BLINK(serial_port="tcp:localhost:5762", baud_rate=57600)  # 跑本地模拟的
    # uav = BLINK(serial_port="udpout:localhost:14500", baud_rate=57600)
    uav.connect()
    uav.message_frequency_init()
    uav.start_receiver()
    time.sleep(1)
    mission = [
        (dialect.MAV_CMD_NAV_TAKEOFF, 0, 0, 0, 0, 0, 0, 30),
        (dialect.MAV_CMD_NAV_WAYPOINT, 0, 0, 0, 0, 40.2275472, 116.2705207, 50),
        (dialect.MAV_CMD_NAV_WAYPOINT, 0, 0, 0, 0, 40.2252372, 116.2646627, 70),
        (dialect.MAV_CMD_NAV_SCRIPT_TIME, 16, 0, 10, 0, 0, 0, 0),
        (dialect.MAV_CMD_NAV_WAYPOINT, 0, 0, 0, 0, 40.2212887, 116.2678814, 50),
        (dialect.MAV_CMD_NAV_WAYPOINT, 0, 0, 0, 0, 40.2237627, 116.2732458, 70),
        (dialect.MAV_CMD_DO_JUMP, 1, 100, 0, 0, 40.2237627, 116.2732458, 70)
    ]
    # uav.upload_mission(mission)
    # uav.send_text('ababa')
    # 主线程保持运行
    while True:
        time.sleep(1)  # 保持主线程活跃

'''
ATTITUDE
VFR_HUD             X
AOA_SSA             X
GLOBAL_POSITION_INT
SYS_STATUS
POWER_STATUS        X
MEMINFO             X
MISSION_CURRENT
SERVO_OUTPUT_RAW    X
RC_CHANNELS
RAW_IMU             X
SCALED_IMU2         X
SCALED_IMU3         X
SCALED_PRESSURE     X
SCALED_PRESSURE2    X
GPS_RAW_INT         X
SYSTEM_TIME
AHRS
WIND                X
EKF_STATUS_REPORT   X
VIBRATION           X
BATTERY_STATUS      X
MCU_STATUS          X

'''
