"""XR 手柄 UDP 接收节点。

从 UDP JSON 数据包中解析左右手柄位姿、握持键和扳机值，并发布为
`xr_rm_interfaces/XrController` 消息，供遥操作和夹爪节点订阅。
"""

import json
import socket
from collections.abc import Iterable, Mapping
from typing import Any

import rclpy
from rclpy.node import Node

from xr_rm_interfaces.msg import XrController


class UdpControllerReceiver(Node):
    """将轻量级 XR UDP 数据包转换成 ROS2 左/右手柄消息。"""

    _VALID_POSE_SOURCES = {"pxr_predict", "unity_xr", "xrobotoolkit", "none"}

    def __init__(self) -> None:
        super().__init__("udp_controller_receiver")

        self.declare_parameter("udp_host", "0.0.0.0")
        self.declare_parameter("udp_port", 15000)
        self.declare_parameter("topic", "/xr/right_controller")
        self.declare_parameter("left_topic", "/xr/left_controller")
        self.declare_parameter("right_topic", "")
        self.declare_parameter("default_hand", "right")
        self.declare_parameter("timer_hz", 100.0)
        self.declare_parameter("quat_order", "xyzw")

        udp_host = self.get_parameter("udp_host").value
        udp_port = int(self.get_parameter("udp_port").value)
        legacy_topic = str(self.get_parameter("topic").value)
        left_topic = str(self.get_parameter("left_topic").value)
        right_topic = str(self.get_parameter("right_topic").value) or legacy_topic
        timer_hz = float(self.get_parameter("timer_hz").value)

        self._default_hand = self._normalize_hand(
            self.get_parameter("default_hand").value,
            fallback="right",
        )
        self._quat_order = str(self.get_parameter("quat_order").value).lower()

        # UDP 接收使用非阻塞模式，避免没有手柄数据时卡住 ROS2 定时器。
        self._controller_publishers = {
            "left": self.create_publisher(XrController, left_topic, 10),
            "right": self.create_publisher(XrController, right_topic, 10),
        }
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._socket.bind((udp_host, udp_port))
        self._socket.setblocking(False)

        self.create_timer(1.0 / timer_hz, self._poll_socket)
        self.get_logger().info(
            "正在监听 XR UDP 数据："
            f"{udp_host}:{udp_port} -> left={left_topic}, right={right_topic}"
        )

    def _poll_socket(self) -> None:
        while True:
            try:
                data, _ = self._socket.recvfrom(4096)
            except BlockingIOError:
                return

            try:
                payload = json.loads(data.decode("utf-8"))
                controller_payloads = list(self._iter_controller_payloads(payload))
            except (UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
                self.get_logger().warn(f"XR 数据包格式错误：{exc}", throttle_duration_sec=1.0)
                continue

            for hand, controller_payload in controller_payloads:
                try:
                    msg = self._payload_to_msg(controller_payload, hand)
                except (KeyError, TypeError, ValueError) as exc:
                    self.get_logger().warn(f"XR 手柄字段错误：{exc}", throttle_duration_sec=1.0)
                    continue
                self._controller_publishers[msg.hand].publish(msg)

    def _iter_controller_payloads(self, payload: Any) -> Iterable[tuple[str, Mapping[str, Any]]]:
        if isinstance(payload, list):
            for item in payload:
                yield from self._iter_controller_payloads(item)
            return

        if not isinstance(payload, Mapping):
            raise TypeError("expected JSON object or array")

        if self._contains_pose(payload):
            hand = self._normalize_hand(payload.get("hand", self._default_hand))
            yield hand, payload
            return

        controllers = payload.get("controllers", payload.get("controller"))
        if isinstance(controllers, Mapping):
            # 兼容两类格式：controllers 本身就是单个手柄，或按 left/right 分组。
            if self._contains_pose(controllers):
                hand = self._normalize_hand(controllers.get("hand", self._default_hand))
                yield hand, controllers
                return
            for hand_key, controller_payload in controllers.items():
                if isinstance(controller_payload, Mapping):
                    hand = self._normalize_hand(controller_payload.get("hand", hand_key))
                    yield hand, controller_payload
            return

        if isinstance(controllers, list):
            for controller_payload in controllers:
                yield from self._iter_controller_payloads(controller_payload)
            return

        for hand_key in ("left", "right", "left_controller", "right_controller"):
            controller_payload = payload.get(hand_key)
            if isinstance(controller_payload, Mapping):
                hand = self._normalize_hand(controller_payload.get("hand", hand_key))
                yield hand, controller_payload

    def _payload_to_msg(self, payload: Mapping[str, Any], hand: str) -> XrController:
        pos, quat = self._extract_pose(payload)
        if len(pos) != 3 or len(quat) != 4:
            raise ValueError("expected pos[3] and quat[4]")

        msg = XrController()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = str(payload.get("frame_id", "xr_world"))
        msg.hand = hand
        grip = self._as_bool(payload.get("grip", payload.get("grip_button", False)))
        pose_valid = self._optional_bool(payload.get("pose_valid"), default=True)
        pose_source = self._validate_pose_diagnostics(payload, hand)
        if not pose_valid:
            self.get_logger().warn(
                f"{hand} XR pose_valid=false，强制 grip=false"
                f" source={pose_source or 'missing'}"
                f" tracking_state={payload.get('tracking_state', 'missing')}"
                f" controller_status={payload.get('controller_status', 'missing')}",
                throttle_duration_sec=1.0,
            )
            grip = False

        msg.grip = grip
        msg.trigger = self._clamp_float(payload.get("trigger", 0.0), 0.0, 1.0)
        msg.pose.position.x = float(pos[0])
        msg.pose.position.y = float(pos[1])
        msg.pose.position.z = float(pos[2])

        # Unity/OpenXR 项目里的四元数顺序可能不同，这里用参数统一转换成 ROS 的 xyzw。
        if self._quat_order == "wxyz":
            w, x, y, z = quat
        else:
            x, y, z, w = quat
        msg.pose.orientation.x = float(x)
        msg.pose.orientation.y = float(y)
        msg.pose.orientation.z = float(z)
        msg.pose.orientation.w = float(w)
        return msg

    def _validate_pose_diagnostics(self, payload: Mapping[str, Any], hand: str) -> str:
        has_any_diagnostic = any(
            key in payload for key in ("pose_valid", "pose_source", "tracking_state", "controller_status")
        )
        if has_any_diagnostic and "pose_valid" not in payload:
            self.get_logger().warn(
                f"{hand} XR 诊断字段缺少 pose_valid，按旧协议处理 grip",
                throttle_duration_sec=5.0,
            )
        if has_any_diagnostic and "pose_source" not in payload:
            self.get_logger().warn(
                f"{hand} XR 诊断字段缺少 pose_source",
                throttle_duration_sec=5.0,
            )

        pose_source_value = payload.get("pose_source", "")
        pose_source = str(pose_source_value).strip()
        if pose_source and pose_source not in self._VALID_POSE_SOURCES:
            self.get_logger().warn(
                f"{hand} XR 未知 pose_source={pose_source!r}",
                throttle_duration_sec=5.0,
            )
        return pose_source

    @staticmethod
    def _contains_pose(payload: Mapping[str, Any]) -> bool:
        return any(key in payload for key in ("pos", "position", "p", "pose"))

    def _extract_pose(self, payload: Mapping[str, Any]) -> tuple[list[float], list[float]]:
        pose = payload.get("pose")
        # 支持多种常见字段命名，方便 Unity/OpenXR/Python 调试脚本接入同一节点。
        if isinstance(pose, Mapping):
            pos = pose.get("position", pose.get("pos", pose.get("p")))
            quat = pose.get("orientation", pose.get("quat", pose.get("q", [0.0, 0.0, 0.0, 1.0])))
        elif isinstance(pose, list) and len(pose) >= 7:
            pos = pose[:3]
            quat = pose[3:7]
        else:
            pos = payload.get("pos", payload.get("position", payload.get("p")))
            quat = payload.get("quat", payload.get("orientation", payload.get("q", [0.0, 0.0, 0.0, 1.0])))

        if pos is None:
            raise KeyError("missing pos/position/p")
        return self._vector3(pos), self._quaternion(quat)

    @staticmethod
    def _vector3(value: Any) -> list[float]:
        if isinstance(value, Mapping):
            return [float(value[key]) for key in ("x", "y", "z")]
        vector = list(value)
        if len(vector) != 3:
            raise ValueError("expected 3D position")
        return [float(item) for item in vector]

    def _quaternion(self, value: Any) -> list[float]:
        if isinstance(value, Mapping):
            if self._quat_order == "wxyz":
                return [float(value[key]) for key in ("w", "x", "y", "z")]
            return [float(value[key]) for key in ("x", "y", "z", "w")]
        quat = list(value)
        if len(quat) != 4:
            raise ValueError("expected quaternion")
        return [float(item) for item in quat]

    def _normalize_hand(self, hand: Any, fallback: str | None = None) -> str:
        value = str(hand).lower()
        if value in ("left", "l", "left_controller"):
            return "left"
        if value in ("right", "r", "right_controller"):
            return "right"
        fallback_hand = fallback or self._default_hand
        self.get_logger().warn(
            f"未知手柄标识 {hand!r}，使用默认 {fallback_hand}",
            throttle_duration_sec=1.0,
        )
        return fallback_hand

    @staticmethod
    def _as_bool(value: Any) -> bool:
        if isinstance(value, str):
            return value.lower() in ("1", "true", "yes", "on")
        return bool(value)

    @classmethod
    def _optional_bool(cls, value: Any, default: bool) -> bool:
        if value is None:
            return default
        return cls._as_bool(value)

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


def main(args: list[str] | None = None) -> None:
    rclpy.init(args=args)
    node = UdpControllerReceiver()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()
