import rclpy
from rclpy.node import Node
import can
import math
from sensor_msgs.msg import JointState

class CallbackListener(can.Listener):
    def __init__(self, user_callback):
        super().__init__()
        self.user_callback = user_callback

    def on_message_received(self, msg: can.Message):
        # 这是在 Notifier 的线程里调用的 -> 回调应该尽量快
        try:
            self.user_callback(msg)
        except Exception as e:
            # 捕获回调内异常，避免中断 Notifier 分发
            print(f"Callback error: {e}")


class JointSensorNode(Node):
    def __init__(self):
        super().__init__('joint_sensor_node')
        self.get_logger().info('正在初始化CAN接口...')
        self.bus = can.interface.Bus(channel='can0', bustype='socketcan')
        # ROS2 publisher: 发布 JointState 到话题 'joint_states'
        self.publisher_ = self.create_publisher(JointState, 'joint_sensor_states', 10)

        # 可配置的 CAN ID -> joint name 列表映射
        # 默认示例：ID 0x123 包含两个关节角，每个角用 int16 存储（单位 centi-degrees，除以100得到度）
        # 你可以根据真实硬件修改此映射或解析函数
        self.can_to_joints = {
            0x01: ['joint1'],
        }

        self.listener = CallbackListener(self.listen_can_callback)
        self.notifier = can.Notifier(self.bus, [self.listener])
    
    def listen_can_callback(self, msg: can.Message):
        # 用户自定义的回调函数
        try:
            self.get_logger().debug(f"Received: id=0x{msg.arbitration_id:X} dlc={msg.dlc} data={msg.data.hex()} ts={msg.timestamp}")

            # 解析角度数据（默认按每 2 字节 int16，单位 centi-degrees）
            angles = self.parse_angles_from_can(msg.data)

            # 根据 can id 找到 joint 名称列表
            names = self.can_to_joints.get(msg.arbitration_id, [f'joint_{i}' for i in range(len(angles))])

            # 对长度不匹配做适配（截断或填充 0）
            if len(names) != len(angles):
                # 截断或填充名称或角度，保证二者等长
                if len(names) > len(angles):
                    names = names[:len(angles)]
                else:
                    # 增加默认名称
                    names = names + [f'joint_{i}' for i in range(len(names), len(angles))]

            # 发布 JointState（位置使用弧度）
            js = JointState()
            js.header.stamp = self.get_clock().now().to_msg()
            js.name = names
            js.position = angles
            self.publisher_.publish(js)

            self.get_logger().debug(f"Published JointState: ids=0x{msg.arbitration_id:X} names={names} positions={angles}")
        except Exception as e:
            self.get_logger().error(f"Error handling CAN msg: {e}")

    def shutdown(self):
        self.notifier.stop()
        self.bus.shutdown()

    def parse_angles_from_can(self, data: bytes):
        angles = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]  # 默认两个关节角，单位弧度
        if data[0] == 0xaa:  # 假设 0xaa 开头表示特殊格式，使用不同解析方式
            raw = (data[2] << 8) | data[3]
            deg = (raw / 4096.0) * 360.0  # 例如特殊格式直接除以 4096 得到度
            rad = math.radians(deg) - 3.14
            self.get_logger().info(f"Special format detected: raw={raw} deg={deg} rad={rad}")
            angles[0] = rad  # 假设特殊格式只包含一个关节角，放在第一个位置
        return angles
        

def main():
    rclpy.init()
    node = JointSensorNode()
    rclpy.spin(node)
    node.shutdown()
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()