"""
ROS2 → ZMQ bridge
Subscribes to /cmd_vel and holosoma/other_input (Python 3.8 + ROS2 Foxy)
Forwards to Holosoma via ZMQ PUB socket (Python 3.10 hsinference)

Run: source /opt/ros/foxy/setup.bash
     python3.8 ~/Models_marcus/ros2_zmq_bridge.py
"""
import json, time
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import TwistStamped
from std_msgs.msg import String
import zmq

ZMQ_PORT = 5556

class ROS2ZMQBridge(Node):
    def __init__(self):
        super().__init__('marcus_zmq_bridge')

        # ZMQ PUB socket
        self._ctx  = zmq.Context()
        self._sock = self._ctx.socket(zmq.PUB)
        self._sock.bind(f"tcp://127.0.0.1:{ZMQ_PORT}")
        time.sleep(0.3)  # let subscribers connect

        # ROS2 subscribers
        self.create_subscription(
            TwistStamped, 'cmd_vel', self._vel_cb, 10)
        self.create_subscription(
            String, 'holosoma/other_input', self._cmd_cb, 10)

        self.get_logger().info(
            f"ROS2→ZMQ bridge ready on tcp://127.0.0.1:{ZMQ_PORT}")

    def _vel_cb(self, msg: TwistStamped):
        data = {"vel": {
            "vx":   msg.twist.linear.x,
            "vy":   msg.twist.linear.y,
            "vyaw": msg.twist.angular.z,
        }}
        self._sock.send_string(json.dumps(data))
        self.get_logger().info(
            f"Vel → vx={data['vel']['vx']:.2f} "
            f"vy={data['vel']['vy']:.2f} "
            f"vyaw={data['vel']['vyaw']:.2f}")

    def _cmd_cb(self, msg: String):
        data = {"cmd": msg.data}
        self._sock.send_string(json.dumps(data))
        self.get_logger().info(f"Cmd → {msg.data}")

def main():
    rclpy.init()
    node = ROS2ZMQBridge()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == "__main__":
    main()