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

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.declare_parameter('can_interface', 'can0')
        self.get_logger().info('正在初始化CAN接口...')
        can_interface = self.get_parameter('can_interface').value
        self.bus = can.interface.Bus(channel=can_interface, bustype='socketcan')
        # ROS2 publisher: 发布 JointState 到话题 'joint_states'
        self.publisher_ = self.create_publisher(JointState, 'joint_sensor_states', 10)

        self.joint_names = [f'joint{i}' for i in range(1, 8)]
        # 7路传感器 CAN ID -> 7个关节索引
        self.can_to_joint_index = {
            0x01: 0,
            0x02: 1,
            0x03: 2,
            0x04: 3,
            0x05: 4,
            0x06: 5,
            0x07: 6,
        }
        self.current_positions = [0.0] * 7
        self.current_velocities = [0.0] * 7
        self.has_position = [False] * 7
        self.last_wrapped_positions = [0.0] * 7
        self.has_wrapped_position = [False] * 7
        # 传感器原始单位：'centi-rad' 表示原始值/100后得到弧度；'centi-deg' 表示原始值/100后得到角度
        self.raw_angle_unit = 'centi-rad'

        # 每个关节零点偏置（弧度），发布角度时会先减去该偏置
        self.declare_parameter('zero_offsets', [0.0] * 7)
        # 每个关节转向系数，1.0 为正向，-1.0 为反向
        self.declare_parameter('sensor_directions', [1.0] * 7)
        self.declare_parameter('auto_zero_on_first_sample', False)
        self.declare_parameter('calibrate_now', False)
        self.zero_offsets = self._load_zero_offsets_from_param()
        self.sensor_directions = self._load_sensor_directions_from_param()
        self.auto_zero_on_first_sample = bool(self.get_parameter('auto_zero_on_first_sample').value)
        self.auto_zero_done = [False] * 7
        self.add_on_set_parameters_callback(self.on_set_parameters)

        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().info(f"Received: id=0x{msg.arbitration_id:X} dlc={msg.dlc} data={msg.data.hex()} ts={msg.timestamp}")

            joint_index = self.can_to_joint_index.get(msg.arbitration_id)
            if joint_index is None:
                self.get_logger().debug(f"忽略未映射CAN ID: 0x{msg.arbitration_id:X}")
                return

            # 每条CAN报文默认对应一个关节角和速度
            angle, vel = self.parse_sensor_from_can(msg.data)
            direction = self.sensor_directions[joint_index]
            angle = self.wrap_to_pi(angle * direction)
            vel = vel * direction

            self.last_wrapped_positions[joint_index] = angle
            self.has_wrapped_position[joint_index] = True

            if self.auto_zero_on_first_sample and not self.auto_zero_done[joint_index]:
                self.zero_offsets[joint_index] = angle
                self.auto_zero_done[joint_index] = True

            angle = self.wrap_to_pi(angle - self.zero_offsets[joint_index])
            angle = self.unwrap_joint_angle(joint_index, angle)
            self.current_positions[joint_index] = angle
            self.current_velocities[joint_index] = vel

            # 发布完整7关节状态（位置/速度单位均为弧度）
            js = JointState()
            js.header.stamp = self.get_clock().now().to_msg()
            js.name = self.joint_names
            js.position = self.current_positions
            js.velocity = self.current_velocities
            self.publisher_.publish(js)

            self.get_logger().info(
                f"Published JointState: id=0x{msg.arbitration_id:X} "
                f"joint={self.joint_names[joint_index]} pos={angle:.6f} vel={vel:.6f}"
            )
        except Exception as e:
            self.get_logger().error(f"Error handling CAN msg: {e}")

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

    def _load_zero_offsets_from_param(self):
        raw_offsets = self.get_parameter('zero_offsets').value
        if not isinstance(raw_offsets, (list, tuple)) or len(raw_offsets) != 7:
            self.get_logger().warn('参数 zero_offsets 长度不是7,已使用默认零偏。')
            return [0.0] * 7
        return [float(v) for v in raw_offsets]

    def _load_sensor_directions_from_param(self):
        raw_dirs = self.get_parameter('sensor_directions').value
        if not isinstance(raw_dirs, (list, tuple)) or len(raw_dirs) != 7:
            self.get_logger().warn('参数 sensor_directions 长度不是7,已使用默认方向。')
            return [1.0] * 7

        dirs = [float(v) for v in raw_dirs]
        for i, v in enumerate(dirs):
            if abs(v) < 1e-9:
                self.get_logger().warn(f'sensor_directions[{i}] 不能为0，已回退为1.0')
                dirs[i] = 1.0
        return dirs

    def calibrate_zero_from_current(self):
        updated = 0
        for i in range(7):
            if self.has_wrapped_position[i]:
                self.zero_offsets[i] = self.last_wrapped_positions[i]
                self.has_position[i] = False
                updated += 1

        if updated > 0:
            self.set_parameters([
                rclpy.parameter.Parameter(
                    'zero_offsets',
                    rclpy.Parameter.Type.DOUBLE_ARRAY,
                    self.zero_offsets,
                )
            ])
            self.get_logger().info(f'零点校准完成，已更新 {updated}/7 个关节零偏。')
        else:
            self.get_logger().warn('零点校准失败：当前还没有可用的关节采样数据。')

    def on_set_parameters(self, params):
        for param in params:
            if param.name == 'zero_offsets':
                if not isinstance(param.value, (list, tuple)) or len(param.value) != 7:
                    return SetParametersResult(successful=False, reason='zero_offsets 必须是长度为7的数组')

            if param.name == 'sensor_directions':
                if not isinstance(param.value, (list, tuple)) or len(param.value) != 7:
                    return SetParametersResult(successful=False, reason='sensor_directions 必须是长度为7的数组')
                for i, v in enumerate(param.value):
                    if abs(float(v)) < 1e-9:
                        return SetParametersResult(successful=False, reason=f'sensor_directions[{i}] 不能为0')

            if param.name == 'auto_zero_on_first_sample':
                self.auto_zero_on_first_sample = bool(param.value)

            if param.name == 'calibrate_now' and bool(param.value):
                self.calibrate_zero_from_current()

        for param in params:
            if param.name == 'zero_offsets':
                self.zero_offsets = [float(v) for v in param.value]
                self.has_position = [False] * 7

            if param.name == 'sensor_directions':
                self.sensor_directions = [float(v) for v in param.value]
                self.has_position = [False] * 7

        return SetParametersResult(successful=True)

    def wrap_to_pi(self, angle_rad: float) -> float:
        # 归一化到 [-pi, pi)，解决 0~2pi 到 ±pi 的过零映射问题
        return (angle_rad + math.pi) % (2.0 * math.pi) - math.pi

    def unwrap_joint_angle(self, joint_index: int, wrapped_angle: float) -> float:
        # 基于上一帧做相位展开，避免在 -pi/+pi 边界发生数值跳变
        if not self.has_position[joint_index]:
            self.has_position[joint_index] = True
            return wrapped_angle

        prev = self.current_positions[joint_index]
        prev_wrapped = self.wrap_to_pi(prev)
        delta = wrapped_angle - prev_wrapped

        if delta > math.pi:
            delta -= 2.0 * math.pi
        elif delta < -math.pi:
            delta += 2.0 * math.pi

        return prev + delta

    def raw_to_rad(self, raw_value: int) -> float:
        if self.raw_angle_unit == 'centi-deg':
            return math.radians(raw_value / 100.0)
        # 默认按 centi-rad 处理，适配你当前 0~6.28 弧度输出
        return raw_value / 100.0

    def parse_sensor_from_can(self, data: bytes):
        if len(data) == 0:
            return 0.0, 0.0

        # 特殊格式：0xAA + 16bit角度 + 16bit速度
        if data[0] == 0xAA:
            if len(data) < 5:
                return 0.0, 0.0
            raw_deg = (data[1] << 8) | data[2]
            raw_vel = (data[3] << 8) | data[4]
            vel_deg = raw_vel / 100.0
            deg = (raw_deg / 4096.0) * 360.0
            rad = self.wrap_to_pi(math.radians(deg))
            vel_rad = math.radians(vel_deg)
            return rad, vel_rad

        if len(data) < 4:
            return 0.0, 0.0

        # 通用格式：前2字节为角度，后2字节为速度，默认按 centi-rad 解析
        raw_deg = int.from_bytes(data[0:2], byteorder='big', signed=True)
        rad = self.wrap_to_pi(self.raw_to_rad(raw_deg))

        raw_vel = int.from_bytes(data[2:4], byteorder='big', signed=True)
        vel_rad = self.raw_to_rad(raw_vel)

        return rad, vel_rad
        

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

if __name__ == '__main__':
    main()