#!/usr/bin/env python3
"""
Run ArduPilot SITL + Gazebo Simulation.

This script orchestrates the full ArduPilot SITL + Gazebo simulation stack:
1. Launches ArduPilot SITL (sim_vehicle.py with Gazebo model)
2. Starts MAVProxy for GCS connectivity
3. Runs MAVLink bridge + controllers

Usage:
    python run_ardupilot.py                    # Start everything
    python run_ardupilot.py --pattern circular # With rover movement
    python run_ardupilot.py --no-sitl          # Just controllers (SITL already running)
"""

import argparse
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

import rclpy
from rclpy.executors import MultiThreadedExecutor

# Check for pymavlink
try:
    from pymavlink import mavutil
    MAVLINK_AVAILABLE = True
except ImportError:
    MAVLINK_AVAILABLE = False

from mavlink_bridge import MAVLinkBridge
from drone_controller import DroneController
from rover_controller import RoverController, MovementPattern


def find_sim_vehicle() -> str:
    """Find sim_vehicle.py in common ArduPilot locations."""
    locations = [
        os.path.expanduser("~/ardupilot/Tools/autotest/sim_vehicle.py"),
        "/opt/ardupilot/Tools/autotest/sim_vehicle.py",
        os.environ.get("ARDUPILOT_HOME", "") + "/Tools/autotest/sim_vehicle.py",
    ]
    
    for loc in locations:
        if os.path.exists(loc):
            return loc
    
    # Try finding it in PATH
    import shutil
    if shutil.which("sim_vehicle.py"):
        return "sim_vehicle.py"
    
    return None


def parse_args():
    parser = argparse.ArgumentParser(
        description='Run ArduPilot SITL + Gazebo + MAVLink Bridge',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run_ardupilot.py                       # Full stack with Gazebo
  python run_ardupilot.py --pattern circular    # Moving rover
  python run_ardupilot.py --no-sitl             # SITL already running
  python run_ardupilot.py --mavproxy-port 14551 # Custom port
        """
    )
    
    # SITL options
    parser.add_argument(
        '--no-sitl', action='store_true',
        help='Skip launching SITL (assume already running)'
    )
    parser.add_argument(
        '--sitl-path', type=str,
        help='Path to sim_vehicle.py'
    )
    parser.add_argument(
        '--vehicle', type=str, default='ArduCopter',
        choices=['ArduCopter', 'ArduPlane', 'APMrover2'],
        help='Vehicle type (default: ArduCopter)'
    )
    parser.add_argument(
        '--frame', type=str, default='gazebo-iris',
        help='Vehicle frame (default: gazebo-iris)'
    )
    
    # MAVProxy settings
    parser.add_argument(
        '--mavproxy-port', type=int, default=14550,
        help='MAVProxy output port (default: 14550)'
    )
    parser.add_argument(
        '--sitl-port', type=int, default=5760,
        help='SITL connection port (default: 5760)'
    )
    
    # Rover options
    parser.add_argument(
        '--pattern', type=str, default='stationary',
        choices=['stationary', 'linear', 'circular', 'random', 'square'],
        help='Rover movement pattern (default: stationary)'
    )
    parser.add_argument(
        '--speed', '-s', type=float, default=0.5,
        help='Rover speed in m/s (default: 0.5)'
    )
    parser.add_argument(
        '--amplitude', '-a', type=float, default=2.0,
        help='Rover movement amplitude in meters (default: 2.0)'
    )
    parser.add_argument(
        '--no-rover', action='store_true',
        help='Disable rover controller'
    )
    
    # Other
    parser.add_argument(
        '--console', action='store_true',
        help='Show MAVProxy console'
    )
    parser.add_argument(
        '--map', action='store_true',
        help='Show MAVProxy map'
    )
    
    args, _ = parser.parse_known_args()
    return args


class SITLProcess:
    """Manages the ArduPilot SITL process."""
    
    def __init__(
        self,
        sim_vehicle_path: str,
        vehicle: str = 'ArduCopter',
        frame: str = 'gazebo-iris',
        sitl_port: int = 5760,
        mavproxy_port: int = 14550,
        console: bool = False,
        show_map: bool = False
    ):
        self.sim_vehicle_path = sim_vehicle_path
        self.vehicle = vehicle
        self.frame = frame
        self.sitl_port = sitl_port
        self.mavproxy_port = mavproxy_port
        self.console = console
        self.show_map = show_map
        self.process = None
    
    def start(self) -> bool:
        """Start SITL process."""
        if not self.sim_vehicle_path:
            print("ERROR: sim_vehicle.py not found!")
            print("Please set --sitl-path or ARDUPILOT_HOME environment variable")
            return False
        
        cmd = [
            sys.executable, self.sim_vehicle_path,
            '-v', self.vehicle,
            '-f', self.frame,
            '--model', 'JSON',
            '--sitl-instance-args', f'--sim-address=127.0.0.1',
            '--out', f'udp:127.0.0.1:{self.mavproxy_port}',
        ]
        
        if self.console:
            cmd.append('--console')
        if self.show_map:
            cmd.append('--map')
        
        print(f"Starting SITL: {' '.join(cmd)}")
        
        try:
            self.process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT
            )
            return True
        except Exception as e:
            print(f"Failed to start SITL: {e}")
            return False
    
    def stop(self):
        """Stop SITL process."""
        if self.process:
            self.process.terminate()
            self.process.wait(timeout=5.0)
            self.process = None


def wait_for_sitl(host: str = '127.0.0.1', port: int = 14550, timeout: float = 60.0) -> bool:
    """Wait for SITL to be ready."""
    print(f"Waiting for SITL on {host}:{port}...")
    start = time.time()
    
    while time.time() - start < timeout:
        try:
            conn = mavutil.mavlink_connection(f'udpin:{host}:{port}', timeout=1)
            if conn.wait_heartbeat(timeout=1):
                print("SITL is ready!")
                conn.close()
                return True
        except Exception:
            pass
        time.sleep(1)
    
    print("Timeout waiting for SITL")
    return False


def main():
    args = parse_args()
    
    if not MAVLINK_AVAILABLE:
        print("ERROR: pymavlink is required!")
        print("Install with: pip install pymavlink")
        return 1
    
    print("=" * 60)
    print("  ArduPilot SITL + Gazebo Simulation")
    print("=" * 60)
    print(f"  Vehicle: {args.vehicle}")
    print(f"  Frame: {args.frame}")
    print(f"  Rover Pattern: {args.pattern}")
    print(f"  MAVProxy Port: {args.mavproxy_port}")
    print("=" * 60)
    print()
    
    sitl = None
    
    # Start SITL if not skipped
    if not args.no_sitl:
        sim_vehicle = args.sitl_path or find_sim_vehicle()
        
        if not sim_vehicle:
            print("=" * 60)
            print("WARNING: sim_vehicle.py not found!")
            print()
            print("Please either:")
            print("1. Install ArduPilot: https://ardupilot.org/dev/docs/building-setup-linux-ubuntu-apt.html")
            print("2. Set ARDUPILOT_HOME environment variable")
            print("3. Use --sitl-path to specify location")
            print("4. Use --no-sitl if SITL is already running")
            print("=" * 60)
            print()
            print("Continuing without starting SITL...")
        else:
            sitl = SITLProcess(
                sim_vehicle,
                vehicle=args.vehicle,
                frame=args.frame,
                sitl_port=args.sitl_port,
                mavproxy_port=args.mavproxy_port,
                console=args.console,
                show_map=args.map
            )
            
            if not sitl.start():
                return 1
            
            print("\nWaiting for SITL to initialize...")
            time.sleep(5)  # Give SITL time to start
    
    # Wait for SITL connection
    if not wait_for_sitl(port=args.mavproxy_port, timeout=30):
        print("Could not connect to SITL!")
        if sitl:
            sitl.stop()
        return 1
    
    # Initialize ROS
    rclpy.init()
    
    nodes = []
    executor = MultiThreadedExecutor(num_threads=4)
    
    try:
        # Create MAVLink bridge
        bridge = MAVLinkBridge(
            sitl_port=args.mavproxy_port
        )
        nodes.append(bridge)
        executor.add_node(bridge)
        print("[OK] MAVLink Bridge started")
        
        # Create drone controller
        drone = DroneController()
        nodes.append(drone)
        executor.add_node(drone)
        print("[OK] Drone Controller started")
        
        # Create rover controller
        if not args.no_rover:
            pattern = MovementPattern[args.pattern.upper()]
            rover = RoverController(
                pattern=pattern,
                speed=args.speed,
                amplitude=args.amplitude
            )
            nodes.append(rover)
            executor.add_node(rover)
            print("[OK] Rover Controller started")
        
        print()
        print("=" * 60)
        print("System ready!")
        print()
        print("To arm and takeoff, use MAVProxy commands:")
        print("  mode guided")
        print("  arm throttle")
        print("  takeoff 5")
        print()
        print("Or use the bridge.arm() and bridge.takeoff() methods")
        print("=" * 60)
        print()
        print("Press Ctrl+C to stop.")
        print()
        
        # Handle Ctrl+C
        def signal_handler(sig, frame):
            print("\nShutting down...")
            executor.shutdown()
        
        signal.signal(signal.SIGINT, signal_handler)
        
        executor.spin()
        
    except Exception as e:
        print(f"Error: {e}")
        raise
    finally:
        print("Cleaning up...")
        
        # Stop nodes
        for node in nodes:
            node.destroy_node()
        
        # Stop ROS
        if rclpy.ok():
            rclpy.shutdown()
        
        # Stop SITL
        if sitl:
            sitl.stop()
        
        print("Shutdown complete.")
    
    return 0


if __name__ == '__main__':
    sys.exit(main())