"""sitl_support.py —— SITL / 真机入口共用的 MAVSDK 小工具。

从根目录 `pod_sitl_test.py` 拆出：原先两个现行入口
（`task_runner_sitl.py` / `task_runner_real.py`）都从这个"测试脚本"里 import
`now_ms` / `wait_until_ready` 等，语义上很别扭。现归位为独立支持模块，
入口只依赖它，`l1_sitl_test.py` 则作为纯粹的 L1 集成测试脚本。

内容（刻意保持无业务逻辑）：
  * `RATE_*`：遥测流速率**唯一来源**（各订阅方共用；PX4 速率由最后一次设置决定）
  * `now_ms`：单调毫秒时钟（与 PodBridge / TaskRunner 同源）
  * `TelemetryHub`：后台订阅 NED 位置 + 欧拉角，供 VirtualPod 几何合成
  * `wait_until_ready` / `set_telemetry_rates` / `print_flight_mode`：连接准备与诊断
"""

import asyncio
import math
import time

from pod_bridge.virtual_pod import CarrierView

# 遥测流速率（Hz）——**唯一来源**，供各订阅方（Px4Gateway / TelemetryHub）共用。
# PX4 的消息速率由最后一次 SET_MESSAGE_INTERVAL 决定，多个订阅方用不同值会互相
# 覆盖；速率越高链路负荷越大（实测压出 MAVSDK 回调队列积压告警）。
# 取值依据见 set_telemetry_rates()。
RATE_ATTITUDE_HZ = 20.0     # 控制律 attitude_stale_ms=300 要用姿态
RATE_POSITION_HZ = 10.0     # 起飞到位 / 到参考点距离
RATE_POSITION_VELOCITY_NED_HZ = 10.0  # 虚拟吊舱几何（TelemetryHub）
RATE_GPS_HZ = 2.0
RATE_VELOCITY_HZ = 2.0
RATE_IN_AIR_HZ = 2.0


def now_ms() -> int:
    """单调毫秒（与 Px4Gateway / PodBridge / TaskRunner 共用同一时钟源）。"""
    return int(time.monotonic() * 1000.0)


class TelemetryHub:
    """后台订阅 NED 位置/姿态，供 VirtualPod 几何合成与锚点记录。"""

    def __init__(self) -> None:
        self.yaw_rad = 0.0
        self.north_m = 0.0
        self.east_m = 0.0
        self.down_m = 0.0
        self.attitude_time_ms = 0
        self._tasks: list = []

    def start(self, drone) -> None:
        self._tasks.append(asyncio.create_task(self._attitude_loop(drone)))
        self._tasks.append(asyncio.create_task(self._position_loop(drone)))

    async def _attitude_loop(self, drone) -> None:
        async for euler in drone.telemetry.attitude_euler():
            self.yaw_rad = math.radians(euler.yaw_deg)
            self.attitude_time_ms = now_ms()

    async def _position_loop(self, drone) -> None:
        async for pos_vel in drone.telemetry.position_velocity_ned():
            pos = pos_vel.position
            self.north_m = pos.north_m
            self.east_m = pos.east_m
            self.down_m = pos.down_m

    async def stop(self) -> None:
        for task in self._tasks:
            task.cancel()
        for task in self._tasks:
            try:
                await task
            except asyncio.CancelledError:
                pass
        self._tasks.clear()

    def view(self) -> CarrierView:
        return CarrierView(north_m=self.north_m, east_m=self.east_m,
                           down_m=self.down_m, yaw_rad=self.yaw_rad)


async def print_flight_mode(drone):
    async for flight_mode in drone.telemetry.flight_mode():
        print(f"[模式] 当前飞行模式: {flight_mode}")
        return flight_mode


async def wait_until_ready(drone):
    """等待连接 + 位置估计就绪（各阶段限时 30 s）。"""
    print("等待无人机连接...")
    try:
        async with asyncio.timeout(30.0):
            async for state in drone.core.connection_state():
                if state.is_connected:
                    print("-- 已连接")
                    break
    except TimeoutError:
        raise RuntimeError("30 s 内未建立 MAVSDK 连接（检查 SITL 是否在运行）") from None

    print("等待全球位置估计就绪...")
    try:
        async with asyncio.timeout(30.0):
            async for health in drone.telemetry.health():
                if health.is_global_position_ok and health.is_home_position_ok:
                    print("-- 位置估计 OK")
                    break
    except TimeoutError:
        raise RuntimeError("30 s 内位置估计未就绪") from None


async def set_telemetry_rates(drone):
    """按**实际消费需求**设定遥测流速率（不是越高越好）。

    取值依据（2026-10-15 复核）：
      * `attitude_euler` 20Hz —— 控制律 `attitude_stale_ms=300` 要用姿态；20Hz 足够
      * `position` 10Hz —— 任务层只需「起飞到位」与「到参考点距离」，10Hz 足够
      * `position_velocity_ned` 10Hz —— 仅虚拟吊舱几何用（真机链路不订阅）
      * `gps_info` / `velocity_ned` / `in_air` 2Hz —— 状态量，变化极慢

    原值（position/attitude 各 50Hz + pvn 30Hz ≈ 130 msg/s）在 CPU 紧张的宿主机上
    会压出 MAVSDK「User callback queue slow」告警；下调后链路负荷降一个量级。
    失败沿用默认值并告警。
    """
    for stream, rate in (("position", RATE_POSITION_HZ),
                         ("attitude_euler", RATE_ATTITUDE_HZ),
                         ("position_velocity_ned", RATE_POSITION_VELOCITY_NED_HZ),
                         ("gps_info", RATE_GPS_HZ),
                         ("velocity_ned", RATE_VELOCITY_HZ),
                         ("in_air", RATE_IN_AIR_HZ)):
        try:
            await getattr(drone.telemetry, f"set_rate_{stream}")(rate)
        except Exception as e:
            print(f"! 设置 {stream} 速率失败（沿用默认值）: {e}")