# Copyright (c) 2022-2026, The Mindbot Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#
# Adapted from XRoboToolkit (https://github.com/NVIDIA-AI-IOT/XRoboToolkit-Teleop-Sample-Python)
# Original XR client by Zhigen Zhao.

"""XR device client for PICO 4 Ultra via XRoboToolkit SDK.

This module wraps the ``xrobotoolkit_sdk`` C++ extension to provide a
Pythonic interface for retrieving pose, button, joystick, and hand-tracking
data from PICO XR headsets over the local network.

Usage::

    from xr_utils import XrClient

    client = XrClient()
    headset_pose = client.get_pose("headset")
    left_trigger = client.get_key_value("left_trigger")
    client.close()
"""

import numpy as np
import xrobotoolkit_sdk as xrt


class XrClient:
    """Client for PICO XR devices via the XRoboToolkit SDK."""

    def __init__(self):
        """Initialize the SDK and start listening for XR data."""
        xrt.init()
        print("[XrClient] XRoboToolkit SDK initialized.")

    # ------------------------------------------------------------------
    # Pose
    # ------------------------------------------------------------------

    def get_pose(self, name: str) -> np.ndarray:
        """Return the 7-DoF pose [x, y, z, qx, qy, qz, qw] of a device.

        Args:
            name: One of ``"left_controller"``, ``"right_controller"``, ``"headset"``.

        Returns:
            np.ndarray of shape (7,).
        """
        _POSE_FUNCS = {
            "left_controller": xrt.get_left_controller_pose,
            "right_controller": xrt.get_right_controller_pose,
            "headset": xrt.get_headset_pose,
        }
        if name not in _POSE_FUNCS:
            raise ValueError(
                f"Invalid pose name: {name}. Choose from {list(_POSE_FUNCS)}"
            )
        return np.asarray(_POSE_FUNCS[name](), dtype=np.float64)

    # ------------------------------------------------------------------
    # Analog Keys (trigger / grip)
    # ------------------------------------------------------------------

    def get_key_value(self, name: str) -> float:
        """Return an analog value (0.0~1.0) for a trigger or grip.

        Args:
            name: One of ``"left_trigger"``, ``"right_trigger"``,
                  ``"left_grip"``, ``"right_grip"``.
        """
        _KEY_FUNCS = {
            "left_trigger": xrt.get_left_trigger,
            "right_trigger": xrt.get_right_trigger,
            "left_grip": xrt.get_left_grip,
            "right_grip": xrt.get_right_grip,
        }
        if name not in _KEY_FUNCS:
            raise ValueError(
                f"Invalid key name: {name}. Choose from {list(_KEY_FUNCS)}"
            )
        return float(_KEY_FUNCS[name]())

    # ------------------------------------------------------------------
    # Buttons
    # ------------------------------------------------------------------

    def get_button(self, name: str) -> bool:
        """Return boolean state of a face / menu button.

        Args:
            name: One of ``"A"``, ``"B"``, ``"X"``, ``"Y"``,
                  ``"left_menu_button"``, ``"right_menu_button"``,
                  ``"left_axis_click"``, ``"right_axis_click"``.
        """
        _BTN_FUNCS = {
            "A": xrt.get_A_button,
            "B": xrt.get_B_button,
            "X": xrt.get_X_button,
            "Y": xrt.get_Y_button,
            "left_menu_button": xrt.get_left_menu_button,
            "right_menu_button": xrt.get_right_menu_button,
            "left_axis_click": xrt.get_left_axis_click,
            "right_axis_click": xrt.get_right_axis_click,
        }
        if name not in _BTN_FUNCS:
            raise ValueError(
                f"Invalid button name: {name}. Choose from {list(_BTN_FUNCS)}"
            )
        return bool(_BTN_FUNCS[name]())

    # ------------------------------------------------------------------
    # Joystick
    # ------------------------------------------------------------------

    def get_joystick(self, side: str) -> np.ndarray:
        """Return joystick (x, y) for left or right controller.

        Args:
            side: ``"left"`` or ``"right"``.

        Returns:
            np.ndarray of shape (2,).
        """
        if side == "left":
            return np.asarray(xrt.get_left_axis(), dtype=np.float64)
        elif side == "right":
            return np.asarray(xrt.get_right_axis(), dtype=np.float64)
        raise ValueError(f"Invalid side: {side}. Choose 'left' or 'right'.")

    # ------------------------------------------------------------------
    # Hand Tracking
    # ------------------------------------------------------------------

    def get_hand_tracking(self, side: str) -> np.ndarray | None:
        """Return hand tracking joint poses or None if inactive.

        Args:
            side: ``"left"`` or ``"right"``.

        Returns:
            np.ndarray of shape (27, 7) or None.
        """
        if side == "left":
            if not xrt.get_left_hand_is_active():
                return None
            return np.asarray(xrt.get_left_hand_tracking_state(), dtype=np.float64)
        elif side == "right":
            if not xrt.get_right_hand_is_active():
                return None
            return np.asarray(xrt.get_right_hand_tracking_state(), dtype=np.float64)
        raise ValueError(f"Invalid side: {side}. Choose 'left' or 'right'.")

    # ------------------------------------------------------------------
    # Timestamp
    # ------------------------------------------------------------------

    def get_timestamp_ns(self) -> int:
        """Return the current XR system timestamp in nanoseconds."""
        return int(xrt.get_time_stamp_ns())

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def close(self):
        """Shut down the SDK connection."""
        xrt.close()
        print("[XrClient] SDK connection closed.")