import math
from typing import Dict, List, Optional, Set, Tuple, Union
import time
from .steppable_system import SteppableSystem
from .motors.robstride import RobstrideRunMode, RobstrideMotorMode
from .motors.motors_manager import (
    MotorsManager, MotorMap, CalibrationWarnings, CalibrationState, MotorCalibrationError
)
from ..utils.trajectory import Trajectory
from ..utils.logger import error, warning
from feathersdk.utils.feathertypes import *
from .motors.robstride.robstride_motor import TimestampedValue

MOTORS_TO_MOVE_SLOW: Set[str] = {"EpR", "EpL", "SpR", "SpL", "SrR", "SrL"}
DEFAULT_KP = 30
DEFAULT_KD = 5

CALIBRATION_POSITIONS_TO_HOLD = {
    "WrR": {
        "WwR": 90 * math.pi / 180,
    },
    "WrL": {
        "WwL": 90 * math.pi / 180,
    },
    "SpR": {
        "EpR": 150 * math.pi / 180,
    },
    "SpL": {
        "EpL": -150 * math.pi / 180,
    },
    "SrR": {
        "EpR": 150 * math.pi / 180,
        "SpR": -60 * math.pi / 180,
    },
    "SrL": {
        "EpL": -150 * math.pi / 180,
        "SpL": 60 * math.pi / 180,
    },
    "SwR": {
        "SrR": -15 * math.pi / 180,
    },
    "SwL": {
        "SrL": 15 * math.pi / 180,
    }
}

def _confirm_calibration(system_name: str) -> bool:
    """Print a caution message and ask user for confirmation before calibrating.
    
    Args:
        system_name: Name of the system being calibrated (for display purposes)
        
    Returns:
        True if user confirms (y/yes), False otherwise
    """
    print("\n" + "=" * 60)
    print("⚠️  CAUTION: CALIBRATION WARNING ⚠️")
    print("=" * 60)
    print(f"You are about to calibrate: {system_name}")
    print()
    print("During calibration, the robot arms will move automatically. This is inherently unsafe.")
    print("Please ensure:")
    print("  • The robot is in a safe position")
    print("  • No obstacles are in the arm's range of motion")
    print("  • No Personnel are within the arm's range of motion")
    print("  • If the calibration failed, you will see another message on what needs to be done for safety.")
    print("=" * 60)
    
    response = input("Do you want to proceed with calibration? (y/yes to confirm): ").strip().lower()
    
    if response in ("y", "yes"):
        print("Proceeding with calibration...\n")
        return True
    else:
        print("Calibration cancelled.\n")
        return False


def _show_calibration_failure_recovery(system_name: str, error: Exception, failed_motor: Optional[str] = None) -> None:
    """Display recommended recovery actions when calibration fails.
    
    Args:
        system_name: Name of the system that failed calibration
        error: The exception that was raised
        failed_motor: Optional name of the specific motor that failed
    """
    print("\n" + "!" * 60)
    print("🚨  CALIBRATION FAILURE - IMMEDIATE ACTION REQUIRED 🚨")
    print("!" * 60)
    print(f"System: {system_name}")
    if failed_motor:
        print(f"Failed motor: {failed_motor}")
    print(f"Error: {error}")
    print()
    print("⚠️  WARNING: The robot may be in an unsafe state!")
    print()
    print("RECOMMENDED ACTIONS - READ ALL THE MESSAGES BEFORE TAKING ACTION:")
    print("-" * 40)
    print("  1. DO NOT leave the robot unattended")
    print("  2. Some motors may still be in POSITION HOLD mode")
    print("     and actively resisting movement")
    print("  3. To safely disable motors, run:")
    print("       robot.manipulation_platform.right_arm.disable()")
    print("       robot.manipulation_platform.left_arm.disable()")
    print("  4. If motors cannot be disabled via software:")
    print("       • Use the hardware emergency stop")
    print("  5. Manually support the arm before disabling motors")
    print("     to prevent it from falling due to gravity")
    print("  6. Investigate the error before retrying calibration")
    print("-" * 40)
    print("BEFORE RETRYING:")
    print("  • Check for mechanical obstructions")
    print("  • Verify motor connections and power")
    print("  • Ensure the arm is in a safe starting position")
    print("  • Review error logs for root cause")
    print("!" * 60 + "\n")


class ArmSystem(SteppableSystem):
    """Single-arm motor control: trajectories, enable/disable, and calibration.

    Each arm registers its motors with a shared :class:`MotorsManager`, loads
    calibration on construction, and runs per-motor trajectories in :meth:`_step`
    at ``operating_frequency`` (Hz), with periodic CAN pings staggered across motors.
    """

    def __init__(self, motors_manager: MotorsManager, motors_map: MotorMap, name: str, calibration_order: List[str]):
        super().__init__()
        self.motors_manager = motors_manager
        self.motors_manager.add_motors(motors_map, family_name=name)
        self.motors_manager.find_motors(list(motors_map.keys()))
        self.motors = motors_map
        self.motor_list = list(self.motors.values())
        for motor_name in self.motors.keys():
            self.motors_manager.ping(motor_name)
        self.motors_trajectories = {motor_name: Trajectory(motor) for motor_name, motor in self.motors.items()}
        self.name = name
        self.calibration_order = calibration_order
        self.calibration_state = CalibrationState.UNINITIALIZED
        self.calibration_warnings: Dict[str, CalibrationWarnings] = {}
        # 300 hz is the maximum speed of the Robstride motors.
        self.operating_frequency = 300
        self.MOTORS_FORCE_REFRESH_HZ = 10
        # Try to load existing calibration
        self._load_calibration()

    def _load_calibration(self) -> bool:
        """Load calibration data for all motors in this arm.
        
        Returns True if all motors loaded successfully, False otherwise.
        """
        all_loaded = True
        for motor_name in self.motors.keys():
            motor = self.motors[motor_name]
            if not motor.load_calibration_state():
                all_loaded = False
        
        if all_loaded:
            for motor_name in self.motors.keys():
                self.motors_manager.recover_from_power_cycle(motor_name)
            self.calibration_state = CalibrationState.HEALTHY
            print(f"{self.name}: Calibration data loaded successfully")
        else:
            self.calibration_state = CalibrationState.UNCALIBRATED
            print(f"{self.name}: Calibration data not found or incomplete, recalibration required")

        return all_loaded

    def _hold_motors_at_current_position(self, calibrating_motor_name: str, verbose: bool = False) -> Tuple[dict, dict]:
        """Hold all arm motors except ``calibrating_motor_name`` at fixed positions.

        Args:
            calibrating_motor_name (str): Motor currently being calibrated; it is excluded from hold.
            verbose (bool): If True, print hold positions during calibration.

        Returns:
            Tuple[dict, dict]: ``(hold_positions, reset_positions)`` — per-motor encoder targets used
            for hold and the positions to restore when releasing hold.
        """
        all_arm_motors = list(self.motors.keys())
        motor_names_to_hold = [m for m in all_arm_motors if m != calibrating_motor_name]
        hold_positions = {}
        reset_positions = {}
        for motor_name in motor_names_to_hold:
            try:
                # Read current position 
                hold_position = self.motors_manager.read_param_sync(motor_name, "mech_pos")
                reset_position = hold_position
                # We want to reset to current position when we release the position hold
                reset_positions[motor_name] = reset_position

                if calibrating_motor_name in CALIBRATION_POSITIONS_TO_HOLD:
                    if motor_name in CALIBRATION_POSITIONS_TO_HOLD[calibrating_motor_name]:
                        hold_position = CALIBRATION_POSITIONS_TO_HOLD[calibrating_motor_name][motor_name]
                
                # Set to position mode and enable
                self.motors_manager.set_run_mode(motor_name, RobstrideRunMode.Position)
                self.motors_manager.enable(motor_name)
                
                # Set target to current position to hold
                should_move_slowly = motor_name in MOTORS_TO_MOVE_SLOW and hold_position != reset_position
                self.motors_manager._move_motor_safely_for_calibration(
                    motor_name,
                    hold_position,
                    move_slowly=should_move_slowly,
                    start_position=reset_position
                )
                hold_positions[motor_name] = hold_position

                if verbose:
                    print(f"  Holding {motor_name} at position {hold_position:.4f}")
                    
            except Exception as e:
                print(f"  Warning: Failed to hold {motor_name}: {e}")
                raise e
                
        return hold_positions, reset_positions

    def _release_position_hold(self, calibrating_motor_name: str, hold_positions: dict, reset_positions: dict, verbose: bool = False) -> None:
        """Move held motors back to reset positions, then disable them.
        
        Args:
            calibrating_motor_name (str): Motor being calibrated; excluded from release.
            hold_positions (dict): Per-motor positions used during hold (see :meth:`_hold_motors_at_current_position`).
            reset_positions (dict): Per-motor positions to restore before disable.
            verbose (bool): If True, print reset and disable progress.
        """
        all_arm_motors = list(self.motors.keys())
        motor_names_to_release = [m for m in all_arm_motors if m != calibrating_motor_name]
        for motor_name in motor_names_to_release:
            try:
                should_move_slowly = motor_name in MOTORS_TO_MOVE_SLOW and hold_positions[motor_name] != reset_positions[motor_name]
                self.motors_manager._move_motor_safely_for_calibration(
                    motor_name,
                    reset_positions[motor_name],
                    move_slowly=should_move_slowly,
                    start_position=hold_positions[motor_name]
                )
                if verbose:
                    print(f"  Reset position of {motor_name} to {reset_positions[motor_name]:.4f}")
            except Exception as e:
                print(f"  Warning: Failed to release {motor_name}: {e}")
                raise e
        
        for motor_name in motor_names_to_release:
            try:
                self.motors_manager.disable(motor_name)
                if verbose:
                    print(f"  Disabled {motor_name}")
            except Exception as e:
                print(f"  Warning: Failed to disable {motor_name}: {e}")
                raise e

    def _recalibrate(
        self,
        motor_names: Optional[List[str]] = None,
        step_time: float = 0.01,
        verbose: bool = False,
        _skip_confirmation: bool = False
    ) -> None:
        """Recalibrate arm motors using the unified calibrate_motor method.
        
        Motors are calibrated in order from bottom-most (wrist) to top-most (shoulder)
        to ensure stable calibration as lower joints are calibrated first.
        
        During calibration of each motor, all OTHER motors in the arm are held in
        position mode to resist gravity and inertial forces from the calibrating
        motor's movement. This prevents physically coupled joints from drifting.
        
        Args:
            motor_names: Optional list of specific motors to calibrate. 
                        If None, calibrates all motors in the arm.
            step_time: Delay between steps in seconds (default 0.01)
            verbose: Enable debug output (default False)
            _skip_confirmation: Internal flag to skip user confirmation prompt.
                               Set to True when called from BimanualManipulationPlatform
                               to avoid duplicate prompts. (default False)
        
        Raises:
            Exception: If calibration is already in progress or user cancels
            MotorCalibrationError: If calibration fails for any motor
        """
        # Ask for user confirmation unless called from parent platform
        if not _skip_confirmation:
            if not _confirm_calibration(self.name):
                raise Exception(f"Calibration cancelled by user for {self.name}")
        
        if self.calibration_state == CalibrationState.RECALIBRATING:
            raise Exception(f"Error: {self.name} is already calibrating, please wait for calibration to complete.")
        
        # Determine which motors to calibrate
        if motor_names is None:
            motors_to_calibrate = self.calibration_order
        else:
            # Validate motor names and sort by calibration order
            for name in motor_names:
                if name not in self.motors:
                    raise ValueError(f"Motor {name} is not part of {self.name}. Valid motors: {list(self.motors.keys())}")
            # Sort by calibration order (bottom to top)
            motors_to_calibrate = [m for m in self.calibration_order if m in motor_names]
        
        try:
            self.calibration_state = CalibrationState.RECALIBRATING
            print(f"{self.name}: Starting calibration for motors: {motors_to_calibrate}")

            for motor in self.motors.values():
                motor._disable_compliance_mode_for_calibration()
            
            for motor_name in motors_to_calibrate:
                print(f"{self.name}: Calibrating {motor_name} (holding other motors in position)...")
                
                # Enable position hold on all other motors
                hold_positions, reset_positions = self._hold_motors_at_current_position(motor_name, verbose=verbose)
                time.sleep(0.1)  # Short delay for position hold to stabilize
                
                # Calibrate the target motor
                warnings = self.motors_manager._calibrate_motor(
                    motor_name,
                    torque_threshold_for_limit_detection=6.0,
                    step_time=step_time,
                    verbose=verbose,
                    move_slowly_during_all_turns=motor_name in MOTORS_TO_MOVE_SLOW,
                )
                self.calibration_warnings[motor_name] = warnings
                # Release position hold on other motors
                self._release_position_hold(motor_name, hold_positions, reset_positions, verbose=verbose)

                time.sleep(0.1)  # Small delay between motors
            
            self.calibration_state = CalibrationState.HEALTHY
            print(f"{self.name}: Calibration completed successfully!")
            
        except MotorCalibrationError as e:
            self.calibration_state = CalibrationState.CALIBRATION_FAILED
            _show_calibration_failure_recovery(self.name, e, failed_motor=e.motor_name)
            raise Exception(f"Error: {self.name} calibration failed on motor {e.motor_name}: {e}")
        except Exception as e:
            self.calibration_state = CalibrationState.CALIBRATION_FAILED
            _show_calibration_failure_recovery(self.name, e)
            raise Exception(f"Error: {self.name} calibration failed: {e}")
        finally:
            for motor in self.motors.values():
                motor._enable_compliance_mode_after_calibration()

    def is_calibrated(self) -> bool:
        """Check if the arm is calibrated and ready for use."""
        return self.calibration_state == CalibrationState.HEALTHY

    def get_state(self) -> dict:
        """Get synchronous snapshot state for all motors in this arm.

        Uses the last cached feedback values on each motor (no extra CAN read).

        Returns:
            dict: Maps each motor name to a dict with:

                - ``temp``: temperature (scalar from last feedback)
                - ``torque``: torque (scalar from last feedback)
                - ``angle``: calibrated angle in radians
                - ``velocity``: velocity in rad/s
        """
        state = {}
        for motor_name, motor in self.motors.items():
            state[motor_name] = {
                "temp": motor.temp[0],
                "torque": motor.torque[0],
                "angle": motor.calibrated_angle[0],
                "velocity": motor.velocity[0],
            }
        return state

    async def get_state_async(self, motor_name: str) -> Dict[str, TimestampedValue]:
        """Fetch fresh feedback for one motor and return timestamped fields.

        Args:
            motor_name (str): Name of a motor belonging to this arm.

        Returns:
            dict[str, TimestampedValue]: Keys ``temp``, ``torque``, ``angle`` (calibrated),
            and ``velocity``.
        """
        await self.motors_manager.read_current_state_async(motor_name)
        motor = self.motors_manager.motors[motor_name]
        return {
            "temp": motor.temp,
            "torque": motor.torque,
            "angle": motor.calibrated_angle,
            "velocity": motor.velocity
        }

    def go_to(self, motors_targets: Dict[str, float], total_time: float):
        """Go to a target position for all motors in the arm.
        
        Args:
            motors_targets (Dict[str, float]): Motor name → target **calibrated** angle (rad).
            total_time (float): Duration (s) over which trajectories should complete.
        """
        for motor_name, calibrated_target_position in motors_targets.items():
            self.motors_trajectories[motor_name].update_trajectory(calibrated_target_position, int(total_time * self.operating_frequency))

    def freeze_current_arm_state(self, kp: Union[float, Dict[str, float]] = DEFAULT_KP, kd: Union[float, Dict[str, float]] = DEFAULT_KD):
        """Hold the arm at the current pose using operation-mode gains.

        Sets ``last_kp`` / ``last_kd`` on each motor (scalar or per-motor dict), then
        commands an immediate move (short ``total_time``) to current calibrated angles.

        Args:
            kp (float | Dict[str, float]): Position gain or per-motor overrides. Default 30.
            kd (float | Dict[str, float]): Damping gain or per-motor overrides. Default 5.
        """
        calibrated_targets = {}
        for m in self.motors.values():
            if isinstance(kp, dict):
                if m.motor_name in kp:
                    m.last_kp = kp[m.motor_name]
                else:
                    m.last_kp = DEFAULT_KP
            else:
                m.last_kp = kp
            if isinstance(kd, dict):
                if m.motor_name in kd:
                    m.last_kd = kd[m.motor_name]
                else:
                    m.last_kd = DEFAULT_KD
            else:
                m.last_kd = kd
            calibrated_targets[m.motor_name] = m.calibrated_angle[0]
        self.go_to(calibrated_targets, 0.01)

    def home(self, time: float = 2.0):
        """Move all joints toward calibrated angle zero over ``time`` seconds.

        Args:
            time (float): Trajectory duration in seconds (default 2.0).
        """
        self.go_to({m.motor_name: 0 for m in self.motors.values()}, time)

    def _step(self):
        curr_remainder = int(self.current_step % (self.operating_frequency / self.MOTORS_FORCE_REFRESH_HZ))
        if curr_remainder < len(self.motor_list):
            # Don't overload the CAN bus by pinging all motors at the same time.
            m = self.motor_list[curr_remainder]
            try:
                self.motors_manager.ping(m.motor_name)
            except Exception as e:
                warning(f"Failed to ping {m.motor_name}: {e}")
        for m in self.motors.values():
            if m.is_runable() and not self.motors_trajectories[m.motor_name].is_at_end():
                try:
                    self.motors_manager.operation_command(m.motor_name, 0.0, m.convert_calibrated_angle_to_angle(self.motors_trajectories[m.motor_name].step()), 0.0, m.last_kp, m.last_kd)
                except Exception as e:
                    warning(f"Failed to send operation command to {m.motor_name}: {e}")

    def enable(self):
        """Enable all arm motors in operation mode with retries.

        Resets trajectories first. Acquires each motor's ``enable_lock`` until Run mode
        is confirmed, then releases locks for motors that succeeded.

        Raises:
            Exception: If this arm's family is still in compliance mode, or if any motor
                fails to enter Run mode after ``max_retries``.
        """
        if self.name in self.motors_manager.motor_families_compliance_threads:
            error(f"Error: Robot cannot be enabled when still {self.name} is in compliance mode.")
            raise Exception(f"Error: Robot cannot be enabled when still {self.name} is in compliance mode.")
        self.reset_trajectories()
        max_retries = 10
        is_all_enabled = False
        failed_motor: List[str] = []
        for m in self.motors.values():
            m.enable_lock.acquire()
        for i in range(max_retries):
            motors_enabled = 0
            if i > 0:
                warning(f"Warning: Enabled failed for {failed_motor} after {i} attempts.")
            failed_motor = []
            for motor_name in self.motors.keys():
                self.motors_manager.enable(motor_name)
            time.sleep(0.033)
            for motor_name in self.motors.keys():
                if self.motors_manager.motors[motor_name].mode.value == RobstrideMotorMode.Run:
                    motors_enabled += 1
                    self.motors[motor_name].enable_lock.release()
                else:
                    failed_motor.append(motor_name)
            is_all_enabled = motors_enabled == len(self.motors.keys())
            if is_all_enabled:
                break
        if not is_all_enabled:
            for m in self.motors.values():
                m.enable_lock.release()
            error(f"Error: Failed to enable {', '.join(failed_motor)} after {max_retries} retries.")
            raise Exception(f"Error: Failed to enable {', '.join(failed_motor)} after {max_retries} retries.")

    def disable(self):
        """Disable this arm via the motors manager (stops compliance mode for this family)."""
        self.motors_manager.stop_compliance_mode(self.name)

    def reset_trajectories(self):
        """Clear trajectory state and align internal targets with current encoder angles.

        Sets each motor's ``target_position`` to the raw ``angle`` feedback and resets
        its trajectory planner so the next :meth:`go_to` defines a fresh motion.
        """
        for m in self.motors.values():
            # motor will still point to last target_position
            m.target_position = m.angle[0]
            self.motors_trajectories[m.motor_name].reset()

    def capture_position_sync(self, is_calibrated_angle: bool = False):
        """Read ``mech_pos`` from each motor via synchronous parameter read.

        Args:
            is_calibrated_angle (bool): If True, convert raw mechanical position to
                calibrated angle; if False, return raw ``mech_pos``.

        Returns:
            dict[str, float]: Motor name → position (rad).
        """
        positions = {}
        for motor_name in self.motors.keys():
            if is_calibrated_angle:
                positions[motor_name] = self.motors[motor_name].convert_angle_to_calibrated_angle(self.motors_manager.read_param_sync(motor_name, "mech_pos"))
            else:
                positions[motor_name] = self.motors_manager.read_param_sync(motor_name, "mech_pos")
        return positions

    def capture_position(self, is_calibrated_angle: bool = False):
        """Return cached angles from last feedback (no extra CAN read).

        Args:
            is_calibrated_angle (bool): If True, use ``calibrated_angle``; else raw ``angle``.

        Returns:
            dict[str, float]: Motor name → angle (rad).
        """
        positions = {}
        for m in self.motors.values():
            if is_calibrated_angle:
                positions[m.motor_name] = m.calibrated_angle[0]
            else:
                positions[m.motor_name] = m.angle[0]
        return positions