import time
import struct
import can


def can_id_single_motor(motor_id: int) -> int:
    """
    单电机命令：标识符 = 0x140 + ID(1~32)
    """
    assert 1 <= motor_id <= 32
    return 0x140 + motor_id


def send_run(bus: can.Bus, motor_id: int):
    """
    电机运行命令：DATA[0]=0x88，其余 0
    """
    arb_id = can_id_single_motor(motor_id)
    data = bytes([0x88, 0, 0, 0, 0, 0, 0, 0])
    msg = can.Message(arbitration_id=arb_id, data=data, is_extended_id=False)
    bus.send(msg)


def send_stop(bus: can.Bus, motor_id: int):
    """
    电机停止命令：DATA[0]=0x81，其余 0
    """
    arb_id = can_id_single_motor(motor_id)
    data = bytes([0x81, 0, 0, 0, 0, 0, 0, 0])
    msg = can.Message(arbitration_id=arb_id, data=data, is_extended_id=False)
    bus.send(msg)


def send_speed_cmd_a2(bus: can.Bus, motor_id: int, speed_dps: float):
    """
    速度闭环控制命令1（0xA2）：
    DATA[0]=0xA2
    DATA[4..7]=speedControl(int32, little-endian), 0.01 dps/LSB

    speed_dps: 目标角速度 (deg/s)
    """
    speed_control = int(round(speed_dps / 0.01))  # 0.01 dps/LSB
    arb_id = can_id_single_motor(motor_id)

    data = bytearray(8)
    data[0] = 0xA2
    data[1] = 0x00
    data[2] = 0x00
    data[3] = 0x00
    data[4:8] = struct.pack("<i", speed_control)  # int32 little-endian

    msg = can.Message(arbitration_id=arb_id, data=data, is_extended_id=False)
    bus.send(msg)


def send_speed_cmd_ad(bus: can.Bus, motor_id: int, speed_dps: float, iq_limit: int):
    """
    速度闭环控制命令2（0xAD，带力矩/转矩电流限制）：
    DATA[0]=0xAD
    DATA[2..3]=iqControl(int16, little-endian), 范围 -2048~2048
    DATA[4..7]=speedControl(int32, little-endian), 0.01 dps/LSB
    """
    iq_limit = int(iq_limit)
    if not (-2048 <= iq_limit <= 2048):
        raise ValueError("iq_limit must be in [-2048, 2048]")

    speed_control = int(round(speed_dps / 0.01))
    arb_id = can_id_single_motor(motor_id)

    data = bytearray(8)
    data[0] = 0xAD
    data[1] = 0x00
    data[2:4] = struct.pack("<h", iq_limit)       # int16 little-endian
    data[4:8] = struct.pack("<i", speed_control)  # int32 little-endian

    msg = can.Message(arbitration_id=arb_id, data=data, is_extended_id=False)
    bus.send(msg)


def read_status2(bus: can.Bus, motor_id: int, timeout=0.2):
    """
    读取电机状态2命令（0x9C），回复里包含：
    temperature(int8), iq/power(int16), speed(int16, 1 dps/LSB), encoder(uint16)
    注意：这里的 speed 是“反馈转速”，分辨率 1 dps/LSB（协议里写 int16，1dps/LSB）
    """
    arb_id = can_id_single_motor(motor_id)

    # 发请求
    req = can.Message(arbitration_id=arb_id, data=bytes([0x9C, 0, 0, 0, 0, 0, 0, 0]), is_extended_id=False)
    bus.send(req)

    # 收回复（同 arb_id）
    t_end = time.time() + timeout
    while time.time() < t_end:
        msg = bus.recv(timeout=timeout)
        if msg is None:
            break
        if msg.arbitration_id != arb_id or len(msg.data) != 8:
            continue
        if msg.data[0] != 0x9C:
            continue

        temperature = struct.unpack("<b", msg.data[1:2])[0]
        iq_or_power = struct.unpack("<h", msg.data[2:4])[0]
        speed_fb_dps = struct.unpack("<h", msg.data[4:6])[0]  # 1 dps/LSB
        encoder = struct.unpack("<H", msg.data[6:8])[0]
        return {
            "temperature_C": temperature,
            "iq_or_power_raw": iq_or_power,
            "speed_fb_dps": speed_fb_dps,
            "encoder": encoder,
        }

    return None


def rpm_to_dps(rpm: float) -> float:
    # 1 rev = 360 deg; rpm -> deg/s
    return rpm * 360.0 / 60.0


def main():
    # 1) 配置你的 CAN 接口（按你的系统改 channel / bustype）
    # 常见：socketcan: channel="can0"
    # bus = can.Bus(interface="socketcan", channel="can0", bitrate=1000000)
    bus = can.Bus(interface="canalystii", channel=0, bitrate=1000000)

    motor_id = 2
    # 2) 使能运行
    send_run(bus, motor_id)
    time.sleep(0.05)

    # 3) 恒速：例如 30 rpm
    target_rpm = 100.0
    target_dps = rpm_to_dps(target_rpm)

    # 推荐：循环周期性刷新速度命令（例如 50~200Hz），更稳

    try:
        last_print = 0.0
        while True:
            # 方案A：纯速度闭环（0xA2）
            send_speed_cmd_a2(bus, 1, target_dps)
            send_speed_cmd_a2(bus, 2, target_dps)
            send_speed_cmd_a2(bus, 3, target_dps)
            # send_speed_cmd_a2(bus, 4, target_dps)

            # 方案B：带转矩电流限制（0xAD），例如限制 iq=300
            # send_speed_cmd_ad(bus, motor_id, target_dps, iq_limit=300)

            time.sleep(0.01)  # 100 Hz

            if time.time() - last_print > 0.1:
                st1 = read_status2(bus, 1)
                st2 = read_status2(bus, 2)
                st3 = read_status2(bus, 3)
                # st4 = read_status2(bus, 4)
                print("status:", [st1, st2, st3])
                last_print = time.time()

    except KeyboardInterrupt:
        send_stop(bus, 1)
        send_stop(bus, 2)
        send_stop(bus, 3)
        # send_stop(bus, 4)


if __name__ == "__main__":
    main()