import time
import struct
import can


def arb_id_single(motor_id: int) -> int:
    # 单电机命令：0x140 + ID(1~32)
    return 0x140 + motor_id


def send_cmd(bus, motor_id: int, data8: bytes):
    msg = can.Message(
        arbitration_id=arb_id_single(motor_id),
        data=data8,
        is_extended_id=False,
    )
    bus.send(msg)


def motor_run(bus, motor_id: int):
    # 电机运行命令：DATA[0]=0x88，其余 0
    send_cmd(bus, motor_id, bytes([0x88, 0, 0, 0, 0, 0, 0, 0]))


def motor_stop(bus, motor_id: int):
    # 电机停止命令：DATA[0]=0x81，其余 0
    send_cmd(bus, motor_id, bytes([0x81, 0, 0, 0, 0, 0, 0, 0]))


def torque_iq_cmd_a1(bus, motor_id: int, iq_control: int):
    """
    转矩闭环控制命令 0xA1：
      DATA[0] = 0xA1
      DATA[4..5] = iqControl (int16, little-endian), 范围 -2048~2048
      其余字节 0
    【协议：0xA1，iqControl:int16，DATA[4]=低字节，DATA[5]=高字节】
    """
    if not (-2048 <= iq_control <= 2048):
        raise ValueError("iq_control must be in [-2048, 2048]")

    data = bytearray(8)
    data[0] = 0xA1
    data[4:6] = struct.pack("<h", int(iq_control))  # little-endian int16
    send_cmd(bus, motor_id, bytes(data))


if __name__ == "__main__":
    # 前提：你已在系统里把 can0 配置并拉起，例如：
    # sudo ip link set can0 type can bitrate 1000000 restart-ms 100
    # sudo ip link set can0 up
    bus = can.Bus(interface="canalystii", channel=0, bitrate=1000000)

    motor_id = 2
    iq = 40  # 示例：给一个正向转矩电流指令（注意范围 -2048~2048）

    try:
        motor_run(bus, motor_id)
        time.sleep(0.05)

        # 建议周期性刷新（例如 100~500 Hz），并在“零力矩”时及时归零
        dt = 0.01  # 100 Hz
        while True:
            torque_iq_cmd_a1(bus, 1, iq)
            torque_iq_cmd_a1(bus, 2, iq)
            torque_iq_cmd_a1(bus, 3, iq)
            # torque_iq_cmd_a1(bus, 4, iq)
            time.sleep(dt)

    except KeyboardInterrupt:
        # 松手/退出时，建议先发 0 力矩，再 stop
        try:
            torque_iq_cmd_a1(bus, 1, 0)
            torque_iq_cmd_a1(bus, 2, 0)
            torque_iq_cmd_a1(bus, 3, 0)
            # torque_iq_cmd_a1(bus, 4, 0)
            time.sleep(0.05)
        except Exception:
            pass
        motor_stop(bus, motor_id)
        bus.shutdown()