import os
import time
import enum
import sys
from feathersdk.utils.math_utils import degrees_to_radians
from feathersdk.comms import CommsManager, SocketResult
from feathersdk.comms.system import get_all_physical_can_interfaces
from .robstride import RobstrideMotor, MotorMap, RobstrideRunMode, RobstrideMotorMode, RobstrideOperationCommand, \
    UnsafeCommandError, handle_robstride_motor_message
from feathersdk.utils import progressbar
from feathersdk.utils.logger import debug, info, warning, error
import math
from feathersdk.utils.feathertypes import *
import asyncio
from feathersdk.utils import constants
from feathersdk.utils.common import make_exception_pickleable, timediff, currtime, dont_allow_on_fork, FeatherThread
import numpy as np
from feathersdk.robot.instance_lock import check_still_own_lock

DRY_RUN = True
COMPLIANCE_MODE_KP = 30
COMPLIANCE_MODE_KD = 5

class CalibrationState(enum.Enum):
    """Represents the calibration and health state of a motor system."""
    UNINITIALIZED = "uninitialized"      # Before any checks are performed
    UNCALIBRATED = "uncalibrated"        # No calibration data exists
    OUT_OF_RANGE = "out_of_range"        # Has calibration but motors are outside valid range
    HEALTHY = "healthy"                  # Calibrated and in range, ready to use
    RECALIBRATING = "recalibrating"      # Currently in the process of recalibration
    CALIBRATION_FAILED = "calibration_failed"  # Recalibration attempt failed


class CalibrationWarnings(enum.Flag):
    RANGE_MISMATCH = 1
    HOMING_POSITION_NONZERO = 2
    LOWER_LIMIT_MISMATCH = 4
    UPPER_LIMIT_MISMATCH = 8
    NONE = 0

    def __str__(self):
        if self == CalibrationWarnings.NONE:
            return "None"
        warnings = []
        if self & CalibrationWarnings.RANGE_MISMATCH:
            warnings.append("Range mismatch: The expected range of the motor is different from the total range.")
        if self & CalibrationWarnings.HOMING_POSITION_NONZERO:
            warnings.append("Calibration Homing position is expected to be close to zero after calibration.")
        if self & CalibrationWarnings.LOWER_LIMIT_MISMATCH:
            warnings.append("Lower limit does not match the min range of the motor.")
        if self & CalibrationWarnings.UPPER_LIMIT_MISMATCH:
            warnings.append("Upper limit does not match the max range of the motor.")
        if len(warnings) > 0:
            warnings.append("If the warnings show up consistently after multiple calibrations, review the motor config and mechanical constraints..")

        return ", ".join(warnings)

@make_exception_pickleable
class MotorNotFoundError(Exception):
    """Raised when attempting to access a motor that does not exist in the motors map."""
    def __init__(self, message: str, motor_name_or_id: Union[str, int]):
        super().__init__(message)
        self.motor_name_or_id = motor_name_or_id
        # For consistency with other exceptions, also provide motor_name (as string representation)
        self.motor_name = str(motor_name_or_id)


@make_exception_pickleable
class MotorCalibrationError(Exception):
    """Raised when motor calibration fails or motor doesn't respond as expected during calibration."""
    def __init__(self, message: str, motor: 'RobstrideMotor'):
        super().__init__(message)
        self.motor = motor
        self.motor_name = motor.motor_name


class SafeStopPolicy(enum.Enum):
    DISABLE_MOTORS = 1
    COMPLIANCE_MODE = 2


class FamilyConfig(TypedDict, total=False):
    safe_stop_policy: SafeStopPolicy


class MotorsManager:
    """
    Handles sending commands to and receiving feedback from motors.
    Handles caching the state of the motors.
    Abstractions to make it easy to work with motors.
    """

    def __init__(self) -> None:
        self.motors: MotorMap = {}
        self.motors_by_id: Dict[int, RobstrideMotor] = {}
        self.motor_families: Dict[str, MotorMap] = {}
        self.motor_families_config: Dict[str, FamilyConfig] = {}
        self.motor_families_compliance_threads: Dict[str, FeatherThread] = {}
        self.aborted = False
        self.comms = CommsManager()
        self.comms.add_callback(lambda result: self.on_motor_message(result))
        self._allow_trigger_compliance_mode: Dict[str, bool] = {}
        self.is_forked = False

    @dont_allow_on_fork
    def on_abort(self) -> None:
        """Stop compliance mode for all motor families."""
        for family_name in self.motor_families_compliance_threads.keys():
            self.stop_compliance_mode(family_name)
        self.aborted = True

    @dont_allow_on_fork
    def add_motors(self, motors_map: MotorMap, family_name: str, family_config: FamilyConfig = {}) -> None:
        """Add a group of motors to the manager as a motor family.

        Args:
            motors_map: Dictionary mapping motor names to RobstrideMotor instances
            family_name: Name identifier for this motor family
            family_config: Configuration options for the motor family
        """
        self.motor_families[family_name] = motors_map
        self.motor_families_config[family_name] = family_config
        for motor_name, motor in motors_map.items():
            self.motors_by_id[motor.motor_id] = motor
            self.motors[motor_name] = motor
            motor.family_name = family_name
            motor.manager = self
    
    @dont_allow_on_fork
    def find_motors(self, names: Optional[List[str]] = None, allow_not_found: bool = False, progress: bool = False) -> None:
        """Find the motors on the physical CAN interfaces.

        Args:
            names (Optional[List[str]]): The names of the motors to find. If not provided, find all motors in the motors map.
            allow_not_found (bool): If True, do not raise an error if a motor is not found.
            progress (bool): If True (and tqdm is installed), show a progress bar.
        """
        if names is None:
            names = list(self.motors.keys())

        for motor_name in progressbar(names, progress=progress):
            motor = self.motors[motor_name]
            if motor.iface is not None:
                continue

            found = False
            for attempt in range(3):
                for can_interface in get_all_physical_can_interfaces():
                    info(f"Testing motor {motor_name} on interface: {can_interface}")
                    try:
                        motor.iface = can_interface
                        motor.update_property("loc_ref", None)
                        self.read_param(motor_name, "loc_ref")
                        time.sleep(0.01)

                        if motor.properties["loc_ref"].value is not None:
                            found = True
                            break
                    except Exception as e:
                        #print(f"failed to read param {motor_name} on interface {can_interface}: {e}")
                        motor.iface = None

                if found:
                    info(f"Motor found on interface: {can_interface}, id: {motor.motor_id}")
                    break
                time.sleep(0.01)

            if not found:
                motor.iface = None
                if not allow_not_found:
                    raise MotorNotFoundError(
                        f"Warning: Could not find motor {motor_name} on any physical CAN interface [New]",
                        motor_name_or_id=motor_name,
                    )

    @dont_allow_on_fork
    def stop_compliance_mode(self, family_name: str) -> None:
        """Stop compliance mode for all motors in the specified family.

        Args:
            family_name (str): Name of the motor family
        """
        family_motors = self.motor_families[family_name]
        for m in family_motors.values():
            m.compliance_mode = False
            m.enable_lock.acquire()
            self.disable(m.motor_name)

        # by setting it to None, trigger_compliance_mode will be called again from handle_msg from m.motor_name.
        if family_name in self.motor_families_compliance_threads:
            self.motor_families_compliance_threads[family_name].join(timeout=0.1)
            if self.motor_families_compliance_threads[family_name].is_alive():
                error(f"Error: Failed to join compliance mode thread for family {family_name} after 0.1 seconds.")
                raise Exception(f"Error: Failed to join compliance mode thread for family {family_name} after 0.1 seconds.")
            del self.motor_families_compliance_threads[family_name]
        
        time.sleep(0.033)
        max_retries = 10
        is_all_disabled = False
        enabled_motor: List[str] = []
        for i in range(max_retries):
            num_disabled = 0
            enabled_motor= []
            for m in family_motors.values():
                if m.mode.value == RobstrideMotorMode.Reset:
                    num_disabled += 1
                else:
                    enabled_motor.append(m.motor_name)
            if num_disabled == len(family_motors.keys()):
                is_all_disabled = True
                break
            else:
                warning(f"Warning: Failed to disable {', '.join(enabled_motor)} after {i} retries.")
                for m in enabled_motor:
                    self.disable(m)
            time.sleep(0.033)
        for m in family_motors.values():
            m.enable_lock.release()
        if not is_all_disabled:
            error(f"Error: Failed to disable {', '.join(enabled_motor)} after {max_retries} retries.")
            raise Exception(f"Error: Failed to disable {', '.join(enabled_motor)} after {max_retries} retries.")
        
        self._allow_trigger_compliance_mode[family_name] = True

    @dont_allow_on_fork
    def compliance_mode_main(self, family_name: str, motor_name: str) -> None:
        """Main loop for compliance mode that adjusts motor positions based on torque feedback.

        Args:
            family_name (str): Name of the motor family
        """
        # disable all motors
        try:
            self.activate_compliance_mode_deceleration_stop(family_name, motor_name)
        except Exception as e:
            error(f"Error in compliance mode main: {e}")
        finally:
            info(f"Final compliance mode main exiting")

    @dont_allow_on_fork
    def activate_compliance_mode_deceleration_stop(self, family_name: str, motor_name: str, safe_stop_policy: SafeStopPolicy = SafeStopPolicy.COMPLIANCE_MODE) -> None:
        info("Compliance mode deceleration stop")
        family_motors = self.motor_families[family_name]
        mkps = {}
        mkds = {}
        mkps_dx = {}
        mkds_dx = {}
        NUM_STEPS_TO_COMPLIANCE_MODE = 60
        TIME_PER_STEP = 0.015
        for m in family_motors.values():
            mkps[m.motor_name] = m.last_kp
            mkds[m.motor_name] = m.last_kd
            mkps_dx[m.motor_name] = (m.last_kp - COMPLIANCE_MODE_KP) / NUM_STEPS_TO_COMPLIANCE_MODE
            mkds_dx[m.motor_name] = (m.last_kd - COMPLIANCE_MODE_KD) / NUM_STEPS_TO_COMPLIANCE_MODE

        target_position = {}
        m = self.motors[motor_name]
        target_position[motor_name] = m.angle[0]
        self._set_target_position_in_compliance_mode(m.motor_name, target_position[m.motor_name], kp = mkps[m.motor_name], kd = mkds[m.motor_name])
        info(f"Urgent target position {m.motor_name} {target_position[m.motor_name]} velocity: {m.velocity[0]} {mkps[m.motor_name]} {mkds[m.motor_name]} {time.time()}")

        for m in family_motors.values():
            m.compliance_mode = True            
            target_position[m.motor_name] = m.angle[0] # self.read_param_sync(m.motor_name, "mech_pos")
            # info(f"reading mech_pos {m.motor_name} {target_position[m.motor_name]} {time.time()}")
        
        for i in range(NUM_STEPS_TO_COMPLIANCE_MODE):
            for m in family_motors.values():
                mkps[m.motor_name] -= mkps_dx[m.motor_name]
                mkds[m.motor_name] -= mkds_dx[m.motor_name]
                # if i < 10:
                #     target_position[m.motor_name] = m.angle[0]
                # self.disable(m.motor_name)
                if safe_stop_policy == SafeStopPolicy.COMPLIANCE_MODE:
                    # self.set_run_mode(m.motor_name, RobstrideRunMode.Position)
                    # self.enable(m.motor_name)
                    
                    self._set_target_position_in_compliance_mode(m.motor_name, target_position[m.motor_name], kp = mkps[m.motor_name], kd = mkds[m.motor_name])
                    debug(f"setting target position {m.motor_name} {target_position[m.motor_name]} curr_pos {m.angle[0]} velocity: {m.velocity[0]} {mkps[m.motor_name]} {mkds[m.motor_name]} {time.time()}")
                    m.last_compliance_pos = target_position[m.motor_name]
            time.sleep(TIME_PER_STEP)

        cnt = 0
        frequency = 30
        step_time = 1.0 / frequency
        while not self.aborted:
            if cnt % frequency / 10 == 0 and not check_still_own_lock():
                error("Compliance mode thread: Lock not owned by this process. Skipping.")
                break
            cnt += 1 
            for m in family_motors.values():
                if cnt % (frequency * 10) == 0:
                    info(f"compliance mode thread is running {m.motor_name} {m.compliance_mode} {time.time()}")
                if m.compliance_mode == False:
                    return  # exit.

                if abs(m.torque[0]) > m.compliance_mode_torque_threshold:
                    proposed_new_pos = m.last_compliance_pos + m.torque[0] * -m.compliance_mode_dx
                    if m.joint_limits is None:
                        m.last_compliance_pos = proposed_new_pos
                    elif m.joint_limits[0] <= proposed_new_pos and m.joint_limits[1] >= proposed_new_pos:
                        m.last_compliance_pos = proposed_new_pos
                # self.write does a time.sleep(0.0002) for each motor.
                self._set_target_position_in_compliance_mode(m.motor_name, m.last_compliance_pos)
                # time.sleep(0.01)
            time.sleep(step_time)

    @dont_allow_on_fork
    def trigger_compliance_mode(self, family_name: str, motor_name: str) -> None:
        """Enable compliance mode for all motors in the specified family.

        Args:
            family_name (str): Name of the motor family
        """
        if not self._allow_trigger_compliance_mode.get(family_name, True):
            warning(f"Compliance mode trigger blocked for family {family_name} because another compliance mode is running.")
            return
        self._allow_trigger_compliance_mode[family_name] = False

        if family_name not in self.motor_families_compliance_threads and not self.aborted:
            if not check_still_own_lock():
                error("Motors manager lock is owned by another process. Exiting program.")
                sys.exit(1)
            info("triggering compliance mode", family_name)
            thread = FeatherThread(name=f"compliance_mode_{family_name}", target=self.compliance_mode_main, args=(family_name, motor_name), daemon=True)
            self.motor_families_compliance_threads[family_name] = thread
            thread.start()

    @dont_allow_on_fork
    def on_motor_message(self, result: SocketResult) -> None:
        """Callback to handle incoming motor messages."""
        if result.is_can():
            self.on_can_message(result)
        elif result.is_tcp():
            self.on_tcp_message(result)
        else:
            raise NotImplementedError(f"Unhandled message type: {result.result_type}")
    
    @dont_allow_on_fork
    def on_can_message(self, result: SocketResult) -> None:
        """Callback to handle incoming CAN messages."""
        if not result.is_extended_can():
            debug(f"Non-extended CAN messages are not supported: {result}")
            return
        handle_robstride_motor_message(self, result)
    
    @dont_allow_on_fork
    def on_tcp_message(self, result: SocketResult) -> None:
        """Callback to handle incoming TCP messages."""
        raise NotImplementedError(f"TCP message handling not implemented")

    @dont_allow_on_fork
    def get_motor(self, motor_name_or_id: Union[str, int]) -> RobstrideMotor:
        """Get a motor instance by name or ID.

        Args:
            motor_name_or_id (Union[str, int]): Motor name or motor ID

        Returns:
            The RobstrideMotor instance

        Raises:
            MotorNotFoundError: If the motor is not found in the motors map
        """
        if isinstance(motor_name_or_id, str):
            if motor_name_or_id not in self.motors:
                raise MotorNotFoundError(f"Motor '{motor_name_or_id}' not found in motors map", motor_name_or_id=motor_name_or_id)
            return self.motors[motor_name_or_id]
        else:
            if motor_name_or_id not in self.motors_by_id:
                raise MotorNotFoundError(f"Motor with ID {motor_name_or_id} not found in motors map", motor_name_or_id=motor_name_or_id)
            return self.motors_by_id[motor_name_or_id]

    @dont_allow_on_fork
    def enable(self, motor_name_or_id: Union[str, int]) -> None:
        """Enable the specified motor.

        Args:
            motor_name_or_id: Motor name or ID
        """
        self.get_motor(motor_name_or_id).enable()

    @dont_allow_on_fork
    def disable(self, motor_name_or_id: Union[str, int]) -> None:
        """Disable the specified motor.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
        """
        self.get_motor(motor_name_or_id).disable()

    @dont_allow_on_fork
    def _disable_active_reporting(self, motor_name_or_id: Union[str, int]) -> None:
        """
        Send a command to disable active reporting, which triggers a feedback frame.
        
        This method sends a SetMotorActiveReporting command with active_flag=0 to the motor.
        While the intended purpose is to disable active reporting, this command also triggers
        the motor to send a feedback frame, which is the primary use case for this method.
        
        Note: Active reporting enable functionality was found to not work as
        documented in testing (TODO for later). This method is primarily used as a side-effect to trigger
        feedback frames rather than to actually control active reporting state.
        
        Args:
            motor_name_or_id: Motor name or ID
        """
        self.get_motor(motor_name_or_id).set_active_reporting(False)

    @dont_allow_on_fork
    async def read_current_state_async(self, motor_name_or_id: Union[str, int]) -> dict:
        """
        Asynchronously trigger a feedback frame from the motor and wait for fresh motor data.
        
        This method sends a command to the motor (via _disable_active_reporting) that triggers
        a feedback response. It then waits for the feedback frame to arrive and be processed,
        retrying the command if necessary.
        
        The method will retry sending the command every 100ms (10 iterations) if no feedback is
        received. After 300ms (30 iterations) with no feedback, a TimeoutError is raised to
        indicate that the motor is not responding. This ensures that stale feedback data is not
        used, as the method guarantees fresh feedback or an error.
        
        Args:
            motor_name_or_id: Name or ID of the motor
            
        Raises:
            TimeoutError: If no feedback frame is received within 300ms (after up to 3 retries)

        Note:
            This method has a side-effect of disabling active reporting, which is not the true purpose.
            Since we are not using active reporting, this is not a problem for now.
            But if we start using active reporting, we need to revisit this. There was a Motor Data
            Save command in the Robstride manual for the motor, however, that did not work as expected.
        """
        motor = self.get_motor(motor_name_or_id)
        
        # Record timestamp before triggering feedback
        update_time = currtime()
        
        # Disable active reporting to trigger a feedback frame
        self._disable_active_reporting(motor_name_or_id)
        
        # Wait for feedback frame to arrive and be processed
        # Check if temp timestamp has been updated (indicating fresh feedback received)
        await asyncio.sleep(0.001)  # Initial pause
        retry_cnt = 10  # Retry command every 10 iterations (10 * 10ms = 100ms)
        max_retry_cnt = 0
        
        # Wait for feedback to be updated (check temp timestamp as indicator)
        # temp[1] is the timestamp; -1 means no feedback received yet
        while motor.temp[1] < update_time:
            await asyncio.sleep(0.01)  # Wait 10ms between checks
            retry_cnt -= 1
            if retry_cnt <= 0:
                # Retry disabling active reporting if no feedback received
                self._disable_active_reporting(motor_name_or_id)
                retry_cnt = 10  # Reset retry counter
                max_retry_cnt += 1
            if max_retry_cnt > 30:  # Timeout after ~0.3 seconds (30 * 0.01s), allows 3 retries
                raise TimeoutError(
                    f"Failed to receive feedback from motor {motor.motor_name} within timeout (300ms). "
                    f"Motor may be disconnected or not responding."
                )
        return {
            "angle": motor.calibrated_angle,
            "angle_internal": motor.angle,
            "velocity": motor.velocity,
            "torque": motor.torque,
            "temp": motor.temp,
            "errors": motor.errors,
            "mode": motor.mode,
            "run_mode": motor.run_mode,
        }


    # TODO: Refactor all these read_* methods with a retry loop to be one single
    # method just waiting for some condition to be met.
    @dont_allow_on_fork
    def read_current_state_sync(self, motor_name_or_id: Union[str, int]) -> dict:
        """
        Asynchronously trigger a feedback frame from the motor and wait for fresh motor data.
        
        This method sends a command to the motor (via _disable_active_reporting) that triggers
        a feedback response. It then waits for the feedback frame to arrive and be processed,
        retrying the command if necessary.
        
        The method will retry sending the command every 100ms (10 iterations) if no feedback is
        received. After 300ms (30 iterations) with no feedback, a TimeoutError is raised to
        indicate that the motor is not responding. This ensures that stale feedback data is not
        used, as the method guarantees fresh feedback or an error.
        
        Args:
            motor_name_or_id: Name or ID of the motor
            
        Raises:
            TimeoutError: If no feedback frame is received within 300ms (after up to 3 retries)

        Note:
            This method has a side-effect of disabling active reporting, which is not the true purpose.
            Since we are not using active reporting, this is not a problem for now.
            But if we start using active reporting, we need to revisit this. There was a Motor Data
            Save command in the Robstride manual for the motor, however, that did not work as expected.
        """
        motor = self.get_motor(motor_name_or_id)
        
        # Record timestamp before triggering feedback
        update_time = currtime()
        
        # Disable active reporting to trigger a feedback frame
        self._disable_active_reporting(motor_name_or_id)
        
        # Wait for feedback frame to arrive and be processed
        # Check if temp timestamp has been updated (indicating fresh feedback received)
        time.sleep(0.001)  # Initial pause
        retry_cnt = 10  # Retry command every 10 iterations (10 * 10ms = 100ms)
        max_retry_cnt = 0
        
        # Wait for feedback to be updated (check temp timestamp as indicator)
        # temp[1] is the timestamp; -1 means no feedback received yet
        while motor.temp[1] < update_time:
            time.sleep(0.01)  # Wait 10ms between checks
            retry_cnt -= 1
            if retry_cnt <= 0:
                # Retry disabling active reporting if no feedback received
                self._disable_active_reporting(motor_name_or_id)
                retry_cnt = 10  # Reset retry counter
                max_retry_cnt += 1
            if max_retry_cnt > 30:  # Timeout after ~0.3 seconds (30 * 0.01s), allows 3 retries
                raise TimeoutError(
                    f"Failed to receive feedback from motor {motor.motor_name} within timeout (300ms). "
                    f"Motor may be disconnected or not responding."
                )
        return {
            "angle": motor.calibrated_angle,
            "angle_internal": motor.angle,
            "velocity": motor.velocity,
            "torque": motor.torque,
            "temp": motor.temp,
            "errors": motor.errors,
            "mode": motor.mode,
            "run_mode": motor.run_mode,
        }

    @dont_allow_on_fork
    def zero_position(self, motor_name_or_id: Union[str, int]) -> None:
        """Set the current position of the motor as zero.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
        """
        self.get_motor(motor_name_or_id).zero_position()

    @dont_allow_on_fork
    def read_param_sync(self, motor_name_or_id: Union[str, bytes], param_id: Union[int, str]) -> float:
        """Synchronously read a parameter from the motor, waiting for the value to be updated.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
            param_id: Parameter ID or name

        Returns:
            The parameter value

        Raises:
            TimeoutError: If the parameter is not read within the timeout period
        """
        motor = self.get_motor(motor_name_or_id)
        update_time = currtime()

        self.read_param(motor_name_or_id, param_id)
        time.sleep(0.001)
        retry_cnt = 50
        max_retry_cnt = 0
        while motor.get_property_timestamp(param_id, default=-1) < update_time:
            time.sleep(0.001)
            retry_cnt -= 1
            if retry_cnt <= 0:
                self.read_param(motor_name_or_id, param_id)    
                retry_cnt = 50
                max_retry_cnt += 1
            if max_retry_cnt > 25:
                raise TimeoutError(f"Failed to read parameter {param_id} for motor {motor_name_or_id} within timeout")
        if param_id == "mech_pos":
            motor.update_angle(motor.get_property_value(param_id))
        return motor.get_property_value(param_id)

    @dont_allow_on_fork
    async def read_param_async(self, motor_name_or_id: Union[str, bytes], param_id: Union[int, str]) -> float:
        """
        Asynchronously reads a motor parameter, waiting for the value to be updated.
        Replaces blocking time.sleep with non-blocking asyncio.sleep.
        """

        # 1. Initial setup and parameter read
        motor = self.get_motor(motor_name_or_id)

        # The 'update_time' from the original code: get the current time before requesting the update.
        update_time = currtime()

        # Initiate the read request (assuming self.read_param is already non-blocking or schedules a background action)
        self.read_param(motor_name_or_id, param_id)

        # Initial pause and setup
        await asyncio.sleep(0.001)  # Replace the first time.sleep(0.01) with non-blocking asyncio.sleep
        retry_cnt = 5
        max_retry_cnt = 0

        # Define the maximum total time to wait for the parameter to update (e.g., 0.5 seconds, as 25 * 0.02 is 0.5s total wait)
        # The original logic used 25 iterations of the main loop. With an inner sleep of 0.01s and an outer sleep/read sequence,
        # it's best to rely on a total timeout based on currtime() rather than a fixed retry count.
        # We will stick close to the original retry/max_retry logic, but use a more explicit maximum duration for robustness.

        # 2. Asynchronous Polling Loop
        while motor.get_property_timestamp(param_id, default=-1) < update_time:

            # Replace time.sleep(0.01) with non-blocking asyncio.sleep
            await asyncio.sleep(0.001)

            retry_cnt -= 1

            if retry_cnt <= 0:
                # Re-request the parameter read
                self.read_param(motor_name_or_id, param_id)
                retry_cnt = 5

                max_retry_cnt += 1

            # Original maximum loop count check (25 * 0.02s delay approx = 0.5s total wait)
            if max_retry_cnt > 25:
                # If a TimeoutError is raised, it's better to use the time-based check below,
                # but for a direct reimplementation, we keep the original logic.
                raise TimeoutError(f"Failed to read parameter {param_id} for motor {motor_name_or_id} within timeout")

        # 3. Return the updated parameter value
        return motor.get_property_value(param_id, default=-1)

    @dont_allow_on_fork
    def read_param(self, motor_name_or_id: Union[str, int], param_id: Union[int, str]) -> None:
        """Read a parameter from the specified motor."""
        self.get_motor(motor_name_or_id).read_param(param_id)
        
    @dont_allow_on_fork
    def set_run_mode(self, motor_name_or_id: Union[str, int], run_mode: RobstrideRunMode) -> None:
        """Set the run mode for the specified motor.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
            run_mode: The run mode to set (Position, Operation, or Torque)
        """
        self.get_motor(motor_name_or_id).set_run_mode(run_mode)

    @dont_allow_on_fork
    def _set_target_position_during_calibration_only(self, motor_name_or_id: Union[str, int], target_position: float) -> None:
        """This method is in general not safe to use because it ignores compliance mode. It should only be used during calibration.
        Args:
            motor_name_or_id: Name or ID of the motor
            target_position: Target position in radians
        """
        if os.environ.get("CALIBRATION_DEBUGGING") == "True":
            info (f"Setting target position {target_position} for motor {motor_name_or_id} during calibration only")
        else:
            self.get_motor(motor_name_or_id)._write_param_checked("loc_ref", target_position, force_run=True)

    def _set_target_position_in_compliance_mode(self, motor_name_or_id: Union[str, int], target_position: float, kp: float = None, kd: float = None) -> None:
        """Set the target position (loc_ref) for the specified motor, bypassing safety checks.
        
        This is an internal method used only by compliance mode to set target positions
        without triggering safety checks. It bypasses motor state validation and position
        safety checks that would normally prevent commands during unsafe conditions.
        
        Args:
            motor_name_or_id: Name or ID of the motor
            target_position: Target position in radians
        """
        motor = self.get_motor(motor_name_or_id)
        self._operation_command(motor_name_or_id, 0, target_position, 0, kp or motor.last_kp, kd or motor.last_kd, from_compliance_mode=True)

    def _move_motor_safely_for_calibration(
        self, 
        motor_name_or_id: Union[str, int], 
        target_position: float, 
        move_slowly: bool = False,
        start_position: Optional[float] = None,
        num_steps: int = 20, 
        slow_step_delay: float = 0.2,
        fast_delay: float = 1.5
    ) -> None:
        """Move a motor to a target position, either slowly (interpolated) or directly.
        
        This method can move a motor either:
        - Slowly: by interpolating between start_position and target_position in small steps
        - Fast: by directly setting the target position
        
        Args:
            motor_name_or_id: Name or ID of the motor
            target_position: Target position in radians
            move_slowly: If True, move slowly using interpolation; if False, move directly
            start_position: Starting position for slow movement (required if move_slowly=True)
            num_steps: Number of interpolation steps when moving slowly (default 20)
            slow_step_delay: Delay between steps when moving slowly in seconds (default 0.2)
            fast_delay: Delay after direct movement in seconds (default 1.5)
        """
        if move_slowly:
            if start_position is None:
                raise ValueError("start_position must be provided when move_slowly=True")
            for target_pos in np.linspace(start_position, target_position, num_steps):
                self._set_target_position_during_calibration_only(motor_name_or_id, target_pos)
                time.sleep(slow_step_delay)
        else:
            self._set_target_position_during_calibration_only(motor_name_or_id, target_position)
            time.sleep(fast_delay)

    @dont_allow_on_fork
    def set_target_position(self, motor_name_or_id: Union[str, int], target_position: float) -> None:
        """Set the target position (loc_ref) for the specified motor.
        
        This is a convenience method that wraps write_param for setting loc_ref.
        It provides a clearer API for position control.
        
        Args:
            motor_name_or_id: Name or ID of the motor
            target_position: Target position in radians
        
        Raises:
            MotorDisabledError: If the motor is in Reset mode (disabled)
            MotorModeInconsistentError: If the motor is not in Position run mode
        """
        self.get_motor(motor_name_or_id).set_target_position(target_position)

    @dont_allow_on_fork
    def set_target_velocity(self, motor_name_or_id: Union[str, int], target_velocity: float) -> None:
        """Set the target velocity (spd_ref) for the specified motor.
        
        This is a convenience method that wraps write_param for setting spd_ref.
        It provides a clearer API for velocity control.
        
        Args:
            motor_name_or_id: Name or ID of the motor
            target_velocity: Target velocity in radians per second
        
        Raises:
            MotorDisabledError: If the motor is in Reset mode (disabled)
            MotorModeInconsistentError: If the motor is not in Speed run mode
        """
        self.get_motor(motor_name_or_id).set_target_velocity(target_velocity)

    @dont_allow_on_fork
    def write_param(self, motor_name_or_id: Union[str, int], param_id: Union[int, str], param_value: Union[float, int]) -> None:
        """Write a parameter value to the specified motor.

        Args:
            motor_name_or_id: Name or ID of the motor
            param_id: Parameter ID
            param_value: Parameter value
        """
        self.get_motor(motor_name_or_id).write_param(param_id, param_value)

    @dont_allow_on_fork
    def operation_command(self, motor_name_or_id: Union[str, int], target_torque: float, target_angle: float, target_velocity: float, kp: float, kd: float) -> None:
        """Send an operation command to a motor with torque, angle, velocity, and PD gains.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
            target_torque: Target torque in Nm
            target_angle: Target angle in radians
            target_velocity: Target velocity in radians per second
            kp: Proportional gain
            kd: Derivative gain
        """
        self._operation_command(motor_name_or_id, target_torque, target_angle, target_velocity, kp, kd, from_compliance_mode=False)

    @dont_allow_on_fork
    def _operation_command(self, motor_name_or_id: Union[str, int], target_torque: float, target_angle: float, target_velocity: float, kp: float, kd: float, from_compliance_mode: bool = False):
        """Send an operation command to a motor with torque, angle, velocity, and PD gains.

        Args:
            motor_name_or_id: Motor name (str) or motor ID (int)
            target_torque: Target torque in Nm
            target_angle: Target angle in radians
            target_velocity: Target velocity in radians per second
            kp: Proportional gain
            kd: Derivative gain
            from_compliance_mode: Whether to the command was sent from compliance mode
        """
        self.send_operation_commands(
            [RobstrideOperationCommand(self.get_motor(motor_name_or_id), target_torque, target_angle, target_velocity, kp, kd)],
            from_compliance_mode=from_compliance_mode
        )

    @dont_allow_on_fork
    def send_operation_commands(self, operation_commands: List[RobstrideOperationCommand], from_compliance_mode: bool = False) -> None:
        """Send operation commands to multiple motors.

        Args:
            operation_commands: List of RobstrideOperationCommand instances to send
        """
        for op in operation_commands:
            op.motor.operation_command(op, from_compliance_mode=from_compliance_mode)

    @dont_allow_on_fork
    def _find_limit(
        self,
        motor_name: str,
        initial_pos: float,
        step_size: float,
        step_time: float,
        max_torque: float,
        threshold_target_distance: float,
        direction: int,  # 1 for forward, -1 for backward
        verbose: bool,
    ) -> float:
        """Helper method to find a limit (upper or lower) by moving the motor until torque limit is reached.
        
        Returns the limit position found.
        """
        motor = self.get_motor(motor_name)
        target_pos = initial_pos
        while True:
            target_pos += direction * step_size
            self._set_target_position_during_calibration_only(motor_name, target_pos)
            time.sleep(step_time)
            if verbose:
                debug("torque, angle", motor.torque[0], motor.angle[0], target_pos)
            if abs(motor.torque[0]) > max_torque:
                limit = motor.angle[0]
                if verbose:
                    debug(f"final_mech_pos (direction={direction})", motor.angle[0], target_pos)
                break
            # To avoid an infinite loop!
            target_distance = abs(target_pos - initial_pos)
            if target_distance > threshold_target_distance:
                curr_pos_from_motor = self.read_param_sync(motor_name, "mech_pos")
                if abs(curr_pos_from_motor - initial_pos) < 0.0175:  # 0.0175 in radians is 1 degree
                    help_string = f"Motor {motor_name} has not moved during calibration (current position: {curr_pos_from_motor}, initial position: {initial_pos})."
                else:
                    help_string = f"Motor {motor_name} has moved to {curr_pos_from_motor} after attempting to reach {target_pos}, but did not hit the limit. Calibration may be faulty."
                raise MotorCalibrationError(help_string, motor=motor)
        return limit

    @dont_allow_on_fork
    def _explore_range_and_move_to_home(
        self,
        motor_name: str,
        torque_threshold_for_limit_detection: float,
        step_size: float = 0.01,
        step_time: float = 0.01,
        verbose: bool = False,
        move_slowly_during_all_turns: bool = False,
    ) -> Tuple[float, float, float, float]:
        """Find the range of motion for a motor and move to the homing position at the end.
        
        The homing position is calculated as:
            homing_pos = lower_limit + (config.homing_pos - config.range.min)
        
        For example, if config has range.min=-10°, range.max=+160°, homing_pos=0°:
            - The homing position is 10° above the lower limit
            - physical_homing = lower_limit + (0 - (-10°)) = lower_limit + 10°
        
        Args:
            motor_name: Name of the motor
            torque_threshold_for_limit_detection: Torque threshold to detect physical limits
            step_size: Incremental position step in radians (default 0.01)
            step_time: Delay between steps in seconds (default 0.01)
            verbose: Enable debug output (default False)
            move_slowly_during_all_turns: Whether to move slowly during exploration turns in the middle of the range (default False), 
                    this should be set to true for motors that could cause jerk when moving too quickly like certain arm motors.

        Returns:
            Tuple of (lower_limit, upper_limit, homing_pos, total_range)
        """
        if step_size < 0.001:
            raise ValueError(f"Step size is too low or negative, got {step_size}, recommended value is 0.01")
        if step_time < 0.001:
            raise ValueError(f"Step time is too low or negative, got {step_time}, recommended value is 0.01")
        if torque_threshold_for_limit_detection <= 0.0:
            raise ValueError(f"Torque threshold for limit detection is not positive, got {torque_threshold_for_limit_detection}, recommended value depends on the platform or motor type.")

        motor = self.get_motor(motor_name)
        expected_range = motor.expected_range
        config_range_min = motor.range["min"]
        config_homing_pos = motor._homing_pos
        
        self.set_run_mode(motor_name, RobstrideRunMode.Position)
        self.enable(motor_name)
        start_pos = self.read_param_sync(motor_name, "mech_pos")

        info(f"Moving the motor {motor_name} to one end to get the range of motion, starting from {start_pos}")
        # Find upper limit (moving forward)
        upper_limit = self._find_limit(
            motor_name=motor_name,
            initial_pos=start_pos,
            step_size=step_size,
            step_time=step_time,
            max_torque=torque_threshold_for_limit_detection,
            threshold_target_distance=expected_range * 2.0,
            direction=1,  # forward
            verbose=verbose,
        )
        
        info(f"Reached one end at upper_limit {upper_limit}, trying to go to the other end by setting initial target to {upper_limit - (expected_range * 0.8)} and exploring the other end")
        self._move_motor_safely_for_calibration(
            motor_name, 
            upper_limit - (expected_range * 0.8),
            move_slowly=move_slowly_during_all_turns,
            start_position=upper_limit
        )
        
        lower_limit_initial_pos = upper_limit - (expected_range * 0.8)
        # Find lower limit (moving backward)
        lower_limit = self._find_limit(
            motor_name=motor_name,
            initial_pos=lower_limit_initial_pos,
            step_size=step_size,
            step_time=step_time,
            max_torque=torque_threshold_for_limit_detection,
            threshold_target_distance=expected_range * 2.5,
            direction=-1,  # backward
            verbose=verbose,
        )

        # Calculate homing position based on config
        # homing_offset = distance from lower limit to homing position
        # For range.min=-10°, homing_pos=0°: homing is 10° above lower limit
        homing_offset_from_lower = config_homing_pos - config_range_min
        homing_pos = lower_limit + homing_offset_from_lower

        if verbose:
            debug(f"Calculated homing position: {homing_pos:.4f} rad")
            debug(f"  (config_homing_pos={config_homing_pos:.4f}, config_range_min={config_range_min:.4f})")
            debug(f"  (homing_offset_from_lower={homing_offset_from_lower:.4f})")

        info(f"Reached the other end at lower_limit {lower_limit}, trying to go to the homing position by setting target position to {homing_pos}")
        self._move_motor_safely_for_calibration(
            motor_name,
            homing_pos,
            move_slowly=move_slowly_during_all_turns,
            start_position=lower_limit,
            fast_delay=2.5,
            slow_step_delay=0.4,
        )

        return lower_limit, upper_limit, homing_pos, abs(upper_limit - lower_limit)

    @dont_allow_on_fork
    def _is_valid_range(self, total_range: float, expected_range: float, tolerance: float = 0.005) -> bool:
        """Check if the found range is within tolerance of the expected range.
        
        Args:
            total_range: The measured total range in radians
            expected_range: The expected range from motor config in radians
            tolerance: Tolerance as a fraction (default 0.1 = 10%)
        
        Returns:
            True if the range is valid (within tolerance), False otherwise
        """
        return total_range >= expected_range * (1.0 - tolerance) and total_range <= expected_range * (1.0 + tolerance)

    @dont_allow_on_fork
    def _calibrate_motor(
        self,
        motor_name: str,
        torque_threshold_for_limit_detection: float,
        step_size: float = 0.01,
        step_time: float = 0.01,
        max_retries: int = 2,
        verbose: bool = False,
        move_slowly_during_all_turns: bool = False,
    ) -> CalibrationWarnings:
        """Calibrate a motor by finding its range and setting the homing position as zero.

        This method is for internal use by high level systems to calibrate motors.
        DO NOT CALL THIS METHOD DIRECTLY. Calling this directly for example
        on arm motors will bypass all safety logic known to arm system and dangerous to either the robot or the user.
        
        This method finds the physical range of the motor, calculates where the homing
        position should be based on the motor config, sets that position as the zero 
        position, and then refines the calibration with a second range discovery.
        
        The second get_range call captures any residual offset after zeroing (since the
        motor may not be exactly at position 0 after the zero command). This residual
        is saved as calibration_homing_pos and used as an internal offset for motor 
        commands - when a user asks to move to homing position, we move to this 
        internal calibration_homing_pos.
        
        Args:
            motor_name: Name of the motor to calibrate
            step_size: Incremental position step in radians (default 0.01)
            step_time: Delay between steps in seconds (default 0.01)
            max_retries: Maximum retry attempts if range is invalid (default 2)
            verbose: Enable debug output (default False)
            move_slowly_during_all_turns: Whether to move slowly during all turns (default False)
        
        Returns:
            CalibrationWarnings after calibration
        
        Raises:
            MotorCalibrationError: If calibration fails.
        """
        motor = self.get_motor(motor_name)
        motor._disable_compliance_mode_for_calibration()
        try:
            expected_range = motor.expected_range
            
            # Step 1: Find initial range and move to homing position
            lower_limit, upper_limit, homing_pos, total_range = self._explore_range_and_move_to_home(
                motor_name,
                torque_threshold_for_limit_detection,
                step_size=step_size,
                step_time=step_time,
                verbose=verbose,
                move_slowly_during_all_turns=move_slowly_during_all_turns,
            )
            
            if verbose:
                debug(f"Initial range discovery: {lower_limit=}, {upper_limit=}, {homing_pos=}, {total_range=}")
            
            # Step 2: Retry if range is not valid
            retries = 0
            while not self._is_valid_range(total_range, expected_range) and retries < max_retries:
                info(f"Warning: Range {total_range:.4f} is not close to expected {expected_range:.4f} radians. Retrying ({retries + 1}/{max_retries})...")
                
                # Zero position and retry
                self._zero_position_twice(motor_name)
                
                lower_limit, upper_limit, homing_pos, total_range = self._explore_range_and_move_to_home(
                    motor_name,
                    torque_threshold_for_limit_detection,
                    step_size=step_size,
                    step_time=step_time,
                    verbose=verbose,
                    move_slowly_during_all_turns=move_slowly_during_all_turns,
                )
                
                if verbose:
                    debug(f"Retry {retries + 1}: {lower_limit=}, {upper_limit=}, {homing_pos=}, {total_range=}")
                
                retries += 1
            
            if not self._is_valid_range(total_range, expected_range):
                raise MotorCalibrationError(
                    f"Failed to find valid range for motor {motor_name} after {max_retries} attempts. "
                    f"Found range: {total_range:.4f} rad, expected: {expected_range:.4f} rad. "
                    f"Please check the motor config and mechanical constraints.",
                    motor=motor
                )
            
            # Step 3: Set homing position as zero
            # Motor is already at homing_pos from get_range_to_homing
            self._zero_position_twice(motor_name)
            
            if verbose:
                debug(f"Set zero position at homing position {homing_pos}")
            
            # Step 4: Do another get_range_to_homing to find refined calibration values
            # This captures any residual offset after zeroing
            lower_limit, upper_limit, calibration_homing_pos, total_range = self._explore_range_and_move_to_home(
                motor_name,
                torque_threshold_for_limit_detection,
                step_size=step_size,
                step_time=step_time,
                verbose=verbose,
                move_slowly_during_all_turns=move_slowly_during_all_turns,
            )
            
            if verbose:
                debug(f"Final calibration: {lower_limit=}, {upper_limit=}, {calibration_homing_pos=}, {total_range=}")
            
            # Step 5: Save calibration data
            # calibration_homing_pos is the actual homing position after zeroing,
            # which may have a small residual offset from 0
            motor.save_calibration_state(lower_limit, upper_limit, calibration_homing_pos, total_range)
            warnings = motor._calibration_checks()

            info(f"Motor {motor_name} calibration complete: lower={lower_limit:.4f}, upper={upper_limit:.4f}, "
                f"homing={calibration_homing_pos:.4f}, range={total_range:.4f}")
            
            return warnings
        finally:
            motor._enable_compliance_mode_after_calibration()

    def _zero_position_twice(self, motor_name: str) -> None:
        """Zero the motor position twice for improved accuracy.

        Args:
            motor_name: Name of the motor to zero
        """
        self.disable(motor_name)
        for _ in range(2):  # Double zero for accuracy
            self.zero_position(motor_name)
            time.sleep(0.1)

    @dont_allow_on_fork
    def recover_from_power_cycle(self, motor_name: str) -> None:
        """Recover motor calibration after a power cycle where the encoder was reset.

        Args:
            motor_name: Name of the motor to recover

        Raises:
            UnsafeCommandError: If the motor is not a dual encoder motor or has no calibration
        """
        motor = self.get_motor(motor_name)
        # recovers the motor from a power cycle where the encoder is reset.
        if not motor.dual_encoder:
            raise UnsafeCommandError("Only dual encoder motors can recover from power cycle.", motor=motor)
        if motor.calibration_time is None:
            raise UnsafeCommandError("Motor does not have a calibration to recover from.", motor=motor)
        
        mech_pos = self.read_param_sync(motor_name, "mech_pos")
        if mech_pos > motor.upper_limit:
            # too high by 2 pi
            motor.upper_limit += 2 * math.pi
            motor.lower_limit += 2 * math.pi
            motor.calibration_homing_pos += 2 * math.pi
            
        if mech_pos < motor.lower_limit:
            # too low by 2 pi
            motor.upper_limit -= 2 * math.pi
            motor.lower_limit -= 2 * math.pi
            motor.calibration_homing_pos -= 2 * math.pi

    @dont_allow_on_fork
    def check_motors_in_range(self, motor_names: List[str], target_angle: Optional[float] = None, raise_error: bool = True) -> bool:
        """Check if motors are within valid range.
        
        Validates that motors' current positions and optionally a target angle are within
        their calibrated limits (upper_limit and lower_limit).
        
        Args:
            motor_names: List of motor names to check
            target_angle: Optional target angle to validate against limits
            raise_error: If True, raise exceptions on violations; if False, return False
            
        Returns:
            True if all motors are in range, False otherwise
            
        Raises:
            ValueError: If target_angle is out of range and raise_error is True
            Exception: If current motor position is out of range and raise_error is True
        """
        for motor_name in motor_names:
            motor = self.get_motor(motor_name)
            
            # Check target angle if provided
            if target_angle is not None:
                if target_angle > motor.upper_limit or target_angle < motor.lower_limit:
                    if raise_error:
                        raise ValueError(
                            f"{target_angle} is out of range for {motor_name}. "
                            f"Range: {motor.lower_limit} to {motor.upper_limit}"
                        )
                    return False
            
            # Check motor current angle. if not in range even with +-2*pi difference, raise error, tell user to recalibrate
            try:
                mech_pos = self.read_param_sync(motor_name, "mech_pos")
            except Exception as e:
                if raise_error:
                    raise
                warning(f"Error reading mech_pos for {motor_name}: {e} in check_motors_in_range")
                return False
            if mech_pos is not None and (mech_pos > motor.upper_limit or mech_pos < motor.lower_limit):
                if raise_error:
                    raise Exception(f"{motor_name} position {mech_pos} is outside of valid range. Please recalibrate.")
                return False
        
        return True

    @dont_allow_on_fork
    def motors_in_mode(self, motor_names: List[str], mode: RobstrideMotorMode) -> bool:
        """Check if all specified motors are in the given run mode.

        Args:
            motor_names: List of motor names to check
            mode: The run mode to check for

        Returns:
            True if all motors are in the specified mode, False otherwise
        """
        for motor_name in motor_names:
            motor = self.get_motor(motor_name)
            if motor.run_mode != mode:
                return False
        return True
        
    @dont_allow_on_fork
    def ping(self, motor_name_or_id: Union[str, int]) -> bool:
        """ 
        Ping the motor to get back the current state of the motor.
        """
        motor = self.get_motor(motor_name_or_id)
        # skip pinging if the motor is locked by another thread.
        if motor.enable_lock.locked():
            return False
        if motor.mode[0] == RobstrideMotorMode.Run:
            motor.enable()
        else:
            motor.disable()
    
    def on_fork(self):
        self.motors: MotorMap = {}
        self.motors_by_id: Dict[int, RobstrideMotor] = {}
        self.motor_families: Dict[str, MotorMap] = {}
        self.motor_families_config: Dict[str, FamilyConfig] = {}
        self.motor_families_compliance_threads = {}
        self.aborted = True
        self._allow_trigger_compliance_mode: Dict[str, bool] = {}
        self.is_forked = True
