import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from pythonosc.udp_client import SimpleUDPClient
from pythonosc import osc_bundle_builder
from pythonosc import osc_message_builder

class JointStateOSC(Node):
    def __init__(self):
        super().__init__('joint_states_osc')

        # Create a ROS 2 subscriber to /joint_states topic
        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1  # Queue size
        )

        self.osc_ip = "127.0.0.1"   # target IP
        self.osc_port = 5005        # target port
        self.osc_client = SimpleUDPClient(self.osc_ip, self.osc_port)

    def joint_states_callback(self, msg):
        """Callback function to handle incoming joint states."""
        header = msg.header
        #print(f"Received joint states at {header.stamp.sec}.{header.stamp.nanosec}")
        joint_names = msg.name
        joint_positions = msg.position
        joint_velocity = msg.velocity
        joint_effort = msg.effort


        bundle = osc_bundle_builder.OscBundleBuilder(timestamp=header.stamp.sec + header.stamp.nanosec * 1e-9)
        #bundle = osc_bundle_builder.OscBundleBuilder(osc_bundle_builder.IMMEDIATELY)

        names = osc_message_builder.OscMessageBuilder(address="/joint_states/names")
        names.add_arg([str(name) for name in joint_names])
        positions = osc_message_builder.OscMessageBuilder(address="/joint_states/positions")
        positions.add_arg([float(pos) for pos in joint_positions])
        velocity = osc_message_builder.OscMessageBuilder(address="/joint_states/velocity")
        velocity.add_arg([float(vel) for vel in joint_velocity])
        effort = osc_message_builder.OscMessageBuilder(address="/joint_states/effort")
        effort.add_arg([float(eff) for eff in joint_effort])

        bundle.add_content(names.build())
        bundle.add_content(positions.build())
        bundle.add_content(velocity.build())
        bundle.add_content(effort.build())

        info = osc_message_builder.OscMessageBuilder(address="/joint_states")
        info.add_arg("joint names:")
        info.add_arg([str(name) for name in joint_names])
        info.add_arg("joint positions:")
        info.add_arg([float(pos) for pos in joint_positions])
        info.add_arg("joint velocity:")
        info.add_arg([float(vel) for vel in joint_velocity])
        info.add_arg("joint effort:")
        info.add_arg([float(eff) for eff in joint_effort])

        bundle.add_content(info.build())

        self.osc_client.send(bundle.build())

def main():
    try: 
        rclpy.init()
        node = JointStateOSC()
        rclpy.spin(node)
    except KeyboardInterrupt: print("")
        #node.destroy_node()
        #rclpy.shutdown()

if __name__ == '__main__':
    main()
