#!/usr/bin/python
import rclpy
from rclpy.node import Node
import argparse
from geometry_msgs.msg import PoseStamped

class CartesianControllerPublisher(Node):

    def __init__(self, robot_name: str, poses: dict):
        super().__init__("cartesian_controller_pose_publisher")
        self.publisher_ = self.create_publisher(
            PoseStamped, 
            "/" + robot_name + "/cartesian_motion_controller/target_frame", 10)
        timer_period = 0.5  # seconds
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.robot_name = robot_name
        self.poses = poses

    def timer_callback(self):
        pose = self.poses.get(self.robot_name, {
            'position': {'x': 0.0, 'y': 0.0, 'z': 0.0},
            'orientation': {'x': 0.0, 'y': 0.0, 'z': 0.0, 'w': 1.0}
        })

        msg = PoseStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "base_link"
        msg.pose.position.x = pose['position']['x']
        msg.pose.position.y = pose['position']['y']
        msg.pose.position.z = pose['position']['z']
        msg.pose.orientation.x = pose['orientation']['x']
        msg.pose.orientation.y = pose['orientation']['y']
        msg.pose.orientation.z = pose['orientation']['z']
        msg.pose.orientation.w = pose['orientation']['w']

        self.publisher_.publish(msg)


def main(args=None):
    rclpy.init(args=args)

    parser = argparse.ArgumentParser(description='ROS2 Minimal Publisher')
    parser.add_argument('--robot-name', type=str, default='arm0', help='Specify the robot name')
    args = parser.parse_args()

    # Define poses for each robot
    poses = {
        'arm2': {
            'position': {'x': -0.3, 'y': 0.0, 'z': 0.45},
            'orientation': {'x': 0.0, 'y': -0.707, 'z': 0.0, 'w': 0.707}
        },
        'arm1': {
            'position': {'x': 0.3, 'y': 0.0, 'z': 0.45},
            'orientation': {'x': 0.0, 'y': 0.707, 'z': 0.0, 'w': 0.707}
        }
        # Add more robots and their poses as needed
    }

    minimal_publisher = CartesianControllerPublisher(args.robot_name, poses)

    rclpy.spin(minimal_publisher)

    minimal_publisher.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()