import logging
import struct
import time
from typing import Optional
from queue import Queue, Empty
from ....utils.files import write_json_file, read_json_file
from ....comms import CommsManager, SocketResult
from ....utils import constants
from .utils import (
    FeedbackFrameInfo,
    MotorErrorCode,
    MotorType,
    MotorConstants,
    uint_to_float,
    float_to_uint,
    ReceiveMode,
)
import os
import time
from feathersdk.utils.common import currtime
from feathersdk.utils.logger import warning, error, info
from feathersdk.utils.constants import MOTORS_CALIBRATION_PATH
import math
from feathersdk.robot.motors.motor import Motor



class ControlMode:
    MIT = "MIT"
    POS_VEL = "POS_VEL"
    VEL = "VEL"

    @classmethod
    def get_id_offset(cls, control_mode: str) -> int:
        if control_mode == cls.MIT:
            return 0x000
        elif control_mode == cls.POS_VEL:
            return 0x100
        elif control_mode == cls.VEL:
            return 0x200
        else:
            raise ValueError(f"Control mode '{control_mode}' not recognized.")


class DamiaoMotor(Motor):
    """Motor firmware interface for Damiao motors.
    
    Handles all low-level CAN communication with the motor firmware.
    """
    
    def __init__(
        self,
        motor_name: str,
        motor_id: int = 7,
        config: dict = {},
        iface: str = "can0",
    ):
        """Initialize the DamiaoMotor.

        Args:
            motor_name: Name of the motor (for calibration storage).
            motor_id: The CAN ID of the motor.
            motor_type: The type of the motor (MotorType.DM4310).
            iface: The CAN interface name (e.g., "can0").
            control_mode: The control mode (ControlMode.MIT, ControlMode.POS_VEL, ControlMode.VEL).
            receive_mode: The receive mode for response IDs (ReceiveMode.p16, etc.).
        """

        self.motor_name = motor_name
        self.motor_id = config.get("id", 7)
        self.iface = iface
        self.control_mode = ControlMode.MIT #so far only MIT is supported
        #use the version from the config
        self.version = config.get("version", "DM4310")
        if self.version == "DM4310":
            self.motor_type = MotorType.DM4310
        else:
            raise ValueError(f"Motor type '{self.version}' not recognized.")
        self.receive_mode = ReceiveMode.p16
        self.cmd_idoffset = ControlMode.get_id_offset(self.control_mode)

        self._homing_pos = config.get("homing_pos", 0) #i radians
        self.range = config.get("range", {
            "min": -math.pi / 2,
            "max": math.pi / 2,
        })

        self.lower_limit = None
        self.upper_limit = None
        self.total_range = None
        self.calibration_homing_pos = None
        self.calibration_time = None

        # Load motor constants from config, with defaults from MotorType
        self.motor_constants = self._get_motor_constants_from_config(config)

        self.dual_encoder: bool = self._get_dual_encoder_property()
        
        # Use provided CommsManager or create a new one
        self.comms = CommsManager()
        
        # Response queue for matching sent messages with received responses
        self._response_queue: Queue[SocketResult] = Queue()

        self.comms.add_callback(self._on_can_message)

        self.load_calibration_state()

    def _get_motor_constants_from_config(self, config: dict) -> MotorConstants:
        """Get MotorConstants from config, with defaults from MotorType.
        
        Args:
            config: Motor configuration dictionary.
            
        Returns:
            MotorConstants: Motor constants with config values if provided, defaults otherwise.
        """
        # Get default constants
        default_const = MotorType.get_motor_constants()
        
        # Get max values from config, fallback to defaults
        max_velocity = config.get("max_velocity", default_const.VELOCITY_MAX)
        max_torque = config.get("max_torque", default_const.TORQUE_MAX)
        
        # Min values are negative of max
        min_velocity = -max_velocity
        min_torque = -max_torque
        
        # Override with config values if provided
        return MotorConstants(
            POSITION_MAX=default_const.POSITION_MAX,
            POSITION_MIN=default_const.POSITION_MIN,
            VELOCITY_MAX=max_velocity,
            VELOCITY_MIN=min_velocity,
            TORQUE_MAX=max_torque,
            TORQUE_MIN=min_torque,
            KP_MAX=config.get("max_kp", default_const.KP_MAX),
            KP_MIN=default_const.KP_MIN,
            KD_MAX=config.get("max_kd", default_const.KD_MAX),
            KD_MIN=default_const.KD_MIN,
        )

    def _get_dual_encoder_property(self) -> bool:
        if self.motor_type == MotorType.DM4310:
            return True
        else:
            return False

    def _on_can_message(self, result: SocketResult) -> None:
        """Callback for receiving CAN messages from CommsManager.
        
        Filters messages for this motor and adds them to the response queue.
        """
        if result.is_error():
            return
        
        # Only process messages from our CAN interface
        if result.socket_name != self.iface:
            return
        
        # Check if this message is a response for our motor
        expected_receive_id = self.receive_mode.get_receive_id(self.motor_id)
        if result.can_id == expected_receive_id:
            self._response_queue.put(result)

    def _get_frame_id(self) -> int:
        """Calculate the Control Frame ID for this motor."""
        return self.cmd_idoffset + self.motor_id

    def _send_message_get_response(
        self,
        frame_id: int,
        data: bytes,
        max_retry: int = 5,
        expected_id: Optional[int] = None,
        timeout: float = 0.2,
    ) -> SocketResult:
        """Send a CAN message and wait for a response.

        Args:
            frame_id: The CAN frame ID to send.
            data: The data payload (8 bytes).
            max_retry: Maximum number of retry attempts.
            expected_id: Expected response CAN ID. If None, uses receive_mode.
            timeout: Timeout for each attempt in seconds.

        Returns:
            SocketResult: The received response message.

        Raises:
            AssertionError: If no response is received after max_retry attempts.
        """
        if expected_id is None:
            expected_id = self.receive_mode.get_receive_id(self.motor_id)
        
        # Ensure data is exactly 8 bytes
        if len(data) != 8:
            data = data[:8] + bytes(8 - len(data))
        
        for attempt in range(max_retry):
            try:
                # Clear any old messages from the queue
                while not self._response_queue.empty():
                    try:
                        self._response_queue.get_nowait()
                    except Empty:
                        break
                
                # Send the message
                self.comms.cansend(
                    self.iface,
                    extended=False,
                    can_id=frame_id,
                    data=data,
                )
                
                # Wait for response
                try:
                    response = self._response_queue.get(timeout=timeout)
                    if response.can_id == expected_id:
                        return response
                except Empty:
                    pass
                
                # Small delay before retry
                time.sleep(0.001)
                
            except Exception as e:
                warning(f"CAN send error (attempt {attempt + 1}/{max_retry}): {e}")
                time.sleep(0.001)
        
        raise AssertionError(
            f"Failed to communicate with motor {self.motor_id} on {self.iface} "
            f"after {max_retry} attempts"
        )

    def parse_recv_message(
        self,
        result: SocketResult,
        ignore_error: bool = False,
    ) -> FeedbackFrameInfo:
        """Parse the received message to extract motor information.

        Args:
            result: The SocketResult from CommsManager.
            ignore_error: If True, don't raise on error codes.

        Returns:
            FeedbackFrameInfo: The current state of the motor.
        """
        data = result.data
        if len(data) < 8:
            raise ValueError(f"Received message has insufficient data: {len(data)} bytes")
        
        error_int = (data[0] & 0xF0) >> 4
        error_hex = hex(error_int)
        error_message = MotorErrorCode.get_error_message(error_int)

        motor_id_of_this_response = self.receive_mode.to_motor_id(result.can_id)
        
        if error_hex != "0x1":
            warning(
                f"Motor id: {motor_id_of_this_response}, error: {error_message} "
                f"on {self.iface}"
            )
            if not ignore_error:
                error(
                    f"Motor id: {motor_id_of_this_response}, error: {error_message} "
                    f"on {self.iface}"
                )
                raise RuntimeError(
                    f"Motor error detected: motor name: {self.motor_name}, motor id: {motor_id_of_this_response}, "
                    f"error: {error_message}"
                )
        
        p_int = (data[1] << 8) | data[2]
        v_int = (data[3] << 4) | (data[4] >> 4)
        t_int = ((data[4] & 0xF) << 8) | data[5]
        temperature_mos = data[6]
        temperature_rotor = data[7]

        position = uint_to_float(p_int, self.motor_constants.POSITION_MIN, self.motor_constants.POSITION_MAX, 16)
        velocity = uint_to_float(v_int, self.motor_constants.VELOCITY_MIN, self.motor_constants.VELOCITY_MAX, 12)
        torque = uint_to_float(t_int, self.motor_constants.TORQUE_MIN, self.motor_constants.TORQUE_MAX, 12)
        temperature_mos = float(temperature_mos)
        temperature_rotor = float(temperature_rotor)

        return FeedbackFrameInfo(
            id=motor_id_of_this_response,
            error_code=error_hex,
            error_message=error_message,
            position=position,
            velocity=velocity,
            torque=torque,
            temperature_mos=temperature_mos,
            temperature_rotor=temperature_rotor,
        )

    def motor_on(self) -> FeedbackFrameInfo:
        """Turn on the motor.

        Returns:
            FeedbackFrameInfo: The motor state after enabling.
        """
        current_level = logging.getLogger().getEffectiveLevel()
        logging.getLogger().setLevel(logging.ERROR)

        frame_id = self.motor_id  # For motor_on, use motor_id directly
        data = bytes([0xFF] * 7 + [0xFC])

        try:
            response = self._send_message_get_response(frame_id, data)
            motor_info = self.parse_recv_message(response, ignore_error=True)
            
            if int(motor_info.error_code, 16) != MotorErrorCode.normal:
                while int(motor_info.error_code, 16) != MotorErrorCode.normal:
                    info(f"Motor {self.motor_id} error: {motor_info.error_message}")
                    self.clear_error()
                    time.sleep(0.01)
                    info(f"Motor {self.motor_id} error cleaned")
                    
                    response = self._send_message_get_response(frame_id, data)
                    motor_info = self.parse_recv_message(response, ignore_error=True)
            else:
                info(f"Motor {self.motor_id} is already on")
            
            logging.getLogger().setLevel(current_level)
            motor_info = self.parse_recv_message(response)
            return motor_info
        except Exception as e:
            logging.getLogger().setLevel(current_level)
            raise

    def clear_error(self) -> None:
        """Clear motor errors."""
        frame_id = self.motor_id
        data = bytes([0xFF] * 7 + [0xFB])
        info("Clearing error")
        
        for _ in range(3):
            try:
                self.comms.cansend(
                    self.iface,
                    extended=False,
                    can_id=frame_id,
                    data=data,
                )
            except Exception as e:
                warning(f"Error clearing motor error: {e}")

    def motor_off(self) -> None:
        """Turn off the motor."""
        frame_id = self._get_frame_id()
        data = bytes([0xFF] * 7 + [0xFD])
        self._send_message_get_response(frame_id, data)

    def save_zero_position(self) -> None:
        """Save the current position as zero position."""
        frame_id = self._get_frame_id()
        data = bytes([0xFF] * 7 + [0xFE])
        
        try:
            self._send_message_get_response(frame_id, data, timeout=2.0)
        except AssertionError:
            pass
        
        # Check if set zero position success
        current_state = self.set_control(pos=0, vel=0, kp=0, kd=0, torque=0)
        diff = abs(current_state.position)
        if diff < 0.01:
            info(
                f"Motor {self.motor_id} set zero position success, current position: {current_state.position}"
            )

    def set_control(
        self,
        pos: float = 0.0,
        vel: float = 0.0,
        kp: float = 0.0,
        kd: float = 0.0,
        torque: float = 0.0,
    ) -> FeedbackFrameInfo:
        """Set the control of the motor and return its status.

        Args:
            pos: The target position value.
            vel: The target velocity value.
            kp: The proportional gain value.
            kd: The derivative gain value.
            torque: The target torque value.

        Returns:
            FeedbackFrameInfo: The current state of the motor.
        """
        frame_id = self._get_frame_id()

        # Prepare the CAN message
        data = bytearray(8)
        if self.control_mode == ControlMode.MIT:
            pos_tmp = float_to_uint(pos, self.motor_constants.POSITION_MIN, self.motor_constants.POSITION_MAX, 16)
            vel_tmp = float_to_uint(vel, self.motor_constants.VELOCITY_MIN, self.motor_constants.VELOCITY_MAX, 12)
            kp_tmp = float_to_uint(kp, self.motor_constants.KP_MIN, self.motor_constants.KP_MAX, 12)
            kd_tmp = float_to_uint(kd, self.motor_constants.KD_MIN, self.motor_constants.KD_MAX, 12)
            tor_tmp = float_to_uint(torque, self.motor_constants.TORQUE_MIN, self.motor_constants.TORQUE_MAX, 12)
            data[0] = (pos_tmp >> 8) & 0xFF
            data[1] = pos_tmp & 0xFF
            data[2] = (vel_tmp >> 4) & 0xFF
            data[3] = ((vel_tmp & 0xF) << 4) | (kp_tmp >> 8)
            data[4] = kp_tmp & 0xFF
            data[5] = (kd_tmp >> 4) & 0xFF
            data[6] = ((kd_tmp & 0xF) << 4) | (tor_tmp >> 8)
            data[7] = tor_tmp & 0xFF

        # Send the CAN message and get response
        response = self._send_message_get_response(frame_id, bytes(data), max_retry=15)

        # Parse the received message
        motor_info = self.parse_recv_message(response)
        return motor_info

    def get_state(self) -> FeedbackFrameInfo:
        """Get the current motor state without releasing position hold.
        
        This method uses motor_on() to get feedback from the motor. Unlike calling
        set_control() with zeros, this approach doesn't release any active position
        hold that may be in effect. The motor should already be enabled, so this
        primarily serves to retrieve the current state (position, velocity, torque, 
        temperatures, etc.) without interfering with ongoing control commands.
        
        Returns:
            FeedbackFrameInfo: The current state of the motor.
        """
        # Use motor_on() to get feedback - this doesn't release position hold
        # like set_control() with zeros would
        return self.motor_on()

    def save_calibration_state(self, lower_limit: float, upper_limit: float, calibration_homing_pos: float, total_range: float):
        """Set the calibration data for this motor.
        
        Args:
            lower_limit: The encoder value at the lower physical limit
            upper_limit: The encoder value at the upper physical limit
            calibration_homing_pos: The encoder value at the homing position after calibration.
                This may have a small residual offset from 0 due to motor drift after zeroing.
            total_range: The total range of motion in radians
        """
        self._disable_compliance_mode = False
        self.calibration_time = currtime()
        self.lower_limit = lower_limit
        self.upper_limit = upper_limit
        self.calibration_homing_pos = calibration_homing_pos
        self.total_range = total_range

        if not self.dual_encoder:
            warning(f"Motor {self.motor_name} does not have a dual encoder, skipping persistence of calibration data to disk.")
            return

        data = {
            "calibration_time": self.calibration_time,
            "lower_limit": self.lower_limit,
            "upper_limit": self.upper_limit,
            "calibration_homing_pos": self.calibration_homing_pos,
            "total_range": self.total_range,
        }
        write_json_file(data, MOTORS_CALIBRATION_PATH + "/" + self.motor_name + ".json")

    def load_calibration_state(self) -> bool:
        """Load the calibration state from the file.
        
        Returns True if the calibration state was loaded successfully, False otherwise.
        Returns False if the file doesn't exist or uses the old calibration format
        (middle_pos instead of calibration_homing_pos), requiring recalibration.
        """
        if not self.dual_encoder:
            return False

        file_path = MOTORS_CALIBRATION_PATH + "/" + self.motor_name + ".json"
        if not os.path.exists(file_path):
            return False
        data = read_json_file(file_path)
        # Check for new calibration format - return False if using old format
        if "calibration_homing_pos" not in data:
            return False
        self.calibration_time = data["calibration_time"]
        self.lower_limit = data["lower_limit"]
        self.upper_limit = data["upper_limit"]
        self.calibration_homing_pos = data["calibration_homing_pos"]
        self.total_range = data["total_range"]
        return True

    def check_in_range(self) -> bool:
        """Check if the motor is in range."""
        cur_pos = self.get_state().position
        if cur_pos is not None and self.lower_limit is not None and self.upper_limit is not None:
            if cur_pos >= self.lower_limit and cur_pos <= self.upper_limit:
                return True
        return False