"""MocapFrame UDP 接收示例（使用 forsense_mocap SDK）。

一个数字选择菜单，提供：全身数据 / 手部数据 / 诊断模式 / 退出。
每个数据界面内按 Ctrl+C 即可退出并回到数字选择界面。

用法:
    python recv_mocap_udp.py                  # 默认端口 15000
    python recv_mocap_udp.py --port 15001     # 指定端口
"""

import argparse
import sys
import time

from forsense_mocap import JOINT_NAMES, UdpMocapReceiver

# 手部关节索引：左手 15-34，右手 44-63
HAND_INDEXES = list(range(15, 35)) + list(range(44, 64))
# 全身：0-63 全部
BODY_INDEXES = list(range(64))


def print_prompt() -> None:
    print("\n===== Forsense Mocap 接收器 =====")
    print("  1) 全身数据")
    print("  2) 手部数据")
    print("  3) 诊断模式")
    print("  4) 退出")
    print("================================")


def run_recv(port: int, indexes: list, title: str) -> None:
    """进入数据接收界面，Ctrl+C 退出回到主菜单。

    使用回调模式（后台线程满速收包，不受控制台打印拖慢），
    并每秒打印一次真实帧率与最新一帧的关节数据。
    """
    state = {"frames": 0, "latest": None}
    last_ts = time.time()
    fps = 0

    def on_frame(frame):
        state["frames"] += 1
        state["latest"] = frame

    # 传入 on_frame -> UdpMocapReceiver 启动后台接收线程
    recv = UdpMocapReceiver(port=port, on_frame=on_frame)
    recv.start()
    print(f"\n[{title}] 接收 0.0.0.0:{port} ... 按 Ctrl+C 返回主菜单")
    try:
        while True:
            time.sleep(1.0)
            now = time.time()
            fps = state["frames"] / max(now - last_ts, 1e-6)
            state["frames"] = 0
            last_ts = now
            frame = state["latest"]
            print("\n[FPS=%d] 来源:%s" % (round(fps), frame.source if frame else "-"))
            if frame is not None:
                print(f"  time_us={frame.timestamp_us}")
                for i in indexes:
                    j = frame.joints[i]
                    print(f"  {JOINT_NAMES[i]:<18} pos={j.position} quat={j.quaternion}")
    except KeyboardInterrupt:
        print("\n已返回主菜单")
    finally:
        recv.stop()


def run_diag(port: int) -> None:
    """进入诊断模式，实时查看连接状态、真实帧率与统计，Ctrl+C 返回主菜单。"""
    state = {"frames": 0, "rejected_prev": 0}
    last_ts = time.time()
    last_connected = False

    def on_frame(frame):
        state["frames"] += 1

    recv = UdpMocapReceiver(port=port, on_frame=on_frame)
    recv.start()
    print(f"\n[诊断模式] 0.0.0.0:{port} ... 按 Ctrl+C 返回主菜单")
    try:
        while True:
            time.sleep(1.0)
            now = time.time()
            fps = state["frames"] / max(now - last_ts, 1e-6)
            state["frames"] = 0
            last_ts = now
            stats = recv.stats
            causes = stats.rejected_causes or {}
            caustr = ", ".join(f"{k}={v}" for k, v in causes.items()) or "-"
            connected = stats.frames > 0 or fps > 0
            if connected and not last_connected:
                print("   连接状态变化: 数据已恢复/正常")
            elif not connected and last_connected:
                print("   连接状态变化: 数据中断（可能设备已断开）")
            last_connected = connected
            conn = "数据正常" if connected else "数据中断（可能设备已断开）"
            print("  连接状态: %-8s | 实时FPS=%d | 累计帧=%d 拒绝=%d | 拒绝原因:%s"
                  % (conn, round(fps), stats.frames, stats.rejected, caustr))
    except KeyboardInterrupt:
        print("\n已返回主菜单")
    finally:
        recv.stop()


def main() -> int:
    parser = argparse.ArgumentParser(description="MocapFrame UDP 接收菜单演示")
    parser.add_argument("--port", type=int, default=15000, help="UDP port (default 15000)")
    args = parser.parse_args()

    while True:
        print_prompt()
        try:
            choice = input("请选择 [1-4]: ").strip()
        except KeyboardInterrupt:
            print("\n再见")
            return 0
        if choice == "1":
            run_recv(args.port, BODY_INDEXES, "全身数据")
        elif choice == "2":
            run_recv(args.port, HAND_INDEXES, "手部数据")
        elif choice == "3":
            run_diag(args.port)
        elif choice == "4":
            print("再见")
            return 0
        else:
            print("输入无效，请重新选择。")


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