import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from osc4py3.as_eventloop import *
from osc4py3 import oscbuildparse

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
        )

        # Open Sound Control (OSC) Client settings
        self.osc_ip = "127.0.0.1"   # Replace with the target IP
        self.osc_port = 8000        # Replace with the target port
        # Start the OSC system
        osc_startup()

        # Make client channels to send packets
        osc_udp_client(self.osc_ip, self.osc_port, "osc_client")

    def joint_states_callback(self, msg):
        """Callback function to handle incoming joint states."""
        header = msg.header
        joint_names = msg.name
        joint_positions = msg.position
        joint_velocity = msg.velocity
        joint_effort = msg.effort

        joint_names_str = "\n- ".join(joint_names)
        joint_positions_str = "\n- ".join(map(str, joint_positions))
        joint_velocity_str = "\n- ".join(map(str, joint_velocity))
        joint_effort_str = "\n- ".join(map(str, joint_effort))

        info = f"""
---
header:
    stamp:
        sec: {header.stamp.sec}
        nanosec: {header.stamp.nanosec}
name:
- {joint_names_str}
position:
- {joint_positions_str}
velocity:
- {joint_velocity_str}
effort:
- {joint_effort_str}
---"""

        # Send the info message
        msg_info = oscbuildparse.OSCMessage("/joint_states", None, [info])
        msg_name = oscbuildparse.OSCMessage("/joint_states/name", None, [i for i in joint_names])
        msg_position = oscbuildparse.OSCMessage("/joint_states/position", None, [i for i in joint_positions])
        msg_velocity = oscbuildparse.OSCMessage("/joint_states/velocity", None, [i for i in joint_velocity])
        msg_effort = oscbuildparse.OSCMessage("/joint_states/effort", None, [i for i in joint_effort])
        
        bun = oscbuildparse.OSCBundle(oscbuildparse.OSC_IMMEDIATELY, [msg_info, msg_name, msg_position, msg_velocity, msg_effort])
        
        osc_send(bun, "osc_client")
        osc_process()
        #print(f"Publishing: {info}")


'''
        # Send each joint state as an OSC message
        for i, name in enumerate(joint_names):
            #msg_sec = oscbuildparse.OSCMessage(f"/joint_states/header/sec", None, [header.stamp.sec])
            #msg_nanosec = oscbuildparse.OSCMessage(f"/joint_states/header/nanosec", None, [header.stamp.nanosec])
            msg_position = oscbuildparse.OSCMessage(f"/joint_states/{name}/position", None, [joint_positions[i]])
            msg_velocity = oscbuildparse.OSCMessage(f"/joint_states/{name}/velocity", None, [joint_velocity[i]])
            msg_effort = oscbuildparse.OSCMessage(f"/joint_states/{name}/effort", None, [joint_effort[i]])

            bun = oscbuildparse.OSCBundle(oscbuildparse.unixtime2timetag(header.stamp.sec + header.stamp.nanosec), [msg_position, msg_velocity, msg_effort])
            #bun = oscbuildparse.OSCBundle(oscbuildparse.OSC_IMMEDIATELY , [msg_position, msg_velocity, msg_effort])
            osc_send(bun, "osc_client")
            osc_process()
            #print(f"OSC bundle sent for joint {name}")
'''

def main():
    rclpy.init()
    node = JointStateOSC()
    print(f"Publishing joint states to OSC on {node.osc_ip}:{node.osc_port}...")
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        print("shutting down...")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()