import math
from typing import Dict
from .steppable_system import SteppableSystem

from .neck_platform import NeckPlatform
from .motors.robstride import RobstrideMotor, RobstrideRunMode, RobstrideMotorMode
from .motors.motors_manager import MotorsManager
from .motors.robstride.robstride_motor import TimestampedValue
import time
from feathersdk.comms.comms_manager import UnknownInterfaceError
from .battery_system import BatterySystem
from .battery_system import PowerEvent
from feathersdk.utils.common import currtime, timediff
from feathersdk.utils.logger import error


MAX_POSITION_VELOCITY = 20  # 0 - 20 rad/s (default is 10 rad/s)
MAX_POSITION_ACCELERATION = 30  # 0 - 1000 rad/s^2 (default is 10 rad/s^2) (50 stays within rated torque limits)


class RobstrideNeckPlatform(NeckPlatform):
    """Platform for controlling Robstride neck motors (pitch and yaw).
    
    This class manages the neck platform with two motors: pitch (NpC) and yaw (NwC).
    It handles initialization, calibration, movement control, and power management.
    The platform supports position control with configurable velocity and acceleration limits.
    """
    YAW_MOTOR_NAME = "NwC"
    PITCH_MOTOR_NAME = "NpC"

    def __init__(self, motors_manager: MotorsManager, power: BatterySystem = None, config: dict = {}):
        """Initialize the Robstride neck platform.
        
        Args:
            motors_manager (MotorsManager): The motors manager instance for controlling motors.
            power (BatterySystem, optional): Battery system for power event handling. If provided,
                the platform will check if recalibration is needed after power cycles.
            config (dict): Configuration dictionary containing motor settings. Must include
                a "motors" key with motor configurations including:
                - id: Motor ID in hexadecimal string format
                - motor_config: Motor configuration dictionary
                - default_velocity_max: Default maximum velocity in rad/s
                - default_acceleration_max: Default maximum acceleration in rad/s^2
                - default_loc_kp: Default location proportional gain
        """
        super().__init__()
        self.motors_manager = motors_manager
        self.power = power
        self.cfg = config
        self.enabled = False
        self.healthy_state = False
        self.recalibrating = False
        self.control_mode = RobstrideRunMode.Position
        self.operation_frequency = 50
        self.calibration_warnings = {}
        self.target_pitch_in_degrees = None
        self.target_yaw_in_degrees = None
        self.last_sent_pitch_in_degrees = None
        self.last_sent_yaw_in_degrees = None

        # vel_max and acc_set are write only parameters, so we need to store them
        self.motor_vel_max = {}
        self.motor_acc_set = {}
        self.default_loc_kp = {}
        self.default_vel_max = {}
        self.default_acc_set = {}
        self.motors : Dict[str, RobstrideMotor] = {}

        for motor_name, motor_dict in self.cfg["motors"].items():
            self.motors[motor_name] = RobstrideMotor(
                int(motor_dict["id"], 16),
                motor_name,
                motor_dict["motor_config"],
            )
            self.default_vel_max[motor_name] = motor_dict["default_velocity_max"]
            self.default_acc_set[motor_name] = motor_dict["default_acceleration_max"]
            self.default_loc_kp[motor_name] = motor_dict["default_loc_kp"]
        self.motors_manager.add_motors(
            self.motors,
            "neck",
        )

        self.motors_manager.find_motors(list(self.motors.keys()))

        for motor_name in self.motors.keys():
            self.set_movement_profile(motor_name, self.default_vel_max[motor_name], self.default_acc_set[motor_name])

        if (
            self.motors[self.PITCH_MOTOR_NAME].calibration_time is not None
            and self.motors[self.YAW_MOTOR_NAME].calibration_time is not None
            and self.motors[self.PITCH_MOTOR_NAME].iface is not None
            and self.motors[self.YAW_MOTOR_NAME].iface is not None
        ):
            if self.power is not None:
                last_powered_up_time = self.power.last_powered_up_time()
                if last_powered_up_time > self.motors[self.PITCH_MOTOR_NAME].calibration_time or last_powered_up_time > self.motors[self.YAW_MOTOR_NAME].calibration_time:
                    self.healthy_state = False
                    error("Neck needs to be recalibrated")
                else:
                    for motor_name in self.motors.keys():
                        self.motors_manager.recover_from_power_cycle(motor_name)
                    self.healthy_state = self.motors_manager.check_motors_in_range([self.PITCH_MOTOR_NAME, self.YAW_MOTOR_NAME], raise_error=False)
                    self._wakeup()
            

    def look_up(self):
        """Move the neck to look straight forward.
        
        Sets pitch to -90 degrees and yaw to 0 degrees (center position).
        This is typically used as a safe homing position.
        """
        self.last_sent_pitch_in_degrees = None
        self.last_sent_yaw_in_degrees = None
        self.move(-90, 0)

    def _wakeup(self):
        """Private method to wake up the neck motors and prepare them for operation.
        
        Raises:
            UnknownInterfaceError: If communication interface is not available, platform is disabled.
            Exception: If wakeup fails for any other reason, platform is disabled.
        """
        print("Waking up neck")
        try:
            self.motors_manager.write_param(self.YAW_MOTOR_NAME, "loc_kp", self.default_loc_kp[self.YAW_MOTOR_NAME])
            self.motors_manager.write_param(self.PITCH_MOTOR_NAME, "loc_kp", self.default_loc_kp[self.PITCH_MOTOR_NAME])
            self.enable_motors()
            #wait for motors to enable
            for i in range(1000):
                if self.motors_manager.motors_in_mode(self.motors.keys(), RobstrideMotorMode.Run):
                    break
                time.sleep(0.001)
                if i % 10 == 0:
                    self.enable_motors()

            self.look_up()
        except UnknownInterfaceError as e:
            print(f"CAN {e}, neck platform will be disabled")
            self.enabled = False
        except Exception as e:
            print(f"Error: Neck wakeup failed: {e}")
            self.enabled = False

    def enable_motors(self):
        """Enable all neck motors in position control mode.
        """
        for motor_name in self.motors:
            self.motors_manager.set_run_mode(motor_name, RobstrideRunMode.Position)
            self.motors_manager.enable(motor_name)
        self.enabled = True

    def sleep_and_disable(self):
        """Put motors to sleep position and disable them safely.
        
        Raises:
            Exception: Re-raises any exceptions that occur during the sleep process,
                but ensures motors are disabled in the finally block.
        """
        # recalibrate with default vel_max and acc_set
        try:
            in_range = self.motors_manager.check_motors_in_range(self.motors.keys(), raise_error=False)
            if self.healthy_state and not self.recalibrating and self.enabled and in_range:

                #set back to default vel_max and acc_set
                for motor_name in self.motors.keys():
                    self.motors_manager.write_param(motor_name, "vel_max", self.default_vel_max[motor_name])
                    self.motors_manager.write_param(motor_name, "acc_set", self.default_acc_set[motor_name])

                # Move both motors to homing position (0°) with safety margin for pitch
                self.motors_manager.set_target_position(
                    self.PITCH_MOTOR_NAME, self.motors[self.PITCH_MOTOR_NAME].calibration_homing_pos
                )
                self.motors_manager.set_target_position(
                    self.YAW_MOTOR_NAME, self.motors[self.YAW_MOTOR_NAME].calibration_homing_pos
                )
                time.sleep(1.0)
                # set back to what the user set
                for motor_name in self.motors.keys():
                    self.set_movement_profile(
                        motor_name,
                        self.motor_vel_max[motor_name],
                        self.motor_acc_set[motor_name],
                    )
                self.target_pitch_in_degrees = None
                self.target_yaw_in_degrees = None
                self.last_sent_pitch_in_degrees = None
                self.last_sent_yaw_in_degrees = None
            else:
                print("Motor does not meet sleep and disable conditions, disabling motor")
        except Exception as e:
            #called on on abort, not best to raise an error here
            print(f"Unexpected Error: Neck sleep and disable failed: {e}")
            raise
        finally:
            self.disable_motors()
            self.enabled = False

    def disable_motors(self):
        """Disable all neck motors.
        """
        for motor_name in self.motors:
            self.motors_manager.disable(motor_name)
        self.enabled = False

    def get_system_state(self):
        """Get the current system state of all neck motors.
        
        Returns:
            Dict[str, float]: Dictionary mapping motor names to their calibrated angles.
                Keys are in the format "{motor_name}_angle" (e.g., "NpC_angle", "NwC_angle").
                Values are the calibrated angle in radians.
        """
        state = {}
        for motor_name in self.motors:
            motor = self.motors_manager.get_motor(motor_name)
            state[f"{motor_name}_angle"] = motor.calibrated_angle.value
        return state

    def _step(self):
        """Private method to execute pending movement commands at operation frequency.
        """
        # We might have to move this into a background thread because _move is blocking 0.001s to execute
        # which means it delays the step loops of all systems by 0.05s assuming a move command is always
        # being sent.        
        self._move(self.target_pitch_in_degrees, self.target_yaw_in_degrees)

    def set_movement_profile(
        self,
        motor_name: str,
        max_velocity: float,
        max_acceleration: float,
        max_jerk=None,
    ):
        """Set the movement profile parameters for a specific motor.
        
        Args:
            motor_name (str): Name of the motor to configure (e.g., "NpC" or "NwC").
            max_velocity (float): Maximum velocity in rad/s. Must be between 0.1 and
                MAX_POSITION_VELOCITY (20 rad/s).
            max_acceleration (float): Maximum acceleration in rad/s^2. Must be between 0.1 and
                MAX_POSITION_ACCELERATION (30 rad/s^2).
            max_jerk (float, optional): Maximum jerk (not supported, must be None).
        
        Raises:
            ValueError: If max_velocity or max_acceleration are out of valid range,
                if max_jerk is not None, or if motor_name is not a valid neck motor.
        """
        # raise an error it not in range [0, MAX_POS_VELOCITY] or [0, MAX_POS_ACCELERATION]
        if max_velocity > MAX_POSITION_VELOCITY or max_velocity < 0.1:
            raise ValueError(f"Error: Max velocity must be between 0.1 and {MAX_POSITION_VELOCITY} rad/s")
        if max_acceleration > MAX_POSITION_ACCELERATION or max_acceleration < 0.1:
            raise ValueError(f"Error: Max acceleration must be between 0.1 and {MAX_POSITION_ACCELERATION} rad/s^2")
        if max_jerk is not None:
            raise ValueError(f"Max Jerk is not supported")

        # Update the stored values for this motor
        if motor_name in self.motors:
            self.motor_vel_max[motor_name] = max_velocity
            self.motor_acc_set[motor_name] = max_acceleration
        else:
            raise ValueError(f"Motor {motor_name} is not a valid neck motor. Valid motors: {list(self.motors.keys())}")

        # Write the parameters to the motor
        self.motors_manager.write_param(motor_name, "vel_max", max_velocity)
        self.motors_manager.write_param(motor_name, "acc_set", max_acceleration)

    def _move(self, pitch_in_degrees: float = None, yaw_in_degrees: float = None, verbose: bool = False):
        """Private method to move neck motors to specified positions.
        
        Args:
            pitch_in_degrees (float, optional): Target pitch angle in degrees. None to keep current target.
                Valid range: -120 to 0 degrees.
            yaw_in_degrees (float, optional): Target yaw angle in degrees. None to keep current target.
                Valid range: -90 to 90 degrees.
            verbose (bool): If True, waits for movement to complete and prints detailed
                movement information including torque, angle, and velocity feedback.
        
        Raises:
            Exception: If movement fails for any reason.
        """
        try:
            start_time = currtime()
            pitch_motor = self.motors[self.PITCH_MOTOR_NAME]
            yaw_motor = self.motors[self.YAW_MOTOR_NAME]
            
            # Calculate target positions using unified pattern
            target_pitch_pos = None
            target_yaw_pos = None
            
            if pitch_in_degrees is not None:
                pitch_in_radians = (pitch_in_degrees * math.pi / 180)
                target_pitch_pos = pitch_motor.calibration_homing_pos + pitch_in_radians
                if verbose:
                    print("trying to go to pitch: ", target_pitch_pos)

                if self.last_sent_pitch_in_degrees != pitch_in_degrees:
                    self.motors_manager.set_target_position(self.PITCH_MOTOR_NAME, target_pitch_pos)
                    self.last_sent_pitch_in_degrees = pitch_in_degrees

            if yaw_in_degrees is not None:
                yaw_in_radians = yaw_in_degrees * math.pi / 180
                target_yaw_pos = yaw_motor.calibration_homing_pos + yaw_in_radians
                if verbose:
                    print("trying to go to yaw: ", target_yaw_pos)
                
                if self.last_sent_yaw_in_degrees != yaw_in_degrees:
                    self.motors_manager.set_target_position(self.YAW_MOTOR_NAME, target_yaw_pos)
                    self.last_sent_yaw_in_degrees = yaw_in_degrees

            # Wait for movement to finish (only in verbose mode)
            while True and verbose:
                self.motors_manager.enable(self.PITCH_MOTOR_NAME)
                self.motors_manager.enable(self.YAW_MOTOR_NAME)
                time.sleep(0.05)
                curr_pitch_pos = pitch_motor.angle[0]
                curr_yaw_pos = yaw_motor.angle[0]
                print(
                    f"Pitch torque: {pitch_motor.torque[0]}, Pitch angle: {pitch_motor.angle[0]}, velocity: {pitch_motor.velocity[0]}\n"
                    f"Yaw torque: {yaw_motor.torque[0]}, Yaw angle: {yaw_motor.angle[0]}, velocity: {yaw_motor.velocity[0]}"
                )
                if (
                    (target_pitch_pos is None or abs(target_pitch_pos - curr_pitch_pos) < 0.02)
                    and (target_yaw_pos is None or abs(target_yaw_pos - curr_yaw_pos) < 0.02)
                ):
                    print(f"Final Pitch Position: {curr_pitch_pos}, Final Yaw Position: {curr_yaw_pos}")
                    print(f"Time taken to complete: {timediff(start_time)}")
                    break
        except Exception as e:
            raise Exception(f"Error: Neck movement failed: {e}")

    def move(self, pitch_in_degrees: float = None, yaw_in_degrees: float = None):
        """Set target position for neck movement.
        
        Args:
            pitch_in_degrees (float, optional): Target pitch angle in degrees. None to keep current target.
                Must be between -120 and 0 degrees.
            yaw_in_degrees (float, optional): Target yaw angle in degrees. None to keep current target.
                Must be between -90 and 90 degrees.
        
        Raises:
            Exception: If neck is not enabled, not healthy, or currently recalibrating.
            ValueError: If pitch or yaw values are outside their valid ranges.
        """
        # Step will read the target_pitch and yaw at operation frequency and call the _move method
        # to prevent overloading the CAN bus with too many commands.
        if not self.enabled:
            raise Exception("Error: Neck must be enabled. Enable with enable_motors() method.")
        if not self.healthy_state:
            raise Exception("Error: Neck must be recalibrated. Recalibrate with recalibrate() method.")
        if self.recalibrating:
            raise Exception("Error: Neck is recalibrating, please wait for calibration to complete.")
        # if pitch in degrees not in range [-120, 0], raise error
        if pitch_in_degrees is not None and (pitch_in_degrees < -120 or pitch_in_degrees > 0):
            raise ValueError(f"Pitch in degrees must be between -120 and 0. Got {pitch_in_degrees}")
        # if yaw in degrees not in range [-90, 90], raise error
        if yaw_in_degrees is not None and (yaw_in_degrees < -90 or yaw_in_degrees > 90):
            raise ValueError(f"Yaw in degrees must be between -90 and 90. Got {yaw_in_degrees}")
        self.target_pitch_in_degrees = pitch_in_degrees
        self.target_yaw_in_degrees = yaw_in_degrees

    def recalibrate(self, verbose: bool = False):
        """Recalibrate neck motors.
        
        Args:
            verbose (bool): If True, prints detailed calibration progress information.
        
        Raises:
            Exception: If neck is already recalibrating, or if calibration fails.
                On failure, healthy_state is set to False and enabled is set to True.
        """
        if self.recalibrating:
            raise Exception("Error: Neck is already recalibrating, please wait for calibration to complete.")
        
        try:
            self.recalibrating = True
            for motor_name in self.motors.keys():
                print(f"Recalibrating {motor_name}")
                # Set default vel_max and acc_set for calibration
                self.motors_manager.write_param(motor_name, "vel_max", self.default_vel_max[motor_name])
                self.motors_manager.write_param(motor_name, "acc_set", self.default_acc_set[motor_name])
                
                # calibrate_motor handles: find range, zero at homing, refine, save calibration
                warnings = self.motors_manager._calibrate_motor(
                    motor_name,
                    torque_threshold_for_limit_detection=1.5,
                    step_time=0.05,
                    verbose=verbose,
                )   
                self.calibration_warnings[motor_name] = warnings
                time.sleep(0.02)

            # set back to what the user set
            self.set_movement_profile(
                self.PITCH_MOTOR_NAME,
                self.motor_vel_max[self.PITCH_MOTOR_NAME],
                self.motor_acc_set[self.PITCH_MOTOR_NAME],
            )
            self.set_movement_profile(
                self.YAW_MOTOR_NAME, self.motor_vel_max[self.YAW_MOTOR_NAME], self.motor_acc_set[self.YAW_MOTOR_NAME]
            )
        except Exception as e:
            self.recalibrating = False
            self.enabled = True
            self.healthy_state = False
            raise Exception(f"Error: Neck recalibration failed: {e}")

        self.healthy_state = True
        self.recalibrating = False
        self.enabled = True
        # After calibration, motors are already at homing position (0°)
        self.target_pitch_in_degrees = 0
        self.target_yaw_in_degrees = 0
        self.last_sent_pitch_in_degrees = None
        self.last_sent_yaw_in_degrees = None

    async def get_position(self, force_update: bool = False) -> Dict[str, float]:
        """Get the current mechanical positions of all neck motors.
        
        Args:
            force_update (bool): Currently unused, reserved for future implementation.
        
        Returns:
            Dict[str, float]: Dictionary mapping motor names to their mechanical positions in radians.
                Keys are motor names (e.g., "NpC", "NwC").
        """
        # to get angle asyncronously from feedback frame need to ping motor for (feedback frame (not async))
        # or enable active reporting (continous polling), so using async mech_pos instead
        positions = {}
        for motor_name in self.motors:
            if motor_name == self.PITCH_MOTOR_NAME:
                result = await self.motors_manager.read_param_async(motor_name, "mech_pos")
            else:
                result = await self.motors_manager.read_param_async(motor_name, "mech_pos")
            positions[motor_name] = result
        return positions

    async def get_state(self, motor_name: str) -> Dict[str, TimestampedValue]:
        """Get the current state of a specific motor.
        
        Args:
            motor_name (str): Name of the motor to query (e.g., "NpC" or "NwC").
        
        Returns:
            Dict[str, TimestampedValue]: Dictionary containing motor state with keys:
                - "temp": Temperature (TimestampedValue)
                - "torque": Torque (TimestampedValue)
                - "angle": Angle (TimestampedValue)
                - "velocity": Velocity (TimestampedValue)
        """
        # get temp, torque, angle, velocity from motor feedback frame
        state = {}
        await self.motors_manager.read_current_state_async(motor_name)
        motor = self.motors_manager.motors[motor_name]
        state["temp"] = motor.temp
        state["torque"] = motor.torque
        state["angle"] = motor.angle
        state["velocity"] = motor.velocity
        return state

    async def health_check(self) -> Dict[str, Dict[str, float]]:
        """Perform a comprehensive health check on all neck motors.
        
        Returns:
            Dict[str, Dict[str, float]]: Nested dictionary mapping motor names to their health parameters.
                Each motor's dictionary contains:
                - "mech_pos": Mechanical position in radians
                - "loc_ref": Location reference (target position) in radians
                - "velocity": Mechanical velocity in rad/s
                - "iqf": Quadrature current feedback in amperes
                - "vel_max": Maximum velocity setting in rad/s
                - "acc_set": Acceleration setting in rad/s^2
                - "loc_kp": Location proportional gain
        """
        # call motors_manager.read_param_async for each motor and return the (mech_pos, loc_ref, velocity, iqf, vel_max, acc_set, loc_kp, spd_kp, spd_ki, spd_filt_gain)
        states = {}
        for motor_name in self.motors:
            result = {
                "mech_pos": await self.motors_manager.read_param_async(motor_name, "mech_pos"),
                "loc_ref": await self.motors_manager.read_param_async(motor_name, "loc_ref"),
                "velocity": await self.motors_manager.read_param_async(motor_name, "mech_vel"),
                "iqf": await self.motors_manager.read_param_async(motor_name, "iqf"),
                "vel_max": self.motor_vel_max[motor_name],
                "acc_set": self.motor_acc_set[motor_name],
                "loc_kp": await self.motors_manager.read_param_async(motor_name, "loc_kp"),
            }
            states[motor_name] = result
        return states

    def on_abort(self):
        """Handle system abort event.
        """
        self.sleep_and_disable()

    def on_power_event(self, event: PowerEvent):
        """Handle power system events.
        
        Args:
            event (PowerEvent): The power event type (POWER_OFF or POWER_RESTORED).
        """
        if event == PowerEvent.POWER_OFF:
            self.healthy_state = False
        elif event == PowerEvent.POWER_RESTORED:
            self.healthy_state = False