import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from std_msgs.msg import Float32MultiArray

# Import your motor library
from raspbot_v2_interface import MotorController  

class MotorControllerNode(Node):
    def __init__(self):
        super().__init__('motor_controller')

        # --- Parameters ---
        self.declare_parameter('wheel_base', 0.3)      # meters between left/right wheels
        self.declare_parameter('max_speed', 1.0)       # max motor speed (library units)

        self.wheel_base = self.get_parameter('wheel_base').value
        self.max_speed  = self.get_parameter('max_speed').value

        # --- Initialize your motor library ---
        self.motors = MotorController()
        self.motors.initialize()

        # --- Subscribers ---
        # cmd_vel is the standard ROS2 topic for velocity commands
        self.cmd_vel_sub = self.create_subscription(
            Twist,
            'cmd_vel',
            self.cmd_vel_callback,
            10  # QoS queue depth
        )

        # Optional: direct per-wheel speed control [FL, FR, RL, RR]
        self.wheel_speeds_sub = self.create_subscription(
            Float32MultiArray,
            'wheel_speeds',
            self.wheel_speeds_callback,
            10
        )

        # --- Publishers ---
        # Publish current wheel speeds for feedback
        self.speed_pub = self.create_publisher(Float32MultiArray, 'current_wheel_speeds', 10)

        # --- Timer for telemetry ---
        self.telemetry_timer = self.create_timer(0.1, self.publish_telemetry)  # 10 Hz

        self.get_logger().info('Motor controller node started')

    def cmd_vel_callback(self, msg: Twist):
        """
        Convert a Twist message (linear.x, angular.z) into
        differential-drive wheel speeds for a 4-wheeled robot.
        """
        linear  = msg.linear.x   # m/s forward
        angular = msg.angular.z  # rad/s rotation

        # Differential drive kinematics
        left_speed  = linear - (angular * self.wheel_base / 2.0)
        right_speed = linear + (angular * self.wheel_base / 2.0)

        # Clamp to max speed
        left_speed  = max(-self.max_speed, min(self.max_speed, left_speed))
        right_speed = max(-self.max_speed, min(self.max_speed, right_speed))

        # Send to motors via your library (FL=FR=left, RL=RR=right)
        try:
            self.motors.set_speed(front_left=left_speed,  front_right=right_speed,
                                  rear_left=left_speed,   rear_right=right_speed)
        except Exception as e:
            self.get_logger().error(f'Motor error: {e}')

    def wheel_speeds_callback(self, msg: Float32MultiArray):
        """Direct per-wheel control: [FL, FR, RL, RR]"""
        if len(msg.data) != 4:
            self.get_logger().warn('wheel_speeds expects exactly 4 values [FL, FR, RL, RR]')
            return
        try:
            self.motors.set_speed(front_left=msg.data[0],  front_right=msg.data[1],
                                  rear_left=msg.data[2],   rear_right=msg.data[3])
        except Exception as e:
            self.get_logger().error(f'Motor error: {e}')

    def publish_telemetry(self):
        """Read and publish current wheel speeds from your library."""
        try:
            speeds = self.motors.get_speeds()  # Expected: [FL, FR, RL, RR]
            msg = Float32MultiArray()
            msg.data = [float(s) for s in speeds]
            self.speed_pub.publish(msg)
        except Exception as e:
            self.get_logger().warn(f'Telemetry read failed: {e}')

    def destroy_node(self):
        """Clean shutdown — stop motors before exiting."""
        self.get_logger().info('Shutting down motors...')
        try:
            self.motors.stop_all()
            self.motors.shutdown()
        except Exception:
            pass
        super().destroy_node()


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


if __name__ == '__main__':
    main()