"""Dual-arm XR teleoperation using the independent RM75 QP controller."""

from __future__ import annotations

import os
import threading
from pathlib import Path

import rclpy
from ament_index_python.packages import get_package_share_directory
from geometry_msgs.msg import PoseStamped
from rclpy.executors import ExternalShutdownException
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Bool, String
from std_srvs.srv import Trigger
from xr_rm_interfaces.msg import XrController


class DualArmQpTeleop(Node):
    def __init__(self) -> None:
        super().__init__("dual_arm_qp_teleop")
        bringup_share = Path(get_package_share_directory("xr_rm_bringup"))
        self.declare_parameter("robot_backend", "mujoco")
        self.declare_parameter("control_rate_hz", 90.0)
        self.declare_parameter("left_controller_topic", "/xr/left_controller")
        self.declare_parameter("right_controller_topic", "/xr/right_controller")
        self.declare_parameter(
            "teleop_config_file", str(bringup_share / "config/dual_arm_rm75.yaml")
        )
        self.declare_parameter(
            "peripheral_config_file",
            str(bringup_share / "config/peripherals_rm75.yaml"),
        )
        self.declare_parameter("show_viewer", True)
        self.declare_parameter("debug_topic_prefix", "/xr_rm/qp")

        backend_name = str(self.get_parameter("robot_backend").value).lower()
        if backend_name != "mujoco":
            raise ValueError(
                "stage 3 implements only robot_backend='mujoco'; "
                "the RealmanRobot contract is reserved for the next stage"
            )
        control_rate_hz = float(self.get_parameter("control_rate_hz").value)
        if control_rate_hz <= 0.0:
            raise ValueError("control_rate_hz must be positive")

        from rm75_ik import DualArmQpTeleopController, load_dual_arm_profiles
        from rm75_ik.mujoco_robot import MujocoRobot

        profiles = load_dual_arm_profiles(
            self.get_parameter("teleop_config_file").value,
            self.get_parameter("peripheral_config_file").value,
        )
        self._robot = MujocoRobot(profiles)
        self._controller = DualArmQpTeleopController(
            self._robot, profiles, control_rate_hz=control_rate_hz
        )
        self._closed = False
        self._viewer_reset_requested = threading.Event()
        self._viewer_reset_key_codes: set[int] = set()
        self._show_viewer = bool(self.get_parameter("show_viewer").value)
        if self._show_viewer:
            from mujoco.glfw import glfw

            self._viewer_reset_key_codes = {glfw.KEY_R, glfw.KEY_HOME}
            self._robot.open_viewer(key_callback=self._on_viewer_key)

        prefix = str(self.get_parameter("debug_topic_prefix").value).rstrip("/")
        self._target_publishers = {
            arm: self.create_publisher(
                PoseStamped, f"{prefix}/{arm}/target_tcp", 10
            )
            for arm in ("left", "right")
        }
        self._current_publishers = {
            arm: self.create_publisher(
                PoseStamped, f"{prefix}/{arm}/current_tcp", 10
            )
            for arm in ("left", "right")
        }
        self._fault_publisher = self.create_publisher(Bool, f"{prefix}/fault", 10)
        self._state_publisher = self.create_publisher(
            String, f"{prefix}/control_state", 10
        )
        self._joint_state_publisher = self.create_publisher(
            JointState, f"{prefix}/joint_states", 10
        )
        self._reset_service = self.create_service(
            Trigger,
            f"{prefix}/reset_to_initial",
            self._on_reset_service,
        )
        self._joint_names = [
            f"{arm}_joint_{joint}"
            for arm in ("left", "right")
            for joint in range(1, 8)
        ]
        self.create_subscription(
            XrController,
            str(self.get_parameter("left_controller_topic").value),
            lambda message: self._on_controller("left", message),
            10,
        )
        self.create_subscription(
            XrController,
            str(self.get_parameter("right_controller_topic").value),
            lambda message: self._on_controller("right", message),
            10,
        )
        self.create_timer(1.0 / control_rate_hz, self._control_tick)

        for arm, diagnostic in self._robot.initial_pose_diagnostics.items():
            self.get_logger().info(
                f"{arm} MuJoCo initialized from initial_tcp_pose: "
                f"position_error={diagnostic.solved_position_error_m:.6f} m, "
                f"orientation_error={diagnostic.solved_orientation_error_rad:.6f} rad"
            )
        self.get_logger().info(
            "dual-arm QP teleop ready: backend=mujoco, grip controls each arm, "
            "R/Home resets both arms, any active-arm fault stops both arms"
        )

    def _now_sec(self) -> float:
        return self.get_clock().now().nanoseconds * 1e-9

    def _on_controller(self, arm: str, message: XrController) -> None:
        try:
            self._controller.update_controller(
                message, self._now_sec(), expected_arm=arm
            )
        except Exception as exc:
            self.get_logger().error(
                f"{arm} controller input rejected: {exc}",
                throttle_duration_sec=1.0,
            )
            self._controller.reject_input(str(exc))

    def _on_viewer_key(self, keycode: int) -> None:
        if keycode in self._viewer_reset_key_codes:
            self._viewer_reset_requested.set()

    def _request_reset(self, source: str):
        result = self._controller.reset_to_initial()
        self.get_logger().info(
            f"dual-arm MuJoCo reset requested by {source}; "
            "waiting for fresh left/right grip release"
        )
        return result

    def _on_reset_service(self, request, response):
        del request
        try:
            result = self._request_reset("ROS service")
        except Exception as exc:
            response.success = False
            response.message = str(exc)
        else:
            response.success = True
            response.message = result.reason
        return response

    def _control_tick(self) -> None:
        if self._viewer_reset_requested.is_set():
            self._viewer_reset_requested.clear()
            try:
                self._request_reset("MuJoCo viewer key")
            except Exception as exc:
                self.get_logger().error(f"dual-arm MuJoCo reset failed: {exc}")
        result = self._controller.step(self._now_sec())
        stamp = self.get_clock().now().to_msg()
        fault = Bool()
        fault.data = result.state.value == "fault"
        self._fault_publisher.publish(fault)
        state_message = String()
        state_message.data = result.state.value
        self._state_publisher.publish(state_message)
        if fault.data:
            self.get_logger().error(
                f"dual-arm QP controller stopped: {result.reason}",
                throttle_duration_sec=1.0,
            )
        if result.targets_tcp:
            for arm, target in result.targets_tcp.items():
                self._target_publishers[arm].publish(self._pose_message(target))
        for arm in ("left", "right"):
            self._current_publishers[arm].publish(
                self._pose_message(self._robot.get_tcp_pose(arm))
            )
        joint_state = self._robot.read_joint_positions().positions_rad
        joint_message = JointState()
        joint_message.header.stamp = stamp
        joint_message.header.frame_id = "dual_rm75"
        joint_message.name = self._joint_names
        joint_message.position = [
            float(value)
            for arm in ("left", "right")
            for value in joint_state[arm]
        ]
        self._joint_state_publisher.publish(joint_message)
        if self._show_viewer and not self._robot.sync_viewer():
            self.get_logger().info("MuJoCo viewer closed; shutting down QP teleop")
            self._close_control()
            if rclpy.ok():
                rclpy.shutdown()

    def _pose_message(self, pose) -> PoseStamped:
        import pinocchio as pin

        quaternion = pin.Quaternion(pose.rotation)
        message = PoseStamped()
        message.header.stamp = self.get_clock().now().to_msg()
        message.header.frame_id = "rm_base"
        message.pose.position.x = float(pose.translation[0])
        message.pose.position.y = float(pose.translation[1])
        message.pose.position.z = float(pose.translation[2])
        message.pose.orientation.x = float(quaternion.x)
        message.pose.orientation.y = float(quaternion.y)
        message.pose.orientation.z = float(quaternion.z)
        message.pose.orientation.w = float(quaternion.w)
        return message

    def _close_control(self) -> None:
        if self._closed:
            return
        self._controller.close()
        self._closed = True

    def destroy_node(self):
        self._close_control()
        return super().destroy_node()


def main(args=None) -> None:
    os.environ.setdefault("MUJOCO_GL", "glfw")
    rclpy.init(args=args)
    node = None
    try:
        node = DualArmQpTeleop()
        rclpy.spin(node)
    except (KeyboardInterrupt, ExternalShutdownException):
        pass
    finally:
        if node is not None:
            try:
                node.destroy_node()
            except KeyboardInterrupt:
                pass
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()
