import can
import math
import struct
from typing import Dict, List, Tuple, Optional

# 广播控制
BROADCAST_TORQUE_CMD_ID        = 0x280
BROADCAST_VELOCITY_CMD_ID      = 0x281
BROADCAST_POSITION_CMD_ID      = 0x282
BROADCAST_MIX_CMD_ID           = 0x288
BROADCAST_CLEAR_ERR_CMD_ID     = 0x9B
BROADCAST_TURN_ON_CMD_ID       = 0x88
BROADCAST_TURN_OFF_CMD_ID      = 0x80
BROADCAST_STOP_CMD_ID          = 0x81
BROADCAST_GET_STATE_CMD_ID     = 0x9C
# 单电机控制
SINGLE_MS_TORQUE_CMD_ID        = 0xA0
SINGLE_MG_TORQUE_CMD_ID        = 0xA1
SINGLE_VELOCITY_CMD_1_ID       = 0xA2
SINGLE_VELOCITY_CMD_2_ID       = 0xAD
SINGLE_MULTITURN_POS_CMD_1_ID  = 0xA3
SINGLE_MULTITURN_POS_CMD_2_ID  = 0xA4
SINGLE_ONETURN_POS_CMD_1_ID    = 0xA5
SINGLE_ONETURN_POS_CMD_2_ID    = 0xA6
READ_MULTITURN_POS_CMD_ID      = 0x92
SET_MULTITURN_POS_CMD_ID       = 0x95
SINGLE_TURN_ON_CMD_ID          = 0x88
SINGLE_TURN_OFF_CMD_ID         = 0x80
SINGLE_STOP_CMD_ID             = 0x81
SINGLE_CLEAR_ERR_CMD_ID        = 0x9B


def _clamp_int16(x: int) -> int:
    return max(-32768, min(32767, int(x)))

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

def _pack_i16_le(x: int) -> Tuple[int, int]:
    x = _clamp_int16(x)
    u = x & 0xFFFF
    return u & 0xFF, (u >> 8) & 0xFF

def build_torque_broadcast(torques: List[int]) -> can.Message:
    """
    广播控制模式速度控制can帧构建
    torques: 电机电流控制量
    """
    # TODO： 不同电机型号的力矩与电流int16数值的换算
    t = [0, 0, 0, 0]
    for i, torque in enumerate(torques):
        t[i] = torque

    d0, d1 = _pack_i16_le(t[0])
    d2, d3 = _pack_i16_le(t[1])
    d4, d5 = _pack_i16_le(t[2])
    d6, d7 = _pack_i16_le(t[3])
    return can.Message(
        arbitration_id=BROADCAST_TORQUE_CMD_ID,
        is_extended_id=False,
        data=[d0, d1, d2, d3, d4, d5, d6, d7],
    )

def build_velocity_broadcast(vels: List[float]) -> can.Message:
    """
    广播控制模式速度控制can帧构建
    vels: 电机速度，单位：弧度/秒
    """
    v = [0, 0, 0, 0]
    for i, vel in enumerate(vels):
        v[i] = int(round(math.degrees(vel)))
    d0, d1 = _pack_i16_le(v[0])
    d2, d3 = _pack_i16_le(v[1])
    d4, d5 = _pack_i16_le(v[2])
    d6, d7 = _pack_i16_le(v[3])
    return can.Message(
        arbitration_id=BROADCAST_VELOCITY_CMD_ID,
        is_extended_id=False,
        data=[d0, d1, d2, d3, d4, d5, d6, d7],
    )

def build_position_broadcast(rads: List[float]) -> can.Message:
    """
    广播控制模式位置控制can帧构建
    rad1: 电机1弧度
    rad2：电机2弧度
    rad2：电机3弧度
    rad2：电机4弧度
    """
    p = [0, 0, 0, 0]
    for i, rad in enumerate(rads):
        p[i] = int(round(math.degrees(rad) / 0.01))
    d0, d1 = _pack_i16_le(p[0])
    d2, d3 = _pack_i16_le(p[1])
    d4, d5 = _pack_i16_le(p[2])
    d6, d7 = _pack_i16_le(p[3])
    return can.Message(
        arbitration_id=BROADCAST_POSITION_CMD_ID,
        is_extended_id=False,
        data=[d0, d1, d2, d3, d4, d5, d6, d7],
    )

def build_turn_on_broadcast() -> can.Message:
    return can.Message(
        arbitration_id=BROADCAST_MIX_CMD_ID,
        is_extended_id=False,
        data=[
            BROADCAST_TURN_ON_CMD_ID, 0x00,
            BROADCAST_TURN_ON_CMD_ID, 0x00,
            BROADCAST_TURN_ON_CMD_ID, 0x00,
            BROADCAST_TURN_ON_CMD_ID, 0x00]
    )

def build_turn_off_broadcast() -> can.Message:
    return can.Message(
        arbitration_id=BROADCAST_MIX_CMD_ID,
        is_extended_id=False,
        data=[
            BROADCAST_TURN_OFF_CMD_ID, 0x00,
            BROADCAST_TURN_OFF_CMD_ID, 0x00,
            BROADCAST_TURN_OFF_CMD_ID, 0x00,
            BROADCAST_TURN_OFF_CMD_ID, 0x00]
    )

def build_stop_broadcast() -> can.Message:
    return can.Message(
        arbitration_id=BROADCAST_MIX_CMD_ID,
        is_extended_id=False,
        data=[
            BROADCAST_STOP_CMD_ID, 0x00,
            BROADCAST_STOP_CMD_ID, 0x00,
            BROADCAST_STOP_CMD_ID, 0x00,
            BROADCAST_STOP_CMD_ID, 0x00
        ]
    )

def build_clear_error_broadcast() -> can.Message:
    return can.Message(
        arbitration_id=BROADCAST_MIX_CMD_ID,
        is_extended_id=False,
        data=[
            BROADCAST_CLEAR_ERR_CMD_ID, 0x00,
            BROADCAST_CLEAR_ERR_CMD_ID, 0x00,
            BROADCAST_CLEAR_ERR_CMD_ID, 0x00,
            BROADCAST_CLEAR_ERR_CMD_ID, 0x00
        ]
    )

def build_get_state_broadcast() -> can.Message:
    return can.Message(
        arbitration_id=BROADCAST_MIX_CMD_ID,
        is_extended_id=False,
        data=[
            BROADCAST_GET_STATE_CMD_ID, 0x00,
            BROADCAST_GET_STATE_CMD_ID, 0x00,
            BROADCAST_GET_STATE_CMD_ID, 0x00,
            BROADCAST_GET_STATE_CMD_ID, 0x00
        ]
    )

def build_torque_single(motor_id: int, iq: int, is_MS: bool) -> can.Message:
    """
    单电机控制模式力矩控制can帧构建
    motor_id: 电机id
    iq： 电流
    """
    # TODO： 不同电机型号的力矩与电流int16数值的换算
    d4, d5 = _pack_i16_le(iq)
    if is_MS:
        return can.Message(
            arbitration_id=_can_id_single_motor(motor_id),
            is_extended_id=False,
            data=[SINGLE_MS_TORQUE_CMD_ID, 0x00, 0x00, 0x00, d4, d5, 0x00, 0x00],
        )
    else:
        return can.Message(
            arbitration_id=_can_id_single_motor(motor_id),
            is_extended_id=False,
            data=[SINGLE_MG_TORQUE_CMD_ID, 0x00, 0x00, 0x00, d4, d5, 0x00, 0x00],
        )

def build_velocity_single(motor_id: int, speed_rps: float) -> can.Message:
    """
    单电机控制模式速度控制can帧构建
    motor_id: 电机id
    speed_rps： 电机转速，单位 弧度/秒
    """
    speed_control = int(round(math.degrees(speed_rps) / 0.01))  # 0.01 dps/LSB
    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
    return  can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data =data,
    )

def build_set_multiturn_position_single(motor_id: int, rad: float) -> can.Message:
    """
    单电机控制模式多圈位置控制can帧构建
    motor_id: 电机id
    rad： 电机多圈角度，单位 弧度
    """
    multiturn_pos = int(round(math.degrees(rad) / 0.01))
    data = bytearray(8)
    data[0] = SET_MULTITURN_POS_CMD_ID
    data[1] = 0x00
    data[2] = 0x00
    data[3] = 0x00
    data[4:8] = struct.pack("<i", multiturn_pos)  # int32 little-endian
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=data,
    )

def build_get_multiturn_position_single(motor_id: int) -> can.Message:
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=[READ_MULTITURN_POS_CMD_ID, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
    )

def build_turn_on_single(motor_id: int) -> can.Message:
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=[SINGLE_TURN_ON_CMD_ID, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
    )

def build_turn_off_single(motor_id: int) -> can.Message:
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=[SINGLE_TURN_OFF_CMD_ID, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
    )

def build_stop_single(motor_id: int) -> can.Message:
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=[SINGLE_STOP_CMD_ID, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
    )

def build_clear_error_single(motor_id: int) -> can.Message:
    return can.Message(
        arbitration_id=_can_id_single_motor(motor_id),
        is_extended_id=False,
        data=[SINGLE_CLEAR_ERR_CMD_ID, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
    )