"""Virtual Robstride motor simulator."""
from feathersdk.robot.motors.robstride import RobstrideRunMode, RobstrideMotorMsg as MotorMsg, RobstrideMotorMode, \
    RobstrideMotorError, RobstrideVersionType
from feathersdk.robot.motors.robstride.params import RobstrideParamList, RobstrideFeedbackMessageBounds as FB_BOUNDS, \
    RobstrideOperationCommand as OP_CMD, get_param_info
from feathersdk.robot.motors.robstride.params import get_robstride_version
from feathersdk.robot.motors.robstride.message_handler import parse_robstride_can_id
from feathersdk.comms import CommsManager, SocketResult
from feathersdk.comms.system import is_can_interface, S_uint16be, S_uint16le
from dataclasses import dataclass
import math
from feathersdk.utils.common import try_log_exception_peacefully, currtime, timediff
from feathersdk.utils.feathertypes import *


DEFAULT_MOTOR_ARMATURE = 0.01  # kg⋅m²
MIN_DAMPING_FOR_STABILITY = 0.1  # Minimum damping to prevent instability
POSITION_ERROR_LOG_THRESHOLD = 10.0  # Log when error changes by this percentage


@dataclass
class VirtualRobstrideOperationState:
    """A dictionary of current operation mode values."""
    target_torque: float
    target_angle: float
    target_velocity: float
    kp: float
    kd: float

    def update(self, target_torque: float, target_angle: float, target_velocity: float, kp: float, kd: float) -> None:
        """Update the operation state with the given values."""
        self.target_torque = target_torque
        self.target_angle = target_angle
        self.target_velocity = target_velocity
        self.kp = kp
        self.kd = kd
    
    @staticmethod
    def default() -> "VirtualRobstrideOperationState":
        return VirtualRobstrideOperationState(
            target_torque=0.0,
            target_angle=0.0,
            target_velocity=0.0,
            kp=0.0,
            kd=0.0,
        )
    
    @staticmethod
    def from_op_command(op_cmd: OP_CMD) -> "VirtualRobstrideOperationState":
        """Create a VirtualRobstrideOperationState from a RobstrideOperationCommand."""
        return VirtualRobstrideOperationState(
            target_torque=op_cmd.target_torque,
            target_angle=op_cmd.target_angle,
            target_velocity=op_cmd.target_velocity,
            kp=op_cmd.kp,
            kd=op_cmd.kd,
        )


@dataclass
class VirtualRobstrideState:
    """A dictionary of current motor state, EG: current angle, velocity, torque, temperature, errors, etc."""
    enabled: bool
    angle: float
    velocity: float
    torque: float
    temp: float
    active_reporting: bool
    last_active_reporting_time: float
    motor_mode: RobstrideMotorMode
    errors: Set[RobstrideMotorError]
    params: RobstrideParamList

    def get_feedback_extra_byte(self) -> int:
        """Build the extra byte for the feedback message. First 2 bit for motor mode, then 6 bit for errors."""
        err_to_bit = {v: k for k, v in RobstrideMotorError._motor_error_bits().items()}
        return self.motor_mode.value << 6 | sum(err_to_bit.get(e, 0) for e in self.errors)
    
    def get_fault_feedback_bytes(self) -> bytes:
        """Get the fault feedback bytes for the fault feedback message.
        
        Note: I'm not exactly sure this is correct endianness since it's hard to trigger a fault manually, but based on
        some previous fault feedback messages we got once, it seems to be little endian.
        """
        err_to_bit = {v: k for k, v in RobstrideMotorError._fault_error_bits().items()}
        return S_uint16le.pack(sum(err_to_bit.get(e, 0) for e in self.errors)) + bytes([0])
    
    @staticmethod
    def default() -> "VirtualRobstrideState":
        return VirtualRobstrideState(
            enabled=False,
            angle=0.0,
            velocity=0.0,
            torque=0.0,
            temp=30.0,
            active_reporting=False,
            last_active_reporting_time=currtime(),
            motor_mode=RobstrideMotorMode.Reset,
            errors=set(),
            params=RobstrideParamList(),
        )
    
    def update(self, **kwargs) -> None:
        """Update the state with the given values."""
        for key, value in kwargs.items():
            if key not in self.__dataclass_fields__:
                raise ValueError(f"Invalid key: {key}")
            setattr(self, key, value)


def default_firmware_version(version: RobstrideVersionType) -> str:
    """Get the default firmware version a virtual motor should have for the given version."""
    version = get_robstride_version(version)
    return f"0.{version.value}.0.0"


def default_broadcast_id(version: RobstrideVersionType, motor_id: int) -> bytes:
    """Get the default broadcast ID for a virtual motor."""
    version = get_robstride_version(version)
    return bytes([0xDE, 0xAD, 0xBE, 0xEF, 0x00, version.value, 0x00, motor_id])


class VirtualRobstride:
    """Virtual Robstride motor simulator.
    
    Args:
        motor_id: The motor ID to simulate. Should be in range [0x00, 0xFF].
        version: The version of the Robstride motor to simulate.
        iface: The interface to use for the communication.

    Attributes:
        motor_id: The unique ID of this motor instance.
        version: The version of Robstride motor this instance is simulating.
        iface: The can interface to use for the communication.
        comms: The comms manager to use for the communication.
        params: A dictionary of current motor parameters and their values. Defaults to the param default values on
            initialization, with 'None' defaults overridden to a guessed default value.
        last_host_id: The last host ID we received a message from.
        state: A dictionary of current motor state, EG: current angle, velocity, torque, temperature, errors, etc.
    """
    def __init__(self, motor_id: int, version: RobstrideVersionType, iface: str):
        self.motor_id = motor_id
        self.version = get_robstride_version(version)
        self.firmware_version = default_firmware_version(self.version)
        self.iface = iface

        self.comms = CommsManager()
        self.comms.add_callback(self.handle_message)
        if not is_can_interface(iface):
            raise ValueError(f"Invalid interface: {iface}")
        self.comms.add_endpoint(iface, skip_duplicates=True)

        self.operation_vals = VirtualRobstrideOperationState.default()
        self.last_host_id = 0xAA

        self.state = VirtualRobstrideState.default()

        self._sent_first_broadcast_frame = False
        
        # Motor dynamics parameters
        self.motor_inertia = DEFAULT_MOTOR_ARMATURE  # kg⋅m² - motor inertia (affects acceleration)
        self.friction_coefficient = 0.5  # N⋅m⋅s/rad - viscous friction for damping
        
        # Logging control
        self._log_feedback_count = 0  # Counter to limit feedback logging
        
        # Position control error tracking
        self._last_error_percentage = None  # Track last error percentage for change detection
    
    def _vmsg_can_id(self, op: MotorMsg, motor_id: int, host_id: Optional[int] = None, e_byte: int = 0) -> int:
        """e_byte is the extra_byte value for the message."""
        host_id = self.last_host_id if host_id is None else host_id
        return (op.value << 24) | (e_byte << 16) | (motor_id << 8) | host_id

    def send_motor_feedback(self, msg_op: Optional[MotorMsg] = MotorMsg.MotorFeedback, send_ver: bool = False) -> None:
        """Send a motor feedback message. If `send_ver` is True, will send the firmware version number instead"""
        data = S_uint16be.map_and_pack(self.state.angle, *FB_BOUNDS.ANGLE_MAP[self.version]) + \
            S_uint16be.map_and_pack(self.state.velocity, *FB_BOUNDS.VELOCITY_MAP[self.version]) + \
            S_uint16be.map_and_pack(self.state.torque, *FB_BOUNDS.TORQUE_MAP[self.version]) + \
            S_uint16be.pack(int(self.state.temp * FB_BOUNDS.TEMP_DIVISOR))
        
        if send_ver:
            data = bytes([0, 0xC4, ord('V')] + list(map(int, self.firmware_version.split('.'))) + [7])
        
        can_id = self._vmsg_can_id(msg_op, self.motor_id, e_byte=self.state.get_feedback_extra_byte())
        self.comms.cansend(self.iface, True, can_id, data)
    
    def send_broadcast_frame(self) -> None:
        """Send a broadcast frame. Done after things like GetDeviceId and SetID."""
        can_id = self._vmsg_can_id(MotorMsg.GetDeviceId, self.motor_id, 0xFE)
        self.comms.cansend(self.iface, True, can_id, data=default_broadcast_id(self.version, self.motor_id))
    
    @try_log_exception_peacefully(capture_locals={'result', 'msg_op', 'motor_id', 'host_id', 'e_byte', 'param_info'})
    def handle_message(self, result: SocketResult):
        """Handle a message from the Robstride motor."""
        # Make sure this is a can message intended for this motor
        if not result.is_can() or not result.is_extended_can():
            return
    
        msg_op, motor_id, host_id, e_byte = parse_robstride_can_id(result.can_id, motor_to_host=False)
        if motor_id != self.motor_id:
            return

        self.last_host_id = host_id

        # Handle the given operation
        if msg_op == MotorMsg.GetDeviceId:
            self.send_broadcast_frame()

        elif msg_op == MotorMsg.OperationControl:
            # Note: these values can never be out of range because we use the entire 65535 range for each value
            self.operation_vals.update(
                target_torque=S_uint16be.unpack_and_map([e_byte, host_id], *FB_BOUNDS.TORQUE_MAP[self.version]),
                target_angle=S_uint16be.unpack_and_map(result.data[:2], *FB_BOUNDS.ANGLE_MAP[self.version]),
                target_velocity=S_uint16be.unpack_and_map(result.data[2:4], *FB_BOUNDS.VELOCITY_MAP[self.version]),
                kp=S_uint16be.unpack_and_map(result.data[4:6], *OP_CMD.KP_BOUNDS_MAP[self.version]),
                kd=S_uint16be.unpack_and_map(result.data[6:8], *OP_CMD.KD_BOUNDS_MAP[self.version]),
            )

        elif msg_op == MotorMsg.Enable:
            self.state.enabled = True
            self.state.motor_mode = RobstrideMotorMode.Run
            self.send_motor_feedback()
        
        elif msg_op == MotorMsg.Disable:
            self.state.enabled = False
            self.state.motor_mode = RobstrideMotorMode.Reset

            # Check if we should send the firmware version number back or not
            self.send_motor_feedback(send_ver=result.data[:2] == bytes([0, 0xC4]))
        
        elif msg_op == MotorMsg.ZeroPos:
            if result.data != bytes([0x01, 0, 0, 0, 0, 0, 0, 0]):
                raise ValueError(f"Invalid zero position command bytes: {result.data}")
            self.state.angle = 0.0
            self.send_motor_feedback()
        
        elif msg_op == MotorMsg.SetID:
            self.motor_id = e_byte
            self.send_broadcast_frame()
        
        elif msg_op == MotorMsg.ReadParam:
            p_info = get_param_info(S_uint16le.unpack(result.data[:2]))
            can_id = self._vmsg_can_id(MotorMsg.ReadParam, self.motor_id, host_id, e_byte=0)
            data = S_uint16le.pack(p_info.id) + bytes([0, 0]) \
                + p_info.dtype.pack(self.state.params[p_info.name].value, pad_end=4)
            self.comms.cansend(self.iface, True, can_id, data)
        
        elif msg_op == MotorMsg.WriteParam:
            p_info = get_param_info(S_uint16le.unpack(result.data[:2]))
            if not p_info.writable:
                raise ValueError(f"Param not writable: {p_info.name}")
            
            self.state.params[p_info.name].value = p_info.unpack_and_clamp(result.data[4:], self.version)
            self.send_motor_feedback()
        
        elif msg_op == MotorMsg.SaveMotorData:
            if result.data != bytes([1, 2, 3, 4, 5, 6, 7, 8]):
                raise ValueError(f"Invalid save motor data command bytes: {result.data}")
            print("Don't know what exactly happens with save motor data")
            self.send_motor_feedback()
        
        elif msg_op == MotorMsg.SetMotorBaudRate:
            if result.data[:-2] != bytes([1, 2, 3, 4, 5, 6]) or result.data[-2] not in [1, 2, 3, 4]:
                raise ValueError(f"Invalid set motor baud rate command bytes: {result.data}")
            print("Will ignore set motor baud rate command")
            self.send_broadcast_frame()
        
        elif msg_op == MotorMsg.SetMotorActiveReporting:
            if result.data[:-2] != bytes([1, 2, 3, 4, 5, 6]) or result.data[-2] not in [0, 1]:
                raise ValueError(f"Invalid set motor active reporting command bytes: {result.data}")
            self.send_motor_feedback(msg_op=MotorMsg.SetMotorActiveReporting)
            self.state.active_reporting = result.data[-2] == 1
            self.state.last_active_reporting_time = currtime()

        elif msg_op == MotorMsg.SetMotorProtocol:
            if result.data[:-2] != bytes([1, 2, 3, 4, 5, 6]) or result.data[-2] not in [0, 1, 2]:
                raise ValueError(f"Invalid set motor protocol command bytes: {result.data}")
            print("Will ignore set motor protocol command")
            self.send_broadcast_frame()
                
        else:
            pass
    
    def step(self, dt: float) -> None:
        """Step the motor simulation."""
        if not self._sent_first_broadcast_frame:
            self.send_broadcast_frame()
            self._sent_first_broadcast_frame = True

        position_error, target_pos = None, None

        if self.state.enabled:
            if self.state.params.run_mode.value == RobstrideRunMode.Position:
                position_error, target_pos = self._simulate_position_mode(dt)
            elif self.state.params.run_mode.value == RobstrideRunMode.Speed:
                position_error, target_pos = self._simulate_speed_mode(dt)
            elif self.state.params.run_mode.value == RobstrideRunMode.Operation:
                position_error, target_pos = self._simulate_operation_mode(dt)
            else:
                raise NotImplementedError(f"Unimplemented run mode: {self.state.params.run_mode}")
        
        # Log position error if applicable
        if position_error is not None and target_pos is not None:
            self._log_position_error(position_error, target_pos)
        
        # Send active reporting message if applicable
        st = self.state.params.epscan_time.value
        at_interval_secs = 0 if st == 0 else (st + 1) * 0.005
        
        if self.state.active_reporting and timediff(self.state.last_active_reporting_time) > at_interval_secs:
            self.send_motor_feedback(msg_op=MotorMsg.SetMotorActiveReporting)
            self.state.last_active_reporting_time = currtime()
        
        # Check for and send fault feedback message if applicable
        self._check_send_fault_feedback()
    
    def _check_send_fault_feedback(self) -> None:
        """Check if we should send a fault feedback message."""
        if len(self.state.errors) == 0:
            return
        
        can_id = self._vmsg_can_id(MotorMsg.FaultFeedback, self.motor_id, self.last_host_id, e_byte=0)
        self.comms.cansend(self.iface, True, can_id, self.state.get_fault_feedback_bytes() + bytes([0, 0, 0, 0, 0]))

        self.state.errors.clear()

    def _update_dynamics(self, dt: float) -> None:
        """Update motor dynamics based on current torque."""
        acceleration = self.state.torque / self.motor_inertia
        self.state.velocity = FB_BOUNDS.VELOCITY_MAP[self.version].clamp(self.state.velocity + acceleration * dt)
        self.state.angle = FB_BOUNDS.ANGLE_MAP[self.version].clamp(self.state.angle + self.state.velocity * dt)

    def _simulate_position_mode(self, dt: float) -> Tuple[Optional[float], Optional[float]]:
        """Simulate position control mode. Returns (position_error, target_pos)."""
        position_error = self.state.params.loc_ref.value - self.state.angle
        
        # PD control: desired velocity based on position error
        desired_velocity = position_error * self.state.params.loc_kp.value
        velocity_error = desired_velocity - self.state.velocity
        
        # Control torque from velocity error (using Kd as velocity gain), and friction damping
        # TODO: Figure out what this `kd` val is supposed to be, because I don't see it in position mode params
        control_torque = velocity_error * max(self.operation_vals.kd, MIN_DAMPING_FOR_STABILITY)
        friction_torque = -self.friction_coefficient * self.state.velocity
        
        # Total torque
        self.state.torque = FB_BOUNDS.TORQUE_MAP[self.version].clamp(control_torque + friction_torque)
        self._update_dynamics(dt)
        
        return position_error, self.state.params.loc_ref.value

    def _simulate_speed_mode(self, dt: float) -> Tuple[Optional[float], Optional[float]]:
        """Simulate speed control mode. Returns (position_error, target_pos)."""
        self.state.velocity = self.state.params.spd_ref.value
        self.state.angle = FB_BOUNDS.ANGLE_MAP[self.version].clamp(self.state.angle + self.state.velocity * dt)
        return None, None

    def _simulate_operation_mode(self, dt: float) -> Tuple[Optional[float], Optional[float]]:
        """Simulate operation control mode. Returns (position_error, target_pos)."""
        position_error = self.operation_vals.target_angle - self.state.angle
        vel_error = self.operation_vals.target_velocity - self.state.velocity
        
        # PD control torque and inherent friction damping (even when Kd=0, this provides stability)
        control_torque = position_error * self.operation_vals.kp + vel_error * self.operation_vals.kd
        friction_torque = -self.friction_coefficient * self.state.velocity
        
        # Total torque: feedforward + control + friction
        new_torque = self.operation_vals.target_torque + control_torque + friction_torque
        self.state.torque = FB_BOUNDS.TORQUE_MAP[self.version].clamp(new_torque)
        self._update_dynamics(dt)
        
        return position_error, self.operation_vals.target_angle

    def _log_position_error(self, position_error: float, target_pos: float):
        """Log position error when it changes significantly."""
        position_range = 8.0 * math.pi
        error_abs = abs(position_error)
        error_percentage = (error_abs / position_range) * 100.0 if position_range > 0 else 0.0
        
        if self._last_error_percentage is not None:
            error_change = abs(error_percentage - self._last_error_percentage)
            if error_change >= POSITION_ERROR_LOG_THRESHOLD:
                print(f"[MOTOR {self.motor_id:02X}]: Position error: {error_percentage:.2f}% "
                      f"(error={position_error:.4f} rad, target={target_pos:.4f} rad, "
                      f"actual={self.state.angle:.4f} rad, change={error_change:.2f}%)")
                self._last_error_percentage = error_percentage
        else:
            # First time, just record the value
            self._last_error_percentage = error_percentage
