import rclpy
from rclpy.node import Node
from rcl_interfaces.srv import GetParameters
import roboticstoolbox as rtb

class URDFRetriever(Node):
    def __init__(self):
        super().__init__('urdf_retriever')

        # Create a parameter client to request the parameter from /robot_state_publisher
        self.client = self.create_client(GetParameters, '/robot_state_publisher/get_parameters')

        while not self.client.wait_for_service(timeout_sec=3.0):
            self.get_logger().warn("Waiting for /robot_state_publisher parameter service...")

        # Create and send a request to get 'robot_description'
        self.request = GetParameters.Request()
        self.request.names = ["robot_description"]

        self.future = self.client.call_async(self.request)
        rclpy.spin_until_future_complete(self, self.future)

        # Extract the URDF
        if self.future.result():
            self.urdf_string = self.future.result().values[0].string_value
            if self.urdf_string:
                self.get_logger().info("Successfully retrieved URDF from /robot_state_publisher!")
            else:
                self.get_logger().error("URDF is empty")
        else:
            self.get_logger().error("Failed to get URDF from /robot_state_publisher.")

def main():
    rclpy.init()
    node = URDFRetriever()
    
    urdf_string = node.urdf_string

    if urdf_string:
        # Save to a file
        urdf_path = "/BA/robot_ur5.urdf"
        with open(urdf_path, "w") as file:
            file.write(urdf_string)

        print(f"URDF saved to {urdf_path}")

        # Load into Robotic Toolbox
        robot = rtb.ERobot.URDF(urdf_path)
        print(robot)

    node.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()
