from .steppable_system import SteppableSystem
from .motors.robstride import RobstrideMotor
from .motors.motors_manager import MotorsManager
from .rh56dftp_hand import create_finger_joints_from_config
from .gripper.i2rt_gripper import I2RTGripper
from .motors.damiao.damaio_motor import DamiaoMotor
from .battery_system import BatterySystem
from .arm_system import ArmSystem, _confirm_calibration
from .hand_system import HandSystem
from .gripper_system import GripperSystem
from feathersdk.utils.logger import info, warning
from feathersdk.utils.feathertypes import *

# Motor calibration order: bottom-most (wrist) to top-most (shoulder)
# Physically: wrist is near the hand (bottom), shoulder connects to torso (top)
# This ensures stable calibration as lower joints are calibrated first
RIGHT_ARM_CALIBRATION_ORDER = ["WwR", "WpR", "WrR", "EpR", "SpR", "SrR", "SwR"]
LEFT_ARM_CALIBRATION_ORDER = ["WwL", "WpL", "WrL", "EpL", "SpL", "SrL", "SwL"]

RIGHT_ARM_FAMILY_NAME = "right_arm"
LEFT_ARM_FAMILY_NAME = "left_arm"

class BimanualManipulationPlatform(SteppableSystem):
    """Coordinates left/right arms, optional hands/grippers, and shared stepping.

    ``subsystems`` lists arms first, then any initialized manipulators, for :meth:`start`
    and :meth:`_step`.
    """

    def __init__(self, motors_manager: MotorsManager, config: dict, power: BatterySystem = None, manipulator_config: dict = {}):

        self.cfg = config
        self.manipulator_cfg = manipulator_config
        self.motors_manager = motors_manager
        self.power = power
        self.operating_frequency = 300

        # Initialize arms from config or with empty motors map
        self.right_arm_motors = {}
        if self.cfg.get(RIGHT_ARM_FAMILY_NAME, False):
            for motor_name, motor_dict in self.cfg[RIGHT_ARM_FAMILY_NAME]["motors"].items():
                self.right_arm_motors[motor_name] = RobstrideMotor(
                    int(motor_dict["id"], 16),
                    motor_name,
                    motor_dict["motor_config"],
                )
            
        self.left_arm_motors = {}
        if self.cfg.get(LEFT_ARM_FAMILY_NAME, False):
            for motor_name, motor_dict in self.cfg[LEFT_ARM_FAMILY_NAME]["motors"].items():
                self.left_arm_motors[motor_name] = RobstrideMotor(
                    int(motor_dict["id"], 16),
                    motor_name,
                    motor_dict["motor_config"],
                )
                
        self.right_arm = ArmSystem(motors_manager, self.right_arm_motors, RIGHT_ARM_FAMILY_NAME, RIGHT_ARM_CALIBRATION_ORDER)
        self.left_arm = ArmSystem(motors_manager, self.left_arm_motors, LEFT_ARM_FAMILY_NAME, LEFT_ARM_CALIBRATION_ORDER)

        self.subsystems = [
            self.right_arm,
            self.left_arm,
        ]

        # Initialize hands and grippers from manipulator config
        self.right_manipulator = None
        self.left_manipulator = None
        self._has_rh56f1_hands = False  # Track if RH56F1 hands are being used

        self._initialize_manipulators()

        if self.right_manipulator:
            self.subsystems.append(self.right_manipulator)
        if self.left_manipulator:
            self.subsystems.append(self.left_manipulator)

    def on_abort(self):
        """Disable both arms and tear down RH56F1 hands if they were initialized.

        RH56F1 EtherCAT masters are released via ``destroy()`` so a new robot instance
        can claim them after abort.
        """
        self.right_arm.disable()
        self.left_arm.disable()
        # Release RH56F1 EtherCAT masters on abort so a new Robot can request them on re-init
        if self._has_rh56f1_hands:
            if self.right_manipulator.is_initialized:
                info("Releasing rh56f1 right hand (destroy)")
                self.right_manipulator.destroy()
            if self.left_manipulator.is_initialized:
                info("Releasing rh56f1 left hand (destroy)")
                self.left_manipulator.destroy()

    def start(self):
        """Start this platform and call ``start()`` on any subsystem that defines it."""
        super().start()
        for subsystem in self.subsystems:
            if hasattr(subsystem, "start"):
                subsystem.start()

    def _initialize_manipulators(self):
        """Populate ``right_manipulator`` / ``left_manipulator`` from ``manipulator_cfg``.

        Sets ``_has_rh56f1_hands`` when RH56F1 is selected (for :meth:`on_abort`).
        """
        if self.manipulator_cfg != {}:
            # Detect manipulator type from config class or by checking available sections
            manipulator_class = self.manipulator_cfg.get("class", "").lower()
            
            # Check for RH56DFTP hands
            if "rh56dftp" in manipulator_class:
                # Initialize right hand if section exists
                right_finger_joints = create_finger_joints_from_config(self.manipulator_cfg, "right")
                if right_finger_joints:
                    self.right_manipulator = HandSystem(self.motors_manager, right_finger_joints, "right_hand")
            
                # Initialize left hand if section exists
                left_finger_joints = create_finger_joints_from_config(self.manipulator_cfg, "left")
                if left_finger_joints:
                    self.left_manipulator = HandSystem(self.motors_manager, left_finger_joints, "left_hand")
        
            if "rh56f1" in manipulator_class:
                # Lazy import - only import RH56F1Hand when actually needed
                # This avoids triggering compilation if RH56F1 is not being used
                from .rh56f1_hand import RH56F1Hand
                
                self._has_rh56f1_hands = True
                # Get config for right and left hands
                right_hand_cfg = self.manipulator_cfg.get("rh56f1_hand_right", {})
                left_hand_cfg = self.manipulator_cfg.get("rh56f1_hand_left", {})
                
                self.right_manipulator = RH56F1Hand(
                    master_index=right_hand_cfg.get("master_index", 0),
                    slave_position=right_hand_cfg.get("slave_position", 0)
                )
                self.left_manipulator = RH56F1Hand(
                    master_index=left_hand_cfg.get("master_index", 1),
                    slave_position=left_hand_cfg.get("slave_position", 0)
                )
                self.right_manipulator.initialize()
                self.left_manipulator.initialize()
                self.right_manipulator.start()
                self.left_manipulator.start()

                #wait for hands to be operational
                left_operational = self.left_manipulator.wait_for_operational(timeout=5.0, check_interval=0.1)
                right_operational = self.right_manipulator.wait_for_operational(timeout=5.0, check_interval=0.1)

                if left_operational and right_operational:
                    info("Both RH56F1 hands are operational!")
                else:
                    warning("Left operational={}, Right operational={}".format(left_operational, right_operational))

            # Check for I2RT grippers
            if "linear4310" in manipulator_class:
                info("Initializing I2RT Linear 4310 grippers...")
                #pass in each arm iface to the gripper if exists, otherwise use can0
                right_can_iface = self.motors_manager.motors["SpR"].iface
                left_can_iface = self.motors_manager.motors["SpL"].iface 
                # Initialize right gripper if section exists
                right_gripper_cfg = self.manipulator_cfg.get("i2rt_gripper_right", {})
                if right_gripper_cfg and "motors" in right_gripper_cfg:
                    right_gripper_motors = {}
                    for motor_name, motor_dict in right_gripper_cfg["motors"].items():
                        if motor_dict.get("type") == "damiao":
                            motor_id = int(motor_dict["id"])
                            motor_config = motor_dict.get("motor_config", {})
                            
                            # Create DamiaoMotor instance
                            damiao_motor = DamiaoMotor(
                                motor_name=motor_name,
                                motor_id=motor_id,
                                config=motor_config,
                                iface=right_can_iface,
                            )
                            
                            # Create I2RTGripper instance
                            i2rt_gripper = I2RTGripper(motor=damiao_motor, power=self.power)
                            right_gripper_motors[motor_name] = i2rt_gripper
                    
                    if right_gripper_motors:
                        self.right_manipulator = GripperSystem(
                            motors_manager=self.motors_manager,
                            motors_map=right_gripper_motors,
                            name="right_gripper",
                        )
                
                # Initialize left gripper if section exists
                left_gripper_cfg = self.manipulator_cfg.get("i2rt_gripper_left", {})
                if left_gripper_cfg and "motors" in left_gripper_cfg:
                    left_gripper_motors = {}
                    for motor_name, motor_dict in left_gripper_cfg["motors"].items():
                        if motor_dict.get("type") == "damiao":
                            motor_id = int(motor_dict["id"])
                            motor_config = motor_dict.get("motor_config", {})
                            
                            # Create DamiaoMotor instance
                            damiao_motor = DamiaoMotor(
                                motor_name=motor_name,
                                motor_id=motor_id,
                                config=motor_config,
                                iface=left_can_iface,
                            )
                            # Create I2RTGripper instance
                            i2rt_gripper = I2RTGripper(motor=damiao_motor, power=self.power)
                            left_gripper_motors[motor_name] = i2rt_gripper
                    
                    if left_gripper_motors:
                        self.left_manipulator = GripperSystem(
                            motors_manager=self.motors_manager,
                            motors_map=left_gripper_motors,
                            name="left_gripper",
                        )

    def disable_arms(self):
        """Disable left and right arms (compliance/stop via each :meth:`ArmSystem.disable`)."""
        self.right_arm.disable()
        self.left_arm.disable()

    def enable_arms(self):
        """Enable left and right arms (operation mode, with retries per :meth:`ArmSystem.enable`)."""
        self.right_arm.enable()
        self.left_arm.enable()

    def freeze_current_arm_state(self, kp: float = 500, kd: float = 5):
        """Hold both arms at their current pose using the given operation gains.

        Args:
            kp: Position gain forwarded to each :meth:`ArmSystem.freeze_current_arm_state`.
            kd: Damping gain forwarded to each arm.
        """
        self.right_arm.freeze_current_arm_state(kp, kd)
        self.left_arm.freeze_current_arm_state(kp, kd)

    def _recalibrate(
        self,
        calibrate_right: bool = True,
        calibrate_left: bool = True,
        motor_names: Optional[List[str]] = None,
        step_time: float = 0.01,
        verbose: bool = False
    ) -> None:
        """Recalibrate arm motors.
        
        Calibrates motors from bottom-most (wrist) to top-most (shoulder) for each arm.
        
        Args:
            calibrate_right: Whether to calibrate right arm (default True)
            calibrate_left: Whether to calibrate left arm (default True)
            motor_names: Optional list of specific motors to calibrate.
                        If None, calibrates all motors in selected arms.
            step_time: Delay between steps (default 0.01s)
            verbose: Enable debug output (default False)
            
        Raises:
            Exception: If user cancels calibration
        """
        if not calibrate_right and not calibrate_left:
            print("No arms selected for calibration.")
            return
            
        if not _confirm_calibration("arms"):
            raise Exception("Calibration cancelled by user")
        
        if motor_names is not None:
            # Split motor names by arm
            right_motors = [m for m in motor_names if m in self.right_arm.motors]
            left_motors = [m for m in motor_names if m in self.left_arm.motors]            
        else:
            right_motors = self.right_arm.motors.keys()
            left_motors = self.left_arm.motors.keys()

        if calibrate_right:
            self.right_arm.recalibrate(
                motor_names=right_motors,
                step_time=step_time,
                verbose=verbose,
                _skip_confirmation=True
            )
        if calibrate_left:
            self.left_arm.recalibrate(
                motor_names=left_motors,
                step_time=step_time,
                verbose=verbose,
                _skip_confirmation=True
            )

    def is_calibrated(self, check_right: bool = True, check_left: bool = True) -> bool:
        """Check if arms are calibrated.
        
        Args:
            check_right: Check right arm calibration state
            check_left: Check left arm calibration state
            
        Returns:
            True if all specified arms are calibrated
        """
        if check_right and not self.right_arm.is_calibrated():
            return False
        if check_left and not self.left_arm.is_calibrated():
            return False
        return True

    def get_state(self) -> dict:
        """Merge :meth:`ArmSystem.get_state` for the right and left arms.

        Returns:
            dict: Motor name → per-motor snapshot (``temp``, ``torque``, ``angle``,
            ``velocity``) as returned by each arm. 
        """
        state = {}
        
        # Collect arm states
        state.update(self.right_arm.get_state())
        state.update(self.left_arm.get_state())
        
        return state

    def _step(self):
        for subsystem in self.subsystems:
            if hasattr(subsystem, "step"):
                subsystem.step()
