urdf-viz 0.35.0

use parking_lot::{Mutex, MutexGuard};
use serde::{Deserialize, Serialize};
use std::{ops, sync::Arc};

use crate::utils::RobotModel;

#[derive(Deserialize, Serialize, Debug, Clone, Default)]
pub struct JointNamesAndPositions {
    pub names: Vec<String>,
    pub positions: Vec<f32>,
}

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct RobotOrigin {
    pub position: [f32; 3],
    pub quaternion: [f32; 4],
}

impl Default for RobotOrigin {
    fn default() -> Self {
        Self {
            position: [0.0; 3],
            quaternion: [0.0, 0.0, 0.0, 1.0],
        }
    }
}

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct PointsAndColors {
    pub id: Option<String>,
    pub points: Vec<[f32; 3]>,
    pub colors: Vec<[f32; 3]>,
}

/// Handle to get and modify the state of the robot.
#[derive(Debug, Default)]
pub struct RobotStateHandle {
    target_joint_positions: Mutex<Option<JointNamesAndPositions>>,
    pub(crate) current_joint_positions: Mutex<JointNamesAndPositions>,
    target_robot_origin: Mutex<Option<RobotOrigin>>,
    pub(crate) current_robot_origin: Mutex<RobotOrigin>,
    point_cloud: Mutex<Option<PointsAndColors>>,
    pub(crate) urdf_text: Option<Arc<Mutex<String>>>,
    robot: Mutex<Option<RobotModel>>,
}

// Wrapper type to prevent parking_lot updates from becoming a breaking change.
macro_rules! impl_guard {
    ($guard_name:ident($target:ident)) => {
        #[derive(Debug)]
        pub struct $guard_name<'a>(MutexGuard<'a, $target>);
        impl ops::Deref for $guard_name<'_> {
            type Target = $target;
            fn deref(&self) -> &Self::Target {
                &self.0
            }
        }
        impl ops::DerefMut for $guard_name<'_> {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.0
            }
        }
    };
}
impl_guard!(JointNamesAndPositionsLockGuard(JointNamesAndPositions));
impl_guard!(RobotOriginLockGuard(RobotOrigin));
impl_guard!(UrdfTextLockGuard(String));

impl RobotStateHandle {
    pub fn current_joint_positions(&self) -> JointNamesAndPositionsLockGuard<'_> {
        JointNamesAndPositionsLockGuard(self.current_joint_positions.lock())
    }

    pub fn current_robot_origin(&self) -> RobotOriginLockGuard<'_> {
        RobotOriginLockGuard(self.current_robot_origin.lock())
    }

    pub fn urdf_text(&self) -> Option<UrdfTextLockGuard<'_>> {
        Some(UrdfTextLockGuard(self.urdf_text.as_ref()?.lock()))
    }

    pub fn set_target_joint_positions(&self, joint_positions: JointNamesAndPositions) {
        *self.target_joint_positions.lock() = Some(joint_positions);
    }

    pub fn set_target_robot_origin(&self, robot_origin: RobotOrigin) {
        *self.target_robot_origin.lock() = Some(robot_origin);
    }

    pub fn set_point_cloud(&self, points_and_colors: PointsAndColors) {
        *self.point_cloud.lock() = Some(points_and_colors);
    }

    pub fn set_robot(&self, robot: RobotModel) {
        // set_robot may change name or number of joints, so reset target_joint_positions.
        *self.target_joint_positions.lock() = None;
        *self.robot.lock() = Some(robot);
    }

    pub fn take_target_joint_positions(&self) -> Option<JointNamesAndPositions> {
        self.target_joint_positions.lock().take()
    }

    pub fn take_target_robot_origin(&self) -> Option<RobotOrigin> {
        self.target_robot_origin.lock().take()
    }

    pub(crate) fn take_point_cloud(&self) -> Option<PointsAndColors> {
        self.point_cloud.lock().take()
    }

    pub(crate) fn take_robot(&self) -> Option<RobotModel> {
        self.robot.lock().take()
    }
}