import math
from typing import List, Optional
import time
from .steppable_system import SteppableSystem
from .motors.robstride import RobstrideRunMode
from .motors.motors_manager import (
    MotorsManager, MotorMap, CalibrationWarnings, CalibrationState, MotorCalibrationError
)

MOTORS_TO_MOVE_SLOW = set[str](["EpR", "EpL", "SpR", "SpL", "SrR", "SrL"])

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.disable_arms()")
    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):
    def __init__(self, motors_manager: MotorsManager, motors_map: MotorMap, name: str, calibration_order: List[str]):
        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.name = name
        self.calibration_order = calibration_order
        self.calibration_state = CalibrationState.UNINITIALIZED
        self.calibration_warnings: dict[str, CalibrationWarnings] = {}
        # 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:
            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]:
        """Enable position hold on specified motors at their current positions.
        
        Args:
            motor_names: List of motor names to hold in position
            verbose: Enable debug output
            
        Returns:
            Dict mapping motor name to its hold position (for reference)
            Dict mapping motor name to its reset position (for release of position 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:
        """Release position hold on specified motors by disabling them.
        
        Args:
            calibrating_motor_name: Name of the motor that is being calibrated
            verbose: Enable debug output
        """
        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 current state of all motors in the arm.
        
        Returns:
            Dict mapping motor name to its state dict containing:
                - 'angle': current angle in radians
                - 'velocity': current velocity in radians/sec
        """
        state = {}
        for motor_name, motor in self.motors.items():
            state[motor_name] = {
                "angle": motor.calibrated_angle[0],
                "velocity": motor.velocity[0],
            }
        return state

    def _step(self):
        pass

