"""XRoboToolkit SDK 到当前 UDP controller JSON 协议的桥接脚本。

该脚本运行在安装了 `xrobotoolkit_sdk` 的 Python 环境中，从官方
XRoboToolkit PC-Service SDK 读取 PICO 左右手柄 pose / grip / trigger，
再发送现有 `udp_controller_receiver` 已兼容的 UDP JSON 包。
"""

import argparse
import json
import math
import os
import socket
import sys
import time
from collections.abc import Callable, Sequence
from typing import Any


IDENTITY_QUAT = [0.0, 0.0, 0.0, 1.0]
ZERO_POS = [0.0, 0.0, 0.0]
POSE_SOURCE = "xrobotoolkit"


def _positive_float(name: str) -> Callable[[str], float]:
    def parser(value: str) -> float:
        number = float(value)
        if number <= 0.0:
            raise argparse.ArgumentTypeError(f"{name} must be > 0")
        return number

    return parser


def _non_negative_int(name: str) -> Callable[[str], int]:
    def parser(value: str) -> int:
        number = int(value)
        if number < 0:
            raise argparse.ArgumentTypeError(f"{name} must be >= 0")
        return number

    return parser


def _clamp_float(value: Any, low: float, high: float) -> float:
    number = float(value)
    return min(max(number, low), high)


class _HysteresisSwitch:
    def __init__(self, *, press_threshold: float, release_threshold: float) -> None:
        self.press_threshold = press_threshold
        self.release_threshold = release_threshold
        self.pressed = False

    def update(self, value: Any) -> bool:
        analog = _clamp_float(value, 0.0, 1.0)
        if self.pressed:
            if analog <= self.release_threshold:
                self.pressed = False
        elif analog >= self.press_threshold:
            self.pressed = True
        return self.pressed


def _pose_to_pos_quat(pose: Any) -> tuple[list[float], list[float]]:
    values = list(pose)
    if len(values) != 7:
        raise ValueError(f"expected XRoboToolkit pose with 7 values, got {len(values)}")

    numbers = [float(value) for value in values]
    if not all(math.isfinite(value) for value in numbers):
        raise ValueError(f"pose contains non-finite values: {numbers!r}")
    return numbers[:3], numbers[3:7]


def _safe_axis(value: Any) -> list[float]:
    try:
        axis = [float(item) for item in list(value)]
    except (TypeError, ValueError):
        return [0.0, 0.0]
    if len(axis) < 2:
        return [0.0, 0.0]
    return [_clamp_float(axis[0], -1.0, 1.0), _clamp_float(axis[1], -1.0, 1.0)]


def _safe_bool(func: Callable[[], Any], default: bool = False) -> bool:
    try:
        return bool(func())
    except Exception:
        return default


def _controller_payload(
    *,
    hand: str,
    pose: Any,
    grip_value: Any,
    trigger_value: Any,
    axis: Any,
    buttons: dict[str, bool],
    grip_pressed: bool,
    trigger_pressed: bool,
    pose_valid: bool = True,
) -> dict[str, Any]:
    pos, quat = _pose_to_pos_quat(pose) if pose_valid else (ZERO_POS.copy(), IDENTITY_QUAT.copy())
    grip_float = _clamp_float(grip_value, 0.0, 1.0)
    trigger_float = _clamp_float(trigger_value, 0.0, 1.0)
    return {
        "hand": hand,
        "grip": pose_valid and grip_pressed,
        "trigger": 1.0 if pose_valid and trigger_pressed else 0.0,
        "pos": pos,
        "quat": quat,
        "pose_valid": pose_valid,
        "pose_source": POSE_SOURCE,
        "grip_value": grip_float,
        "trigger_value": trigger_float,
        "axis": _safe_axis(axis),
        "buttons": buttons,
    }


def _stop_controller_payload(hand: str) -> dict[str, Any]:
    return {
        "hand": hand,
        "grip": False,
        "trigger": 0.0,
        "pos": ZERO_POS.copy(),
        "quat": IDENTITY_QUAT.copy(),
        "pose_valid": False,
        "pose_source": POSE_SOURCE,
        "grip_value": 0.0,
        "trigger_value": 0.0,
        "axis": [0.0, 0.0],
        "buttons": {
            "grip": False,
            "primary": False,
            "secondary": False,
            "menu": False,
            "axis_click": False,
        },
    }


def _packet(seq: int, frame_id: str, controllers: dict[str, dict[str, Any]]) -> dict[str, Any]:
    now = time.monotonic()
    return {
        "t": now,
        "source_time": now,
        "seq": seq,
        "frame_id": frame_id,
        "controllers": controllers,
    }


def _send_packet(sock: socket.socket, host: str, port: int, payload: dict[str, Any]) -> None:
    sock.sendto(json.dumps(payload, separators=(",", ":")).encode("utf-8"), (host, port))


def _send_stop_packets(
    sock: socket.socket,
    *,
    host: str,
    port: int,
    frame_id: str,
    seq: int,
    count: int,
    interval_sec: float,
) -> int:
    controllers = {
        "left": _stop_controller_payload("left"),
        "right": _stop_controller_payload("right"),
    }
    for _ in range(count):
        seq += 1
        _send_packet(sock, host, port, _packet(seq, frame_id, controllers))
        if interval_sec > 0.0:
            time.sleep(interval_sec)
    return seq


def _buttons_payload(
    *,
    grip: bool,
    primary: Callable[[], Any],
    secondary: Callable[[], Any],
    menu: Callable[[], Any],
    axis_click: Callable[[], Any],
) -> dict[str, bool]:
    return {
        "grip": grip,
        "primary": _safe_bool(primary),
        "secondary": _safe_bool(secondary),
        "menu": _safe_bool(menu),
        "axis_click": _safe_bool(axis_click),
    }


def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Read PICO controller state from xrobotoolkit_sdk and send UDP JSON packets."
    )
    parser.add_argument("--host", default="127.0.0.1", help="UDP receiver host.")
    parser.add_argument("--port", type=int, default=15000, help="UDP receiver port.")
    parser.add_argument("--hz", type=_positive_float("hz"), default=90.0, help="Send rate in Hz.")
    parser.add_argument(
        "--grip-threshold",
        type=float,
        default=0.9,
        help="Grip analog press threshold used to set the stable boolean grip field.",
    )
    parser.add_argument(
        "--grip-release-threshold",
        type=float,
        default=0.75,
        help="Grip analog release threshold. Must be <= --grip-threshold.",
    )
    parser.add_argument(
        "--trigger-threshold",
        type=float,
        default=0.95,
        help="Trigger analog press threshold used to set the stable trigger field to 1.0.",
    )
    parser.add_argument(
        "--trigger-release-threshold",
        type=float,
        default=0.75,
        help="Trigger analog release threshold. Must be <= --trigger-threshold.",
    )
    parser.add_argument("--frame-id", default="xr_world", help="frame_id written into UDP packets.")
    parser.add_argument(
        "--stop-packets",
        type=_non_negative_int("stop-packets"),
        default=5,
        help="Number of all-stop packets sent on shutdown or SDK read failure.",
    )
    return parser.parse_args(argv)


def _import_xrobotoolkit_sdk():
    try:
        import xrobotoolkit_sdk as xrt  # type: ignore[import-not-found]
    except ImportError as exc:
        message = (
            "Failed to import xrobotoolkit_sdk. Install the XRoboToolkit SDK into "
            "the ROS Python environment and set LD_LIBRARY_PATH, for example:\n"
            "  export LD_LIBRARY_PATH=/opt/apps/roboticsservice/SDK/x64:${LD_LIBRARY_PATH:-}\n"
            "  ros2 run xr_rm_input xrobotoolkit_to_udp_bridge"
        )
        raise RuntimeError(message) from exc
    return xrt


def _validate_threshold_pair(parser_name: str, press: float, release: float) -> None:
    if release > press:
        raise ValueError(
            f"{parser_name} release threshold ({release:.2f}) must be <= press threshold ({press:.2f})"
        )


def main(argv: Sequence[str] | None = None) -> None:
    args = _parse_args(argv)
    args.grip_threshold = _clamp_float(args.grip_threshold, 0.0, 1.0)
    args.grip_release_threshold = _clamp_float(args.grip_release_threshold, 0.0, 1.0)
    args.trigger_threshold = _clamp_float(args.trigger_threshold, 0.0, 1.0)
    args.trigger_release_threshold = _clamp_float(args.trigger_release_threshold, 0.0, 1.0)
    _validate_threshold_pair("grip", args.grip_threshold, args.grip_release_threshold)
    _validate_threshold_pair("trigger", args.trigger_threshold, args.trigger_release_threshold)

    grip_switches = {
        "left": _HysteresisSwitch(
            press_threshold=args.grip_threshold,
            release_threshold=args.grip_release_threshold,
        ),
        "right": _HysteresisSwitch(
            press_threshold=args.grip_threshold,
            release_threshold=args.grip_release_threshold,
        ),
    }
    trigger_switches = {
        "left": _HysteresisSwitch(
            press_threshold=args.trigger_threshold,
            release_threshold=args.trigger_release_threshold,
        ),
        "right": _HysteresisSwitch(
            press_threshold=args.trigger_threshold,
            release_threshold=args.trigger_release_threshold,
        ),
    }

    stop_interval_sec = min(1.0 / args.hz, 0.02)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    seq = 0
    xrt = _import_xrobotoolkit_sdk()
    sdk_initialized = False

    try:
        xrt.init()
        sdk_initialized = True
        print(
            "XRoboToolkit UDP bridge started: "
            f"pid={os.getpid()}, endpoint={args.host}:{args.port}, hz={args.hz:.1f}, "
            f"grip={args.grip_threshold:.2f}/{args.grip_release_threshold:.2f}, "
            f"trigger={args.trigger_threshold:.2f}/{args.trigger_release_threshold:.2f}",
            flush=True,
        )

        dt = 1.0 / args.hz
        next_tick = time.monotonic()
        while True:
            try:
                left_grip_value = xrt.get_left_grip()
                right_grip_value = xrt.get_right_grip()
                left_trigger_value = xrt.get_left_trigger()
                right_trigger_value = xrt.get_right_trigger()
                left_grip = grip_switches["left"].update(left_grip_value)
                right_grip = grip_switches["right"].update(right_grip_value)
                left_trigger = trigger_switches["left"].update(left_trigger_value)
                right_trigger = trigger_switches["right"].update(right_trigger_value)
                controllers = {
                    "left": _controller_payload(
                        hand="left",
                        pose=xrt.get_left_controller_pose(),
                        grip_value=left_grip_value,
                        trigger_value=left_trigger_value,
                        axis=xrt.get_left_axis(),
                        buttons=_buttons_payload(
                            grip=left_grip,
                            primary=xrt.get_X_button,
                            secondary=xrt.get_Y_button,
                            menu=xrt.get_left_menu_button,
                            axis_click=xrt.get_left_axis_click,
                        ),
                        grip_pressed=left_grip,
                        trigger_pressed=left_trigger,
                    ),
                    "right": _controller_payload(
                        hand="right",
                        pose=xrt.get_right_controller_pose(),
                        grip_value=right_grip_value,
                        trigger_value=right_trigger_value,
                        axis=xrt.get_right_axis(),
                        buttons=_buttons_payload(
                            grip=right_grip,
                            primary=xrt.get_A_button,
                            secondary=xrt.get_B_button,
                            menu=xrt.get_right_menu_button,
                            axis_click=xrt.get_right_axis_click,
                        ),
                        grip_pressed=right_grip,
                        trigger_pressed=right_trigger,
                    ),
                }
            except Exception:
                print(
                    "XRoboToolkit SDK read failed; sending stop packets.",
                    file=sys.stderr,
                    flush=True,
                )
                seq = _send_stop_packets(
                    sock,
                    host=args.host,
                    port=args.port,
                    frame_id=args.frame_id,
                    seq=seq,
                    count=args.stop_packets,
                    interval_sec=stop_interval_sec,
                )
                raise

            seq += 1
            _send_packet(sock, args.host, args.port, _packet(seq, args.frame_id, controllers))
            next_tick += dt
            sleep_sec = next_tick - time.monotonic()
            if sleep_sec > 0.0:
                time.sleep(sleep_sec)
            else:
                next_tick = time.monotonic()

    except KeyboardInterrupt:
        print("XRoboToolkit UDP bridge interrupted; sending stop packets.", flush=True)
    finally:
        seq = _send_stop_packets(
            sock,
            host=args.host,
            port=args.port,
            frame_id=args.frame_id,
            seq=seq,
            count=args.stop_packets,
            interval_sec=stop_interval_sec,
        )
        if sdk_initialized:
            try:
                xrt.close()
            except Exception as exc:
                print(f"XRoboToolkit SDK close failed: {exc}", file=sys.stderr, flush=True)
        sock.close()
        print(f"XRoboToolkit UDP bridge stopped at seq={seq}.", flush=True)


if __name__ == "__main__":
    main()
