#!/usr/bin/env python3

import rclpy
import math
from tf2_ros import Buffer, TransformListener, TransformException, TransformStamped
from geometry_msgs.msg import Pose, Quaternion, Point
from rbs_mill_interfaces.srv import GetGraspingPose
from rclpy.node import Node
from rclpy.service import Service
from scipy.spatial.transform import Rotation as R

class GetGraspPickPoses(Node):
    def __init__(self) -> None:
        super().__init__("get_grasp_pick_poses")
        self.srv: Service = self.create_service(
            GetGraspingPose, "get_grasping_pose", self.get_grasp_pick_poses
        )

        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        
        # Precompute the rotation correction quaternion
        self.rotation_correction = self.euler_to_quat(0, math.pi, math.pi/2)

    def get_grasp_pick_poses(
        self,
        request: GetGraspingPose.Request,
        response: GetGraspingPose.Response,
    ) -> GetGraspingPose.Response:
        if request.action_type not in ("pick", "place"):
            self.get_logger().error(f"Unsupported action type: {request.action_type}")
            response.ok = False
            return response

        # Common logic for pick/place
        tf = self.get_tf_pose(request.pose_name, request.relative_to)
        #tf_base = self.get_tf_pose(request.pose_name, "base_link")
        if tf is None:
            response.ok = False
            return response

        # Apply transformations
        main_pose, pre_pose, post_pose = self.process_tf_and_generate_poses(tf)

        if request.action_type == "pick":
            response.grasp_poses.pose = main_pose
            response.grasp_poses.pre_pose = pre_pose
            response.grasp_poses.post_pose = post_pose
        else:
            response.place_poses.pose = main_pose
            response.place_poses.pre_pose = pre_pose
            response.place_poses.post_pose = post_pose

        response.ok = True
        return response

    def process_tf_and_generate_poses(self, tf: TransformStamped) -> tuple[Pose, Pose, Pose]:
        """Processes TF and generates main poses."""
        # Apply rotation correction
        tf.transform.rotation = self.multiply_quaternions(
            tf.transform.rotation, 
            self.rotation_correction
        )
        pose = self.get_pose_from_tf(tf)
        return pose, self.add_offset(pose, z=0.1), self.add_offset(pose, z=0.1)

    def multiply_quaternions(self, q1: Quaternion, q2: Quaternion) -> Quaternion:
        """Multiplies two quaternions using scipy."""
        r1 = R.from_quat([q1.x, q1.y, q1.z, q1.w])
        r2 = R.from_quat([q2.x, q2.y, q2.z, q2.w])
        result = (r1 * r2).as_quat()  # type: ignore
        return Quaternion(x=result[0], y=result[1], z=result[2], w=result[3])

    def add_offset(self, pose: Pose, x: float = 0.0, y: float = 0.0, z: float = 0.0) -> Pose:
        """Creates a new pose with offset."""
        return Pose(
            position=Point(
                x=pose.position.x + x,
                y=pose.position.y + y,
                z=pose.position.z + z
            ),
            orientation=pose.orientation
        )

    def euler_to_quat(self, roll: float, pitch: float, yaw: float) -> Quaternion:
        """Converts Euler angles to quaternion."""
        q = R.from_euler("xyz", [roll, pitch, yaw], degrees=False).as_quat()
        return Quaternion(x=q[0], y=q[1], z=q[2], w=q[3])

    def get_pose_from_tf(self, tf: TransformStamped) -> Pose:
        """Creates Pose from TransformStamped."""
        return Pose(
            position=Point(
                x=tf.transform.translation.x,
                y=tf.transform.translation.y,
                z=tf.transform.translation.z
            ),
            orientation=Quaternion(
                x=tf.transform.rotation.x,
                y=tf.transform.rotation.y,
                z=tf.transform.rotation.z,
                w=tf.transform.rotation.w
            )
        )

    def get_tf_pose(self, source_frame: str, target_frame: str) -> TransformStamped | None:
        """Gets coordinate transform between frames."""
        try:
            return self.tf_buffer.lookup_transform(
                target_frame,
                source_frame,
                rclpy.time.Time()
            )
        except TransformException as ex:
            self.get_logger().error(
                f"Transform error: {target_frame} -> {source_frame}: {ex}"
            )
            return None

def main():
    rclpy.init()
    node = GetGraspPickPoses()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == "__main__":
    main()