from feathersdk.comms.system import S_uint16le
from feathersdk.utils.common import make_exception_pickleable, currtime
from .params import RobstrideMotorMode, RobstrideMotorError, RobstrideMotorMsg, RobstrideRunMode, \
    RobstrideVersionType, get_param_info, RobstrideOperationCommand, get_robstride_version, _ROBSTRIDE_PARAMS
from .admin import RobstrideAdminControl
from ..motor import Motor
from feathersdk.comms import CommsManager, CanOverloadError, MotorUIDType
from feathersdk.utils.feathertypes import Any, Union, TypedDict, List, TYPE_CHECKING, Optional, TypeVar, Generic
from feathersdk.utils.files import write_json_file, read_json_file
from feathersdk.utils.constants import MOTORS_CALIBRATION_PATH
from dataclasses import dataclass
import math
import time
import os

if TYPE_CHECKING:  # Due to circular import
    from ..motors_manager import MotorsManager


T = TypeVar("T")

@dataclass(frozen=True, slots=True)
class TimestampedValue(Generic[T]):
    value: T
    timestamp: float

    def __iter__(self):
        return iter((self.value, self.timestamp))

    def __len__(self):
        return 2

    def __getitem__(self, i: int):
        return (self.value, self.timestamp)[i]


@make_exception_pickleable
class UnsafeCommandError(Exception):
    """Raised when attempting to execute an unsafe command."""
    def __init__(self, message: str, motor: 'RobstrideMotor'):
        super().__init__(message)
        self.motor = motor
        self.motor_name = motor.motor_name


@make_exception_pickleable
class MotorDisabledError(Exception):
    """Raised when attempting to command a motor that is disabled (in Reset mode)."""
    def __init__(self, message: str, motor: 'RobstrideMotor'):
        super().__init__(message)
        self.motor = motor
        self.motor_name = motor.motor_name


@make_exception_pickleable
class MotorModeInconsistentError(Exception):
    """Raised when attempting to command a motor with an incompatible run mode."""
    def __init__(self, message: str, motor: 'RobstrideMotor', curr_mode: RobstrideRunMode, exp_mode: RobstrideRunMode):
        super().__init__(message)
        self.motor = motor
        self.motor_name = motor.motor_name
        self.current_mode = curr_mode
        self.expected_mode = exp_mode


_MOTOR_NO_DEFAULT = object()


class RobstrideMotorConfig(TypedDict, total=False):
    version: RobstrideVersionType
    max_torque: float  # Max torque in N
    max_position_dx: float  # positional difference in radians
    joint_limits: tuple[float, float]  # joint limits in radians
    max_velocity: float  # Max velocity in radians / seconds


class RobstrideMotor(Motor):
    """Class to control a Robstride motor.
    
    Args:
        motor_id: The motor ID to use for the communication.
        name: The name of the motor.
        config: The configuration for the motor.
        iface: The interface to use for the communication.
    """
    def __init__(self, motor_id: int, name: str, config: RobstrideMotorConfig = {}, iface: Optional[str] = None):
        if "model" in config and "version" not in config:
            config["version"] = config["model"]
        if "version" not in config:
            raise KeyError("\"version\" key is required for RobstrideMotorConfig")
        self.version = get_robstride_version(config["version"])

        if motor_id < 0 or motor_id >= 0xFE:
            raise ValueError(f"Error: motor_id must be between [0, 0xFD], got {motor_id}")

        self.manager: Optional['MotorsManager'] = None

        self.comms = CommsManager()
        self.motor_id: int = motor_id
        self.default_host_id: int = 0xAA
        self.uid: MotorUIDType = MotorUIDType(self.motor_id)
        self.motor_name: str = name
        self.calibration_time = None
        self._homing_pos = config.get("homing_pos", 0)
        self.direction = config.get("direction", 1)

        self.range = config.get("range", {
            "min": -math.pi / 2,
            "max": math.pi / 2,
        })

        self.upper_limit = None
        self.lower_limit = None
        self.middle_pos = None
        self.total_range = None
        
        self.calibrated_angle = (0, -1)
        self.dual_encoder = True

        self.family_name: Union[str, None] = None
        self.iface: Union[str, None] = iface
        self.properties: dict[str, TimestampedValue] = {}
        self.joint_limits = config.get("joint_limits", None)

        if self.joint_limits is not None:
            if (
                self.joint_limits[0] < -math.pi
                or self.joint_limits[1] > math.pi
                and self.joint_limits[0] <= self.joint_limits[1]
            ):
                print("Warning: Joint limits must be from -pi to pi for ", name)
        self.max_torque = config.get("max_torque", None)
        self.max_position_dx = config.get("max_position_dx", None)
        self.max_velocity = config.get("max_velocity", None)
        self.run_mode = RobstrideRunMode.Operation
        self.compliance_mode: bool = False
        self.target_position = 0
        self.last_compliance_pos = 0
        self.compliance_mode_torque_threshold = config.get("compliance_mode_torque_threshold", 1.0)
        self.compliance_mode_dx = config.get("compliance_mode_dx", 0.01)

        self.mode: TimestampedValue[RobstrideMotorMode] = TimestampedValue(RobstrideMotorMode.Reset, -1)
        self.angle: TimestampedValue[float] = TimestampedValue(0, -1)
        self.calibrated_angle: TimestampedValue[float] = TimestampedValue(0, -1)
        self.velocity: TimestampedValue[float] = TimestampedValue(0, -1)
        self.torque: TimestampedValue[float] = TimestampedValue(0, -1)
        self.temp: TimestampedValue[float] = TimestampedValue(0, -1)
        self.errors: TimestampedValue[List[RobstrideMotorError]] = TimestampedValue([], -1)

        self.load_calibration_state()

        self._admin: RobstrideAdminControl = RobstrideAdminControl(self)

    @property
    def expected_range(self) -> float:
        return self.range["max"] - self.range["min"]

    # make a getter for homing_pos
    @property
    def homing_pos(self) -> float:
        if self.middle_pos is not None:
            return self.middle_pos + self._homing_pos
        return self._homing_pos

    def update_feedback(
        self, angle: float, velocity: float, torque: float, temp: float, errors: List[RobstrideMotorError], mode: RobstrideMotorMode
    ):
        last_update = currtime()
        # Raw angle coming from the encoder.
        self.angle = TimestampedValue(angle, last_update)
        if self.middle_pos is not None:
            # Calibrated angle useful for control the robot between different machines to account for encoder diffs.
            # Calibrated angle is a better indicator of the angle set by the user.
            self.calibrated_angle = TimestampedValue(angle - self.middle_pos, last_update)
        else:
            self.calibrated_angle = TimestampedValue(angle, last_update)
        self.velocity = TimestampedValue(velocity, last_update)
        self.torque = TimestampedValue(torque, last_update)
        self.temp = TimestampedValue(temp, last_update)
        self.errors = TimestampedValue(errors, last_update)
        self.mode = TimestampedValue(mode, last_update)

    def update_property(self, property_name: str, value: Union[float, int]):
        self.properties[property_name] = TimestampedValue(value, currtime())
    
    def get_property_value(self, property_name: str, default: Any = _MOTOR_NO_DEFAULT) -> Any:
        """Get the currently stored value for a property in this motor. 
        
        If the property is not found, return the default value. If the default value is not provided, raise a KeyError.
        """
        if property_name in self.properties:
            return self.properties[property_name].value
        if default is not _MOTOR_NO_DEFAULT:
            return default
        raise KeyError(f"Property {property_name} not found in motor {self.motor_name}")

    def get_property_timestamp(self, property_name: str, default: Any = _MOTOR_NO_DEFAULT) -> Any:
        """Get the timestamp of the currently stored value for a property in this motor. 
        
        If the property is not found, return the default value. If the default value is not provided, raise a KeyError.
        """
        if property_name in self.properties:
            return self.properties[property_name].timestamp
        if default is not _MOTOR_NO_DEFAULT:
            return default
        raise KeyError(f"Property {property_name} not found in motor {self.motor_name}")

    def is_safe_position_update(self, target_position: float) -> bool:
        if self.max_position_dx is not None:
            return abs(target_position - self.angle[0]) <= self.max_position_dx
        return True

    def update_target_position(self, target_position: float) -> None:
        """
        Used for safety purposes.
        If max_position_dx is set, the motor will be stopped if the position is outside the max_position_dx.
        """
        if self.compliance_mode:
            raise UnsafeCommandError(message="Commands currently blocked - in compliance mode", motor=self)

        self.target_position = target_position

    def should_trigger_compliance_mode(self):
        if self.compliance_mode or self.mode[0] == RobstrideMotorMode.Reset:
            return False

        if self.max_position_dx is not None:
            if abs(self.angle[0] - self.target_position) > self.max_position_dx:
                print(
                    f"Compliance mode triggered: Motor {self.motor_name} Position {self.angle[0]} is outside max position dx {self.max_position_dx} Target position {self.target_position}"
                )
                return True

        if self.max_torque is not None:
            if abs(self.torque[0]) > self.max_torque:
                print(
                    f"Compliance mode triggered: Motor {self.motor_name} Torque {self.torque[0]} is greater than max torque {self.max_torque}"
                )
                return True

        if self.joint_limits is not None:
            if self.angle[0] > self.joint_limits[1] or self.angle[0] < self.joint_limits[0]:
                print(
                    f"Compliance mode triggered: Motor {self.motor_name} Angle {self.angle[0]} is outside joint limits {self.joint_limits}"
                )
                return True

        if self.max_velocity is not None:
            if abs(self.velocity[0]) > self.max_velocity:
                print(
                    f"Compliance mode triggered: Motor {self.motor_name} Velocity {self.velocity[0]} is greater than max velocity {self.max_velocity}"
                )
                return True

        return False

    def set_calibration(self, lower_limit: float, upper_limit: float, middle_pos: float, total_range: float):
        self.calibration_time = currtime()
        self.lower_limit = lower_limit
        self.upper_limit = upper_limit
        self.middle_pos = middle_pos
        self.total_range = total_range
        self.save_calibration_state()

    def save_calibration_state(self):
        data = {
            "calibration_time": self.calibration_time,
            "lower_limit": self.lower_limit,
            "upper_limit": self.upper_limit,
            "middle_pos": self.middle_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.
        """
        file_path = MOTORS_CALIBRATION_PATH + "/" + self.motor_name + ".json"
        if not os.path.exists(file_path):
            return False
        data = read_json_file(file_path)
        self.calibration_time = data["calibration_time"]
        self.lower_limit = data["lower_limit"]
        self.upper_limit = data["upper_limit"]
        self.middle_pos = data["middle_pos"]
        self.total_range = data["total_range"]
        return True

    def _msg_can_id(self, op: RobstrideMotorMsg, motor_id: int, host_id: Union[int, bytes, None] = None, e_byte: int = 0) -> int:
        """Build a CAN ID for a Robstride motor message.
        
        Args:
            op: The operation to build the CAN ID for.
            motor_id: The motor ID to build the CAN ID for.
            host_id: The host ID to build the CAN ID for. If int, then data will be checked to be within valid bounds.
                If you wish to skip this check, pass it as a bytes object with one element.
            e_byte: The extra byte to build the CAN ID for.
        """
        if isinstance(host_id, bytes):
            assert len(host_id) == 1, f"Error: host_id must be 1 byte, got {len(host_id)}"
            host_id = host_id[0]
        else:
            host_id = self.default_host_id if host_id is None else host_id
            if host_id < 0 or host_id > 0xFD:
                raise ValueError(f"Error: host_id must be between [0, 0xFD], got {host_id}")
        
        if motor_id < 0 or motor_id > 0xFD:
            raise ValueError(f"Error: motor_id must be between [0, 0xFD], got {motor_id}")
        if e_byte < 0 or e_byte > 0xFF:
            raise ValueError(f"Error: e_byte must be between [0, 0xFF], got {e_byte}")
        
        return (op.value << 24) | (e_byte << 16) | (host_id << 8) | motor_id
    
    def get_device_id(self) -> int:
        """Get the device id of the motor."""
        can_id = self._msg_can_id(RobstrideMotorMsg.GetDeviceId, self.motor_id)
        self.comms.cansend(self.iface, True, can_id, bytes([0] * 8), self.uid)

    def enable(self):
        """Send the 'enable' command to the motor."""
        can_id = self._msg_can_id(RobstrideMotorMsg.Enable, self.motor_id)
        self.comms.cansend(self.iface, True, can_id, bytes([0] * 8), self.uid)

    def disable(self, clear_faults: bool = False):
        """Send the 'disable' command to the motor.
        
        Args:
            clear_faults: If True, clear the fault bits in the disable command.
        """
        can_id = self._msg_can_id(RobstrideMotorMsg.Disable, self.motor_id)
        self.comms.cansend(self.iface, True, can_id, bytes([1] + [0] * 7) if clear_faults else bytes([0] * 8), self.uid)
    
    def set_active_reporting(self, active: bool, scan_time: int = 1):
        """Turn on/off active reporting for the motor.
        
        Args:
            active: If True, turn on active reporting. False to disable.
            scan_time: Decides the time between feedback messages. 1 = 10ms, +1 = +5ms
        """
        self.write_param(_ROBSTRIDE_PARAMS.epscan_time, scan_time)

        can_id = self._msg_can_id(RobstrideMotorMsg.SetMotorActiveReporting, self.motor_id)
        data = bytes([1, 2, 3, 4, 5, 6, 1 if active else 0, 8])
        self.comms.cansend(self.iface, True, can_id, data, self.uid)
    
    def zero_position(self):
        """Set the zero position for the motor."""
        can_id = self._msg_can_id(RobstrideMotorMsg.ZeroPos, self.motor_id)
        self.comms.cansend(self.iface, True, can_id, bytes([0x01] + [0] * 7), self.uid)

    def read_param(self, param_id: Union[int, str]) -> float:
        """Read a parameter from the motor."""
        can_id = self._msg_can_id(RobstrideMotorMsg.ReadParam, self.motor_id)
        data = S_uint16le.pack(get_param_info(param_id).id) + bytes([0] * 6)
        self.comms.cansend(self.iface, True, can_id, data, self.uid)

    def write_param(self, p_id: Union[int, str], value: Union[float, int]) -> None:
        self._write_param_checked(p_id, value, force_run=False)

    def _write_param_checked(self, p_id: Union[int, str], value: Union[float, int], force_run: bool = False) -> None:
        """Write a parameter value to the specified motor.

        Args:
            p_id: The parameter ID to write.
            value: The value to write.
            force_run: If True, force the write to be executed, ignoring the safety checks.
        """
        param = get_param_info(p_id)
        if not force_run and not param.is_within_range(value, self.version):
            raise ValueError(f"Value {value} is out of range for param {param.name}: {param.range[self.version]}")
            
        data = S_uint16le.pack(param.id) + bytes([0] * 2) + param.dtype.pack(value, pad_end=4, pad_val=0)
        
        if not force_run and param.name == "loc_ref":
            if self.is_safe_position_update(value):
                self.update_target_position(value)  # can raise an exception to cancel the run.
            else:
                self.trigger_compliance_mode()
                raise UnsafeCommandError(
                    message=f"Safety Error: Target position {value} is too far from current position {self.angle[0]} (max dx: {self.max_position_dx})",
                    motor=self
                ) # stop the program safely.

        can_id = self._msg_can_id(RobstrideMotorMsg.WriteParam, self.motor_id)

        try:
            self.comms.cansend(self.iface, True, can_id, data, self.uid)
        except CanOverloadError:
            time.sleep(0.01)
            self.comms.cansend(self.iface, True, can_id, data, self.uid)
        
    def set_run_mode(self, run_mode: RobstrideRunMode):
        self.write_param("run_mode", run_mode.value)
        self.run_mode = run_mode

    def set_target_position(self, target_position: float) -> None:
        """Set the target position (loc_ref) for the specified motor."""
        # Check if motor is disabled (in Reset mode)
        if self.mode[0] == RobstrideMotorMode.Reset:
            raise MotorDisabledError(
                message=f"Cannot set target position for motor {self.motor_name}: "
                        f"Motor is disabled (in Reset mode). Enable the motor first.",
                motor=self
            )
        
        # Check if motor is in the correct run mode for position control
        if self.run_mode != RobstrideRunMode.Position:
            raise MotorModeInconsistentError(
                f"Cannot set target position for motor {self.motor_name}: "
                f"Motor is in {self.run_mode.name} mode, but Position mode is required. "
                f"Use set_run_mode() to switch to Position mode first.",
                motor=self,
                curr_mode=self.run_mode,
                exp_mode=RobstrideRunMode.Position
            )
        
        self.write_param("loc_ref", target_position)
    
    def set_target_velocity(self, target_velocity: float) -> None:
        """Set the target velocity (spd_ref) for the specified motor."""
        # Check if motor is disabled (in Reset mode)
        if self.mode[0] == RobstrideMotorMode.Reset:
            raise MotorDisabledError(
                message=f"Cannot set target velocity for motor {self.motor_name}: "
                        f"Motor is disabled (in Reset mode). Enable the motor first.",
                motor=self
            )
        
        # Check if motor is in the correct run mode for velocity control
        if self.run_mode != RobstrideRunMode.Speed:
            raise MotorModeInconsistentError(
                f"Cannot set target velocity for motor {self.motor_name}: "
                f"Motor is in {self.run_mode.name} mode, but Speed mode is required. "
                f"Use set_run_mode() to switch to Speed mode first.",
                motor=self,
                curr_mode=self.run_mode,
                exp_mode=RobstrideRunMode.Speed
            )
        
        self.write_param("spd_ref", target_velocity)
    
    def operation_command(self, op: RobstrideOperationCommand, apply_target_position: bool = True):
        """ target_torque range (-60Nm to 60Nm)"""
        can_id, data = op.build_command()

        if apply_target_position:
            if self.is_safe_position_update(op.target_angle):
                self.update_target_position(op.target_angle)  # can raise an exception to cancel the run.
            else:
                self.trigger_compliance_mode()
                raise UnsafeCommandError(
                    message=f"Safety Error: Target position {op.target_angle} is too far from current position {self.angle[0]} (max dx: {self.max_position_dx})",
                    motor=self
                ) # stop the program safely.

        self.comms.cansend(self.iface, True, can_id, data, self.uid)
    
    def trigger_compliance_mode(self):
        self.manager.trigger_compliance_mode(self.family_name)


MotorMap = dict[str, RobstrideMotor]

