code-test-drone/utils/pod_bridge/px4_gateway.py

"""px4_gateway.py —— PX4(MAVSDK) 侧网关：读飞控遥测 + 发 offboard 速度指令。

职责（桥接的飞控侧）：
  * 订阅 PX4 遥测（position / attitude_euler / gps_info / velocity_ned），缓存为
    OwnState 快照，供桥接层构造 表3(载机姿态下发) 与 GuidanceInput(引导输入)；
  * 复用 visual_track_control.MavsdkCommandGateway 发送 offboard 速度设定点
    （不重造三分支映射，直接喂 GuidanceLaw 的 ControlOutput）。

设计：遥测拆为「纯缓存更新方法 update_*（无 IO，可单测）」+「薄异步订阅胶水
start_telemetry/_consume」，便于无硬件 mock 测试。指令发送整体委托给被组合的
MavsdkCommandGateway。
"""

import asyncio
import logging
import math
import time
from dataclasses import dataclass, replace
from typing import Any, Callable, Optional

from visual_track_control.mavsdk_gateway import MavsdkCommandGateway

logger = logging.getLogger(__name__)

_GPS_FIX_3D = 3  # mavsdk GpsFixType.THREE_D_FIX 起视为定位有效


def _enum_int(value: Any) -> int:
    """MAVSDK 枚举（Enum/IntEnum）或裸整数统一转 int。"""
    if hasattr(value, "value"):
        try:
            return int(value.value)
        except (TypeError, ValueError):
            return 0
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


@dataclass
class OwnState:
    """PX4 遥测缓存快照（角度弧度 NED；位置度；高度/速度米制）。"""

    latitude_deg: float = 0.0
    longitude_deg: float = 0.0
    absolute_altitude_m: float = 0.0
    relative_altitude_m: float = 0.0
    roll_rad: float = 0.0
    pitch_rad: float = 0.0
    yaw_rad: float = 0.0
    ground_speed_mps: float = 0.0
    fix_type: int = 0
    num_satellites: int = 0
    has_position: bool = False
    has_attitude: bool = False
    position_valid: bool = False  # has_position 且 GPS 3D 定位
    time_ms: int = 0              # 任意遥测流最近一次更新（毫秒），仅供诊断
    position_time_ms: int = 0     # position 流最近一次更新（毫秒），位置超龄判定专用


class Px4Gateway:
    """drone：mavsdk System（duck-typed，便于注入 mock）。"""

    def __init__(self, drone: Any,
                 now_ms_fn: Optional[Callable[[], int]] = None) -> None:
        self._drone = drone
        self._cmd = MavsdkCommandGateway(drone)
        self._now_ms = now_ms_fn or (lambda: int(time.monotonic() * 1000))
        self._state = OwnState()
        self._tasks: list = []

    # --- 连接（实机）-------------------------------------------------------
    @classmethod
    async def connect(cls, connection_string: str,
                      now_ms_fn: Optional[Callable[[], int]] = None) -> "Px4Gateway":
        """创建 MAVSDK System、连接并等待联通后返回网关。"""
        from mavsdk import System
        drone = System()
        logger.info("[Px4Gateway] 连接飞控 %s ...", connection_string)
        await drone.connect(system_address=connection_string)
        async for state in drone.core.connection_state():
            if state.is_connected:
                logger.info("[Px4Gateway] 飞控已连接")
                break
        return cls(drone, now_ms_fn=now_ms_fn)

    # --- 遥测缓存更新（纯方法，可单测）------------------------------------
    def update_position(self, p: Any) -> None:
        self._state.latitude_deg = float(p.latitude_deg)
        self._state.longitude_deg = float(p.longitude_deg)
        self._state.absolute_altitude_m = float(p.absolute_altitude_m)
        self._state.relative_altitude_m = float(p.relative_altitude_m)
        self._state.has_position = True
        self._recompute_valid()
        now = self._now_ms()
        self._state.time_ms = now
        self._state.position_time_ms = now  # 仅此流刷新，姿态高频流不得代刷

    def update_attitude(self, e: Any) -> None:
        """MAVSDK attitude_euler() 返回 EulerAngle(roll_deg/pitch_deg/yaw_deg，度)，
        转弧度存入 OwnState（下游控制律统一用弧度 NED）。"""
        self._state.roll_rad = math.radians(float(e.roll_deg))
        self._state.pitch_rad = math.radians(float(e.pitch_deg))
        self._state.yaw_rad = math.radians(float(e.yaw_deg))
        self._state.has_attitude = True
        self._state.time_ms = self._now_ms()

    def update_gps(self, g: Any) -> None:
        self._state.fix_type = _enum_int(getattr(g, "fix_type", 0))
        self._state.num_satellites = _enum_int(getattr(g, "num_satellites", 0))
        self._recompute_valid()
        self._state.time_ms = self._now_ms()

    def update_velocity_ned(self, v: Any) -> None:
        """MAVSDK velocity_ned() 返回 VelocityNed(north/east/down_m_s)，
        水平合速度即地速（表3 ground_speed 语义）。"""
        self._state.ground_speed_mps = math.hypot(
            float(v.north_m_s), float(v.east_m_s))
        self._state.time_ms = self._now_ms()

    def own_state(self) -> OwnState:
        """返回当前缓存的快照副本（调用方修改不影响内部状态）。"""
        return replace(self._state)

    def _recompute_valid(self) -> None:
        self._state.position_valid = (self._state.has_position
                                      and self._state.fix_type >= _GPS_FIX_3D)

    # --- 遥测订阅（异步胶水）----------------------------------------------
    async def start_telemetry(self) -> None:
        """订阅四路遥测流，后台任务持续刷新缓存。幂等：已订阅则直接返回。"""
        if self._tasks:
            return  # 已订阅，避免重复起任务（runner 会先订阅，bridge.start 再触发时跳过）
        t = self._drone.telemetry
        self._tasks = [
            asyncio.ensure_future(self._consume(t.position(), self.update_position)),
            asyncio.ensure_future(
                self._consume(t.attitude_euler(), self.update_attitude)),
            asyncio.ensure_future(self._consume(t.gps_info(), self.update_gps)),
            asyncio.ensure_future(
                self._consume(t.velocity_ned(), self.update_velocity_ned)),
        ]
        logger.info("[Px4Gateway] 遥测订阅已启动")

    async def _consume(self, stream: Any, handler: Callable[[Any], None]) -> None:
        try:
            async for msg in stream:
                handler(msg)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("[Px4Gateway] 遥测流中断: %s", exc)

    async def stop_telemetry(self) -> None:
        for task in self._tasks:
            task.cancel()
        if self._tasks:
            await asyncio.gather(*self._tasks, return_exceptions=True)
        self._tasks = []
        logger.info("[Px4Gateway] 遥测订阅已停止")

    # --- offboard 指令（委托 MavsdkCommandGateway）-------------------------
    async def start_offboard(self) -> bool:
        return await self._cmd.start()

    async def send_velocity_heading(self, output: Any,
                                    time_boot_ms: Optional[int] = None) -> bool:
        return await self._cmd.send_velocity_heading(output, time_boot_ms)

    async def send_brake_hover(self, time_boot_ms: Optional[int] = None) -> bool:
        return await self._cmd.send_brake_hover(time_boot_ms)

    async def stop_offboard(self) -> None:
        await self._cmd.stop()