from dataclasses import dataclass
import numpy as np
import enum

import struct
import time
from typing import Any, List, Optional

import can


def uint_to_float(x_int: int, x_min: float, x_max: float, bits: int) -> float:
    """Converts unsigned int to float, given range and number of bits."""
    span = x_max - x_min
    return (x_int * span / ((1 << bits) - 1)) + x_min


def float_to_uint(x: float, x_min: float, x_max: float, bits: int) -> int:
    """Converts a float to an unsigned int, given range and number of bits."""
    span = x_max - x_min
    x = min(x, x_max)
    x = max(x, x_min)
    return int((x - x_min) * ((1 << bits) - 1) / span)

@dataclass
class MotorConstants:
    POSITION_MAX: float = 12.5
    POSITION_MIN: float = -12.5

    VELOCITY_MAX: float = 45
    VELOCITY_MIN: float = -45

    TORQUE_MAX: float = 54
    TORQUE_MIN: float = -54

    ####### Mihgt be used for other motors #######
    CURRENT_MAX: float = 1.0
    CURRENT_MIN: float = -1.0
    KT:float = 1.0
    ##############################

    KP_MAX: float = 500.0
    KP_MIN: float = 0.0
    KD_MAX: float = 5.0
    KD_MIN: float = 0.0

@dataclass
class MotorInfo:
    """Class to represent motor information.

    Attributes:
        id (int): Motor ID.
        target_torque (int): Target torque value.
        vel (float): Motor speed.
        eff (float): Motor current.
        pos (float): Encoder value.
        voltage (float): Motor voltage.
        temperature (float): Motor temperature.

    """

    id: int
    error_code: int
    target_torque: int = 0
    vel: float = 0.0
    eff: float = 0
    pos: float = 0
    voltage: float = -1
    temp_mos: float = -1
    temp_rotor: float = -1


@dataclass
class FeedbackFrameInfo:
    id: int
    error_code: int
    error_message: str
    position: float
    velocity: float
    torque: float
    temperature_mos: float
    temperature_rotor: float


@dataclass
class EncoderInfo:
    encoder = -1
    encoder_raw = -1
    encoder_offset = -1


class MotorErrorCode:
    disabled = 0x0
    normal = 0x1
    over_voltage = 0x8
    under_voltage = 0x9
    over_current = 0xA
    mosfet_over_temperature = 0xB
    motor_over_temperature = 0xC
    loss_communication = 0xD
    overload = 0xE

    # create a dict map error code to error message
    motor_error_code_dict = {
        normal: "normal",
        disabled: "disabled",
        over_voltage: "over voltage",
        under_voltage: "under voltage",
        over_current: "over current",
        mosfet_over_temperature: "mosfet over temperature",
        motor_over_temperature: "motor over temperature",
        loss_communication: "loss communication",
        overload: "overload",
    }
    # covert to decimal
    motor_error_code_dict = {int(k): v for k, v in motor_error_code_dict.items()}

    @classmethod
    def get_error_message(cls, error_code: int) -> str:
        return cls.motor_error_code_dict.get(int(error_code), f"Unknown error code: {error_code}")

class MotorType:
    DM4310 = "DM4310"

    @classmethod
    def get_motor_constants(cls, motor_type: str = None) -> MotorConstants:
        """Get motor constants for DM4310.
        
        Args:
            motor_type: Optional motor type (ignored, kept for compatibility).
        
        Returns:
            MotorConstants for DM4310.
        """
        return MotorConstants(
            POSITION_MAX=12.5,
            POSITION_MIN=-12.5,
            VELOCITY_MAX=30,
            VELOCITY_MIN=-30,
            TORQUE_MAX=7,
            TORQUE_MIN=-7,
            KP_MAX=500.0,
            KP_MIN=0.0,
            KD_MAX=5.0,
            KD_MIN=0.0,
        )


class AutoNameEnum(enum.Enum):
    def _generate_next_value_(name: str, start: int, count: int, last_values: List[str]) -> str:
        return name

class ReceiveMode(AutoNameEnum):
    p16 = enum.auto()
    same = enum.auto()
    zero = enum.auto()
    plus_one = enum.auto()

    def get_receive_id(self, motor_id: int) -> int:
        if self == ReceiveMode.p16:
            return motor_id + 16
        elif self == ReceiveMode.same:
            return motor_id
        elif self == ReceiveMode.zero:
            return 0
        elif self == ReceiveMode.plus_one:
            return motor_id + 1
        else:
            raise NotImplementedError(f"receive_mode: {self} not recognized")

    def to_motor_id(self, receive_id: int) -> int:
        if self == ReceiveMode.p16:
            return receive_id - 16
        elif self == ReceiveMode.same:
            return receive_id
        elif self == ReceiveMode.zero:
            return 0
        else:
            raise NotImplementedError(f"receive_mode: {self} not recognized")

class RawCanInterface:
    def __init__(
        self,
        channel: str = "PCAN_USBBUS1",
        bustype: str = "pcan",
        bitrate: int = 1000000,
        name: str = "default_can_interface",
    ):
        self.bus = can.interface.Bus(bustype=bustype, channel=channel, bitrate=bitrate)
        self.busstate = self.bus.state
        self.name = name

    def close(self) -> None:
        """Shut down the CAN bus."""
        self.bus.shutdown()

    def _send_message_get_response(self, id: int, data: List[int], max_retry: int = 20) -> can.Message:
        self.try_receive_message(id)
        message = can.Message(arbitration_id=id, data=data, is_extended_id=False)
        for _ in range(max_retry):
            try:
                self.bus.send(message)
                response = self._receive_message(data[0])
                return response
            except (can.CanError, AssertionError) as e:
                print(e)
                # print warning in red
                print(
                    "\033[91m"
                    + f"CAN Error {self.name}: Failed to communicate with motor {data[0]} over can bus. Retrying..."
                    + "\033[0m"
                )
            time.sleep(0.002)
        raise AssertionError(
            f"fail to communicate with the motor {id} on {self.name} at can channel {self.bus.channel_info}"
        )

    def try_receive_message(self, motor_id: Optional[int] = None, timeout: float = 0.009) -> can.Message:
        """Try to receive a message from the CAN bus.

        Args:
            timeout (float): The time to wait for a message (in seconds).

        Returns:
            can.Message: The received message, or None if no message is received.
        """
        try:
            return self._receive_message(motor_id, timeout)
        except AssertionError:
            return None

    def _receive_message(self, motor_id: Optional[int] = None, timeout: float = 0.009) -> can.Message:
        """Receive a message from the CAN bus.

        Args:
            timeout (float): The time to wait for a message (in seconds).

        Returns:
            can.Message: The received message.

        Raises:
            AssertionError: If no message is received within the timeout.
        """
        start_time = time.time()
        while (time.time() - start_time) < timeout:
            message = self.bus.recv(timeout=0.002)
            if message:
                return message
        raise AssertionError(
            f"Failed to receive message, {self.name} motor id {motor_id} motor timeout. Check if the motor is powered on or if the motor ID exists."
        )


def bytes_to_uint32(data: bytearray) -> int:
    """Convert the last 4 bytes (data[4:8]) in a bytearray to uint32 (little-endian)."""
    return struct.unpack("<I", data[4:8])[0]


def bytes_to_float32(data: bytearray) -> float:
    """Convert the last 4 bytes (data[4:8]) in a bytearray to float32 (little-endian)."""
    return struct.unpack("<f", data[4:8])[0]


def uint32_to_bytes(value: int) -> bytearray:
    """Convert a uint32 value to a 4-byte bytearray (little-endian).

    Args:
        value (int): A uint32 value to be converted to bytes.

    Returns:
        bytearray: The corresponding 4-byte representation in little-endian order.
    """
    return bytearray(struct.pack("<I", value))


def float32_to_bytes(value: float) -> bytearray:
    """Convert a float32 value to a 4-byte bytearray (little-endian).

    Args:
        value (float): A float32 value to be converted to bytes.

    Returns:
        bytearray: The corresponding 4-byte representation in little-endian order.
    """
    return bytearray(struct.pack("<f", value))


register_addr_map = {
    "KT_value": (1, bytes_to_float32),
    "OT_value": (2, bytes_to_float32),
    "master_id": (7, bytes_to_uint32),
    "id": (8, bytes_to_uint32),
    "timeout": (9, bytes_to_uint32),
    "inertia": (12, bytes_to_float32),
    "sw_ver": (14, bytes_to_uint32),
    "flux": (19, bytes_to_float32),
    "gear_ratio": (20, bytes_to_float32),
    "gear_eff": (30, bytes_to_float32),
}

register_info_map = {"control_mode": {1: "MIT", 2: "pos_speed", 3: "speed", 4: "torque_pos"}}

def get_special_message_response(can_interface: RawCanInterface, motor_id: int, reg_name: str) -> Any:
    """Get the current value of a register from a motor.

    Args:
        can_interface (RawCanInterface): The CAN interface to use.
        motor_id (int): The ID of the motor to read the register value from.
        reg_name (str): The name of the register to read.

    Returns:
        Any: The value read from the register.
    """
    assert reg_name in register_addr_map, f"reg_name {reg_name} not in register_addr_map"
    reg_id, convert_func = register_addr_map[reg_name]
    for _ in range(3):
        try:
            message = can_interface._send_message_get_response(
                0x7FF, [motor_id, 0x00, 0x33, reg_id, 0x00, 0x00, 0x00, 0x00], max_retry=20
            )
            return convert_func(message.data)
        except Exception as e:
            can_interface.try_receive_message(motor_id)
    raise Exception(f"Failed to read {reg_name} of motor {motor_id} after 3 retries")


def write_special_message(can_interface: RawCanInterface, motor_id: int, reg_name: str, data: Any) -> Any:
    """Write a value to a register of a motor.

    Args:
        can_interface (RawCanInterface): The CAN interface to use.
        motor_id (int): The ID of the motor to write the register value for.
        reg_name (str): The name of the register to write.
        data (Any): The value to write to the register.

    Returns:
        Any: The value read from the register after writing.
    """
    assert reg_name in register_addr_map, f"reg_name {reg_name} not in register_addr_map"
    reg_id, convert_func = register_addr_map[reg_name]
    if convert_func is bytes_to_uint32:
        byte_list = uint32_to_bytes(data)
    elif convert_func is bytes_to_float32:
        byte_list = float32_to_bytes(data)
    for _ in range(3):
        try:
            message = can_interface._send_message_get_response(
                0x7FF, [motor_id, 0x00, 0x55, reg_id] + list(byte_list), max_retry=20
            )
            result = convert_func(message.data)
            return result
        except Exception as e:
            can_interface.try_receive_message(motor_id)


def save_to_memory(can_interface: RawCanInterface, motor_id: int, reg_name: str) -> can.Message:
    """Save the current value of a register to the motor's memory.

    Args:
        can_interface (RawCanInterface): The CAN interface to use.
        motor_id (int): The ID of the motor to save the register value for.
        reg_name (str): The name of the register to save.

    Returns:
        can.Message: The response message from the motor.
    """
    assert reg_name in register_addr_map, f"reg_name {reg_name} not in register_addr_map"
    reg_id, convert_func = register_addr_map[reg_name]
    message = can_interface._send_message_get_response(
        0x7FF, [motor_id, 0x00, 0xAA, reg_id, 0x00, 0x00, 0x00, 0x00], max_retry=20
    )
    return message
