import time, os
from .motors.motors_manager import MotorsManager
from .steppable_system import SteppableSystem
from .navigation_platform import ThreeWheelServeDrivePlatform
from .manipulation_platform import BimanualManipulationPlatform
from .battery_system import DualBatterySystem
from .robstride_neck_platform import RobstrideNeckPlatform
from .torso_platform import EZMotionTorsoPlatform
from .teleop_manager import TeleopManager
from .test_system import TestSystem
from feathersdk.comms import CommsManager
from feathersdk.comms.system import get_all_physical_can_interfaces
import threading
import atexit, signal, sys
from .instance_lock import get_robot_instance_lock, release_robot_instance_lock
from feathersdk.utils.common import timediff, currtime, dont_allow_on_fork, FeatherThread, exponential_timeout
from feathersdk.utils.logger import warning, error, info
from feathersdk.robot.instance_lock import check_still_own_lock

import os

class Robot(SteppableSystem):
    """Base robot class for managing multiple systems and control loops."""

    def __init__(self, robot_name: str):
        get_robot_instance_lock(override_lock=False, strict=False)
        self.is_forked = False
        
        self.robot_name = robot_name
        # 10000 Hz
        self.operating_frequency = 10000
        self.systems = {}
        self.should_continue = True
        self._aborted = False  # Track if on_abort has been called to prevent double calls
        self.loop_started = False
        atexit.register(self.on_abort)
        if self.register_signal_handlers:
            signal.signal(signal.SIGTERM, self._handle_exit)
            signal.signal(signal.SIGINT, self._handle_exit)
        os.register_at_fork(before=lambda: info(f"Robot before fork: {self.robot_name}, {os.getpid()}, {threading.get_ident()}"), after_in_child=self._on_fork)

        warning(f"Robot initialized: {self.robot_name}, {os.getpid()}, {threading.get_ident()}")

    @dont_allow_on_fork
    def start(self):
        """Start the robot and all registered systems."""
        super().start(is_root_step=True)
        for system in self.systems.values():
            if system:
                system.start()
        self.loop_started = True

    def _handle_exit(self, signum, frame):
        """Handle system exit signals.
        
        Args:
            signum: Signal number
            frame: Current stack frame
        """
        self.on_abort(signal=signum, frame=frame)

    @dont_allow_on_fork
    def on_abort(self, signal=None, frame=None):
        """Handle system abort and stop all systems."""
        # Prevent double calls (e.g., from both signal handler and atexit)
        # Note: Signal handlers and atexit both run in main thread, so no threading issues
        if self._aborted:
            warning("Robot abort already called")
            return
        self._aborted = True
        
        self.should_continue = False
        warning(f"MAIN ABORT STARTED: {self.robot_name}, {os.getpid()}, {threading.get_ident()}, signal: {signal}")

        for system_name, system in self.systems.items():
            if system:
                try:
                    warning(f"Robot abort: calling on_abort for {system_name}")
                    system.on_abort()
                except Exception as e:
                    error(f"Uncaught Error in base robot on_abort for {system}: {e}")
        
        try:
            warning("Robot abort: calling on_abort for motors manager")
            self.motors_manager.on_abort()
        except Exception as e:
            error(f"Uncaught Error in motors manager on_abort: {e}")
        
        try:
            warning("Robot abort: calling on_abort for comms manager")
            self.comms.on_abort()
        except Exception as e:
            error(f"Uncaught Error in comms manager on_abort: {e}")

        try:
            warning("Robot abort: joinging with robot loop thread")
            if self.robot_loop_thread is not None:
                self.robot_loop_thread.join(timeout=0.3)
                if self.robot_loop_thread.is_alive():
                    raise RuntimeError(f"Robot loop thread did not exit in time")
        except Exception as e:
            error(f"Uncaught Error in robot loop thread join: {e}")

        warning("Robot abort: releasing instance lock.")
        try:
            release_robot_instance_lock(allow_no_lock=True)
        except Exception as e:
            error(f"Uncaught Error releasing robot instance lock: {e}")
        info(f"FeatherRobot final abort: {self.robot_name}, {os.getpid()}, {threading.get_ident()}")

        n = 1
        for i in range(n):
            if len(threading.enumerate()) <= 1:
                break

            frames = sys._current_frames()
            for t in threading.enumerate():
                frames.get(t.ident)
                import traceback
                tb = traceback.format_stack(frame) if frame is not None else "unknown_traceback"
                tn = t._target.__name__ if hasattr(t, "_target") and t._target is not None else "unknown_Target"
                info(f"sleep # {i} - Thread: {t.name}, Ident: {t.ident}, Alive: {t.is_alive()}, Daemon: {t.daemon}, target: {tn}, traceback: \n{tb}")
            if n > 1:
                time.sleep(1.0)
        
        info(f"Exiting robot abort after checking for threads")

    @dont_allow_on_fork
    def robot_loop(self):
        """Run the main robot control loop until aborted."""
        try:
            self.start()
            cnt = 0
            while self.should_continue:
                cnt += 1
                if cnt % self.operating_frequency == 0:
                    if not check_still_own_lock():
                        error("Robot lock not owned by this process. Exiting robot loop.")
                        break
                self.step()
        except Exception as e:
            error(f"Error in robot loop: {e}")
        finally:
            info(f"Final robot loop exiting")

    @dont_allow_on_fork
    def start_loop_in_background(self):
        """Start the robot control loop in a background thread."""
        self.robot_loop_thread = FeatherThread(target=self.robot_loop, daemon=True)
        self.robot_loop_thread.start()

        try:
            exponential_timeout(0.4, lambda: True if self.loop_started else None)
            info(f"Robot loop started")
        except Exception as e:
            error(f"Error in start_loop_in_background: {e}")
            raise


class FeatherRobot(Robot):
    """Feather robot implementation with navigation, manipulation, torso, and neck platforms."""

    def __init__(self, robot_name: str, config: dict = {}):
        """Initialize Feather robot with configured platforms.
        
        Initializes communication manager, battery system, motors manager, and all
        configured platforms (navigation, manipulation, torso, neck). Automatically
        registers systems with power event listeners if battery system is configured.
        
        Args:
            robot_name (str): Name identifier for the robot
            config (dict): Configuration dictionary with optional keys for each platform
        Raises:
            Exception: If E-stop is pressed or initialization fails (unless testing=True)
        """
        try:
            self.cfg = config
            self.register_signal_handlers = config.get("register_signal_handlers", False)
            
            super().__init__(robot_name)
            
            self.is_forked = False
            self.robot_loop_thread = None

            # Initialize the comms manager and add the endpoints for all physical CAN interfaces (needed to find motors)
            self.comms = CommsManager()
            self.comms.add_endpoint(get_all_physical_can_interfaces(), enable_cans=False, allow_no_enable_can=True)

            self.torso = None
            self.navigation_platform = None

            if "power" in self.cfg:
                self.power = DualBatterySystem(self.cfg["power"])
                if self.power.is_estop_pressed():
                    raise Exception("E-stop is pressed.")
            else:
                self.power = None

            self.motors_manager = MotorsManager()

            if "neck" in self.cfg:
                self.neck = RobstrideNeckPlatform(self.motors_manager, self.power, self.cfg["neck"])
            else:
                self.neck = None

            try:
                if "navigation_platform" in self.cfg:
                    self.navigation_platform = ThreeWheelServeDrivePlatform(self.motors_manager, self.cfg["navigation_platform"], self.power)
                    # Find the base motor can interface to use for EZMotion motors. Assume same for all base motors 
                    base_can_iface = self.motors_manager.motors["BwC"].iface
                
                    # Initialize the torso platform, bug requires base can interface.
                    if "torso" in self.cfg:
                        self.torso = EZMotionTorsoPlatform(self.motors_manager, self.power, base_can_iface, self.cfg["torso"])
            except Exception:
                if not config.get("testing", False):
                    raise

            if "manipulation_platform" in self.cfg and "manipulator" in self.cfg:
                self.manipulation_platform = BimanualManipulationPlatform(self.motors_manager, config=self.cfg["manipulation_platform"], power=self.power, manipulator_config=self.cfg["manipulator"])
            elif "manipulation_platform" in self.cfg and "manipulator" not in self.cfg:
                self.manipulation_platform = BimanualManipulationPlatform(self.motors_manager, config=self.cfg["manipulation_platform"], power=self.power)
            else:
                self.manipulation_platform = None

            self.vision_system = None
            self.audio_system = None
            self.listening_system = None
            self.visual_system = None

            if "teleop" in self.cfg:
                self.teleop = TeleopManager(
                    neck=self.neck,
                    torso=self.torso,
                    nav_planner=self.navigation_platform.nav_planner if self.navigation_platform else None,
                    manipulation_platform=self.manipulation_platform,
                    motors_manager=self.motors_manager,
                    config=self.cfg["teleop"],
                )
            else:
                self.teleop = None

            if config.get("testing", False):
                pass
                #self.test_system = TestSystem(self.motors_manager)

            # Initialize monitoring platform if configured
            if "monitoring" in self.cfg:
                from .monitoring import MonitoringPlatform, StepTimingCollector
                self.monitoring = MonitoringPlatform(self.motors_manager, self.cfg["monitoring"])
                self.step_timing_interval = self.cfg["monitoring"].get("step_timing_interval", 1000)
                self._step_timing_collector = self.monitoring.get_collector(StepTimingCollector)
            else:
                self.monitoring = None
                self._step_timing_collector = None

            self.systems = {
                "navigation_platform": self.navigation_platform,
                "manipulation_platform": self.manipulation_platform,
                "power": self.power,
                "torso": self.torso,
                "vision_system": self.vision_system,
                "audio_system": self.audio_system,
                "listening_system": self.listening_system,
                "visual_system": self.visual_system,
                "neck": self.neck,
                "monitoring": self.monitoring,
            }

            # Automatically register systems with on_power_event method to BatterySystem
            if self.power:
                for system_name, system in self.systems.items():
                    # Skip the power system itself to avoid self-registration
                    if system_name != "power" and system and hasattr(system, 'on_power_event') and callable(getattr(system, 'on_power_event')):
                        self.power.add_power_event_listener(system)
            
            # Reset the comms manager and add the endpoints for only the motors/ifaces we will use
            feather_ifaces = set([m.iface for m in self.motors_manager.motors.values()])
            to_remove_ifaces = set(self.comms.endpoints.keys()) - feather_ifaces
            if len(to_remove_ifaces) > 0:
                self.comms.remove_endpoint(list(to_remove_ifaces))
        except:
            import traceback
            error(f"Error in FeatherRobot init: {traceback.format_exc()}")
            self.on_abort()
            raise
    
    @dont_allow_on_fork
    def clear_motor_errors(self) -> bool:
        """Clear all motor errors. Returns True if any errors were cleared, False otherwise."""
        cleared_error = False

        for motor in self.motors_manager.motors.values():
            if len(motor.errors) > 0:
                motor.disable(clear_faults=True)
                motor.enable()
                cleared_error = True
        if self.torso is not None and self.torso.has_error():
            self.torso.clear_errors()
            cleared_error = True
        return cleared_error
    
    def _on_fork(self):
        info(f"FeatherRobot on_fork: {self.robot_name}, {os.getpid()}, {threading.get_ident()}")
        self.is_forked = True
        self.motors_manager.on_fork()
        self.comms.on_fork()
        self.navigation_platform.on_fork()
        self.torso.on_fork()

        self.navigation_platform = None
        self.manipulation_platform = None
        self.power = None
        self.torso = None
        self.vision_system = None
        self.audio_system = None
        self.listening_system = None
        self.visual_system = None
        self.neck = None
        self.monitoring = None
        self.teleop = None

        self.systems = {
            "navigation_platform": self.navigation_platform,
            "manipulation_platform": self.manipulation_platform,
            "power": self.power,
            "torso": self.torso,
            "vision_system": self.vision_system,
            "audio_system": self.audio_system,
            "listening_system": self.listening_system,
            "visual_system": self.visual_system,
            "neck": self.neck,
            "monitoring": self.monitoring,
        }

    @dont_allow_on_fork
    def get_battery_voltage(self):
        """Get battery voltage.
        
        Returns:
            float: Battery voltage in volts (currently returns 0)
        """
        return 0

    @dont_allow_on_fork
    def _step(self):
        """Private method for robot step (no-op for FeatherRobot)."""
        pass

    @dont_allow_on_fork
    def _after_step(self):
        """Step each system with per-platform timing; record root step timing when monitoring enabled."""
        current_time = currtime()
        for system_name, system in self.systems.items():
            if system and system.get_next_step_start_time() <= current_time:
                try:
                    system.step()
                except Exception as e:
                    error(
                        f"Uncaught Error in step started at {system.current_step_start_time} for {system_name}: {e}, ignoring in main loop so that other systems keep working."
                    )
                    # Call _after_step to update the system's current_step_end_time and counter as the exceptions likely happened during the _step() call
                    system._after_step()
                finally:
                    if self._step_timing_collector is not None:
                        self._step_timing_collector.record_system_step_timing(
                            system_name,
                            system.current_step_end_time - system.current_step_start_time,
                            system.current_step_end_time,
                            step_start_delay=system.current_step_start_delay,
                        )
        # Record root step duration after all subsystem stepping so it reflects full cycle time
        # Intentionally do this before calling super()._after_step() to ensure the root step duration does not include the sleep time waiting for the next step.
        end_time = currtime()
        step_duration = end_time - self.current_step_start_time
        if self._step_timing_collector is not None:
            if self.current_step % self.step_timing_interval == 0:
                self._step_timing_collector.record_step_timing(
                    step_duration,
                    end_time,
                    step_start_delay=self.current_step_start_delay,
                )
        super()._after_step()

    @dont_allow_on_fork
    def on_abort(self, signal=None, frame=None):
        """Handle system abort and stop all systems."""
        try:
            super().on_abort(signal=signal, frame=frame)
        except Exception as e:
            error(f"Uncaught Error in base robot on_abort: {e}")


