#!/usr/bin/env python3
"""task_runner_sitl.py —— 任务级状态机 SITL 全流程集成测试（预定流程.md 天空侧）。

用状态机驱动完整流程：WAIT → ARM_TAKEOFF → GOTO(L2 坐标引导) → TRACK(L1 航向伺服)
→ APPROACH(悬停) → RTL(返航落地)。相比 pod_sitl_test.py（固定场景、人工编排），
本脚本把「飞行编排 + 控制律切换」交给 TaskRunner 状态机统一驱动。

场景：
  * 起飞前注入虚拟目标坐标（模拟地面站给坐标）——目标 = home 北 30m / 东 15m；
  * VirtualPod 按无电机语义合成表2-2（框架角恒0、无激光、锁定=目标在视场内）；
  * 后台任务持续刷新注入坐标（模拟地面站持续上报），坐标始终新鲜；
  * GOTO→TRACK 由状态机按「距目标 < switch_radius 且吊舱锁定」守卫切换。

运行前：启动 PX4 SITL（默认本机 14540）。用法：
  python task_runner_sitl.py
  python task_runner_sitl.py --conn udp://:14540 --config utils/config/pod_bridge.json
"""

import argparse
import asyncio
import logging
import math
import sys
import time
from pathlib import Path

_ROOT = Path(__file__).resolve().parent
_UTILS = _ROOT / "utils"
for _p in (_UTILS, _ROOT):
    if str(_p) not in sys.path:
        sys.path.insert(0, str(_p))

from mavsdk import System

from pod_bridge.bridge import PodBridge
from pod_bridge.config import load_from_json
from pod_bridge.ground_link import GroundInjectConfig, GroundLink, GroundTarget
from pod_bridge.guidance import haversine_distance_m
from pod_bridge.px4_gateway import Px4Gateway
from pod_bridge.virtual_pod import (CarrierView, VirtualPod, VirtualTarget,
                                    VirtualTargetConfig)
from task_runner.log import TaskLogConfig, TaskLogger
from task_runner.runner import TaskRunner
from task_runner.state_machine import TaskStateConfig, TaskStateMachine
from pod_sitl_test import (TelemetryHub, now_ms, print_flight_mode,
                          set_telemetry_rates, wait_until_ready)

logger = logging.getLogger("task_sitl")

CONNECTION = "udpin://0.0.0.0:14540"
CONFIG_PATH = _UTILS / "config" / "pod_bridge.json"

# 虚拟目标：相对 home（NED 原点）北 30m / 东 15m / 与悬停同高度（不追垂直）
TARGET_NORTH_M = 30.0
TARGET_EAST_M = 15.0
TARGET_DOWN_M = -5.0          # 与 5m 悬停高度持平 → 垂直通道无动作

TASK_TIMEOUT_S = 180.0         # 全流程总超时
INJECT_PERIOD_S = 0.5          # 注入坐标刷新周期（地面站上报快于 guidance.target_stale_ms）

_DEG_LAT_M = 111320.0


def ned_offset_to_latlon(lat0: float, lon0: float, north_m: float,
                         east_m: float) -> tuple[float, float]:
    """NED 偏移（米）转经纬度（局部平面近似，SITL 小范围够用）。"""
    lat = lat0 + north_m / _DEG_LAT_M
    lon = lon0 + east_m / (_DEG_LAT_M * math.cos(math.radians(lat0)))
    return lat, lon


async def read_home_gps(drone) -> tuple[float, float]:
    """起飞前读 home 经纬度（放虚拟目标 + 注入坐标）。"""
    try:
        async with asyncio.timeout(10.0):
            async for pos in drone.telemetry.position():
                return pos.latitude_deg, pos.longitude_deg
    except TimeoutError:
        raise RuntimeError("10 s 内未取得 home GPS") from None


def build_sm_config() -> TaskStateConfig:
    """SITL 状态机参数（占位值，仅用于仿真验证链路）。"""
    return TaskStateConfig(
        takeoff_altitude_m=5.0,
        switch_radius_m=25.0,      # 目标在 30m，飞进 25m 内切 TRACK
        goto_timeout_s=60.0,
        track_timeout_s=60.0,
        center_hold_s=3.0,
        center_threshold_deg=3.0,  # 脱靶量 < 3° 判居中
        approach_hover_s=5.0,      # 悬停 5s 后 RTL
        loop_frequency_hz=20.0,
    )


async def inject_loop(ground: GroundLink, lat: float, lon: float,
                      stop: asyncio.Event) -> None:
    """后台持续刷新注入坐标（模拟地面站持续上报，保坐标始终新鲜）。"""
    first = True
    while not stop.is_set():
        ground.inject(GroundTarget(latitude_deg=lat, longitude_deg=lon,
                                   abs_alt_m=0.0))
        if first:
            print(f"-- 首次注入坐标 time_ms={ground.latest().time_ms}")
            first = False
        await asyncio.sleep(INJECT_PERIOD_S)


async def run(args) -> int:
    cfg = load_from_json(args.config)
    cfg.ground_inject.enabled = True  # 启用地面注入（骨架期 JSON 默认关）
    cfg.guidance.target_stale_ms = 5000  # 注入坐标超时放宽（表4 默认 1000 对低频注入太紧）
    sm_cfg = build_sm_config()

    drone = System()
    await drone.connect(system_address=args.conn)
    try:
        await wait_until_ready(drone)
    except RuntimeError as e:
        print(f"! {e}")
        return 1
    await set_telemetry_rates(drone)

    home_lat, home_lon = await read_home_gps(drone)
    tgt_lat, tgt_lon = ned_offset_to_latlon(home_lat, home_lon,
                                            TARGET_NORTH_M, TARGET_EAST_M)
    print(f"-- home ({home_lat:.6f}, {home_lon:.6f})，"
          f"注入目标 ({tgt_lat:.6f}, {tgt_lon:.6f})（北 {TARGET_NORTH_M}m "
          f"东 {TARGET_EAST_M}m）")

    # 装配：真 SITL 飞控 + 模拟吊舱 + 注入坐标 + 状态机编排器
    hub = TelemetryHub()
    hub.start(drone)
    target = VirtualTarget(
        VirtualTargetConfig(ahead_m=TARGET_NORTH_M, right_m=TARGET_EAST_M,
                            down_m=TARGET_DOWN_M, arc_radius_m=0.1,
                            arc_speed_mps=0.0, arc_ccw=True),
        (0.0, 0.0, 0.0), 0.0)  # anchor=home NED 原点，yaw=0 → 北30/东15
    pod = VirtualPod(target, hub.view,
                     fov_h_deg=cfg.tracking.camera.fov_h_deg,
                     fov_v_deg=cfg.tracking.camera.fov_v_deg)
    px4 = Px4Gateway(drone)
    ground = GroundLink(cfg.ground_inject, now_ms_fn=now_ms)
    bridge = PodBridge(cfg, px4, pod, now_ms_fn=now_ms,
                       ground_target_fn=lambda: ground.latest())
    sm = TaskStateMachine(sm_cfg)
    task_log = TaskLogger(TaskLogConfig(), now_ms_fn=now_ms)
    runner = TaskRunner(cfg, sm_cfg, px4, bridge, sm, ground, drone,
                        now_ms_fn=now_ms, logger=task_log)

    stop = asyncio.Event()
    if args.ground_station:
        # 地面站模式：UDP 接收地面站注入坐标（地面站页面 POST /api/inject → 14555）
        ground_task = asyncio.ensure_future(ground.run(stop))
    else:
        ground_task = asyncio.ensure_future(inject_loop(ground, tgt_lat, tgt_lon, stop))
    runner_task = asyncio.ensure_future(runner.run(stop))

    print("-- 启动任务状态机全流程（WAIT→…→RTL）")
    if args.ground_station:
        print(f"   [地面站模式] 等待 UDP {cfg.ground_inject.port} 注入坐标…")
    print(f"   状态机: takeoff={sm_cfg.takeoff_altitude_m}m "
          f"switch_radius={sm_cfg.switch_radius_m}m "
          f"center={sm_cfg.center_threshold_deg}°/{sm_cfg.center_hold_s}s "
          f"hover={sm_cfg.approach_hover_s}s")

    # 状态打印 + 总超时保护
    started = time.monotonic()
    last_state = None
    diag_count = 0
    try:
        while not runner.rtl_done:
            await asyncio.sleep(1.0)
            if runner.state is not last_state:
                last_state = runner.state
                print(f"[任务] 进入状态: {runner.state.value}")
            if runner_task.done():
                exc = runner_task.exception()
                if exc is not None:
                    print(f"! runner 异常退出: {type(exc).__name__}: {exc}")
                break
            if runner.state.value == "goto" and diag_count % 3 == 0:
                own = px4.own_state()
                g = ground.latest()
                dist = haversine_distance_m(own.latitude_deg, own.longitude_deg,
                                            g.latitude_deg, g.longitude_deg)
                try:
                    async for m in drone.telemetry.flight_mode():
                        mode = str(m)
                        break
                except Exception:
                    mode = "?"
                print(f"[诊断] pos_valid={own.position_valid} "
                      f"pos_ms={own.position_time_ms} now={now_ms()} "
                      f"mode={mode} yaw={own.yaw_rad:.2f} dist={dist:.1f}m "
                      f"lat={own.latitude_deg:.5f} lon={own.longitude_deg:.5f}")
            diag_count += 1
            if time.monotonic() - started > TASK_TIMEOUT_S:
                print("! 全流程超时，强制返航")
                runner.request_rtl()
        if runner.rtl_done:
            print("-- 任务完成（RTL 落地）")
    except asyncio.CancelledError:
        raise
    finally:
        stop.set()
        ground_task.cancel()
        runner_task.cancel()
        try:
            await runner_task
        except asyncio.CancelledError:
            pass
        try:
            await ground_task
        except asyncio.CancelledError:
            pass
        await hub.stop()

    await print_flight_mode(drone)
    print("\n===== 任务状态机全流程完成 =====")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="任务级状态机 SITL 全流程")
    ap.add_argument("--conn", default=CONNECTION, help="MAVSDK 连接串")
    ap.add_argument("--config", default=str(CONFIG_PATH), help="桥接配置 JSON 路径")
    ap.add_argument("--ground-station", action="store_true",
                    help="地面站模式：UDP 接收地面站注入坐标（而非内置虚拟目标自动注入）")
    args = ap.parse_args()

    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
                        datefmt="%H:%M:%S")
    return asyncio.run(run(args))


if __name__ == "__main__":
    sys.exit(main())