import rclpy
from rclpy.node import Node
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from osc4py3.as_eventloop import *
from osc4py3 import oscmethod as osm
import xml.etree.ElementTree as ET

class ScaledJointTrajectoryPublisher(Node):
    """Node to publish joint trajectories based on OSC messages."""
    def __init__(self, joint_names):
        super().__init__('scaled_joint_trajectory_publisher')
        
        # ROS2 Publisher
        self.publisher = self.create_publisher(
            JointTrajectory, 
            '/scaled_joint_trajectory_controller/joint_trajectory', 
            10
        )

        # Store received joint positions
        self.joint_positions = [0.0] * len(joint_names)
        self.joint_names = joint_names

        osc_startup()
        osc_udp_server("0.0.0.0", 8000, "osc_server")
        print("Server started on 0.0.0.0:8000 \nready to receive messages in the following format: \n   /joint_angles [joint_positions]; optional: duration as last element, default is 3sec")
        # Register OSC handler
        osc_method("/joint_angles", self.joint_angles_handler, argscheme=osm.OSCARG_DATAUNPACK)

    def joint_angles_handler(self, *args):
        """Handles incoming OSC messages for joint positions."""
        if len(args) == len(self.joint_positions): 
            self.joint_positions = args
            self.send_trajectory(self.joint_positions)
        elif len(args) == len(self.joint_positions) + 1:
            self.joint_positions = args[:-1]
            self.send_trajectory(self.joint_positions, args[-1])
            print(f'Duration: {args[-1]}')



    def send_trajectory(self, joint_positions, duration=3.0):
        """Publish a joint trajectory command to move the robot."""
        msg = JointTrajectory()
        msg.joint_names = self.joint_names
        point = JointTrajectoryPoint()
        point.positions = joint_positions  # Updated joint positions
        point.time_from_start.sec = int(duration)
        point.time_from_start.nanosec = int((duration - int(duration)) * 1e9)
        
        msg.points.append(point)
        self.publisher.publish(msg)
        print(f"Updated joint positions: {self.joint_positions}")


def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""

    tree = ET.parse('/BA/robot.urdf')
    root = tree.getroot()
    joint_names = [joint.get('name') for joint in root.findall('joint') if joint.get('type') == 'revolute' or joint.get('type') == 'continuous' or joint.get('type') == 'prismatic']

    rclpy.init()

    node = ScaledJointTrajectoryPublisher(joint_names)
    
    # Run both ROS 2 and OSC Server together
    try:
        while rclpy.ok():
            osc_process()  # Handle one OSC request at a time
            rclpy.spin_once(node, timeout_sec=0.1)  # Process ROS callbacks
    except KeyboardInterrupt:
        print("")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()
