#!/usr/bin/env python3
"""
ArduPilot ROS 2 Launcher - Official DDS Integration.

This script provides a convenient way to launch the ArduPilot SITL simulation
using the official ardupilot_gz packages with ROS 2 DDS support.

Usage:
    python run_ardupilot.py                    # Launch Iris on runway
    python run_ardupilot.py --world maze       # Launch Iris in maze
    python run_ardupilot.py --vehicle rover    # Launch WildThumper rover
    python run_ardupilot.py --mavproxy         # Also start MAVProxy

Prerequisites:
    - ArduPilot ROS 2 packages installed (./setup/install_ardupilot.sh)
    - ROS 2 Humble/Jazzy sourced
    - ~/ardu_ws workspace built and sourced
"""

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


def check_ros2():
    """Check if ROS 2 is available."""
    try:
        subprocess.run(['ros2', '--help'], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


def check_ardupilot_packages():
    """Check if ArduPilot ROS 2 packages are installed."""
    try:
        result = subprocess.run(
            ['ros2', 'pkg', 'list'],
            capture_output=True,
            text=True,
            check=True
        )
        packages = result.stdout
        return 'ardupilot_gz_bringup' in packages or 'ardupilot_sitl' in packages
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


def source_workspace():
    """Source the ArduPilot workspace."""
    ardu_ws = os.path.expanduser("~/ardu_ws")
    setup_bash = os.path.join(ardu_ws, "install", "setup.bash")
    
    if os.path.exists(setup_bash):
        # Update environment by sourcing the workspace
        # This is done by running commands in a sourced shell
        return True
    return False


def get_launch_command(world: str, vehicle: str) -> list:
    """Get the appropriate launch command."""
    
    launch_files = {
        # Copter configurations
        'runway': ('ardupilot_gz_bringup', 'iris_runway.launch.py'),
        'maze': ('ardupilot_gz_bringup', 'iris_maze.launch.py'),
        'iris': ('ardupilot_gz_bringup', 'iris_runway.launch.py'),
        
        # Rover configurations
        'rover': ('ardupilot_gz_bringup', 'wildthumper_playpen.launch.py'),
        'wildthumper': ('ardupilot_gz_bringup', 'wildthumper_playpen.launch.py'),
        
        # SITL only (no Gazebo)
        'sitl': ('ardupilot_sitl', 'sitl_dds_udp.launch.py'),
    }
    
    if vehicle == 'rover':
        key = 'rover'
    else:
        key = world if world in launch_files else 'runway'
    
    package, launch_file = launch_files.get(key, launch_files['runway'])
    
    cmd = ['ros2', 'launch', package, launch_file]
    
    # Add SITL-specific parameters if using sitl_dds_udp
    if launch_file == 'sitl_dds_udp.launch.py':
        cmd.extend([
            'transport:=udp4',
            'synthetic_clock:=True',
            'model:=quad' if vehicle == 'copter' else 'rover',
            'speedup:=1',
        ])
    
    return cmd


def launch_mavproxy(master_port: int = 14550):
    """Launch MAVProxy in a new terminal."""
    mavproxy_cmd = f"mavproxy.py --console --map --master=:{master_port}"
    
    # Try different terminal emulators
    terminals = [
        ['gnome-terminal', '--', 'bash', '-c', mavproxy_cmd],
        ['xterm', '-e', mavproxy_cmd],
        ['konsole', '-e', mavproxy_cmd],
    ]
    
    for term_cmd in terminals:
        try:
            subprocess.Popen(term_cmd)
            return True
        except FileNotFoundError:
            continue
    
    print(f"[WARN] Could not open terminal for MAVProxy")
    print(f"[INFO] Run manually: {mavproxy_cmd}")
    return False


def parse_args():
    parser = argparse.ArgumentParser(
        description='Launch ArduPilot SITL with ROS 2 and Gazebo',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run_ardupilot.py                  # Iris on runway
  python run_ardupilot.py --world maze     # Iris in maze
  python run_ardupilot.py --vehicle rover  # WildThumper rover
  python run_ardupilot.py --mavproxy       # With MAVProxy

Available worlds: runway, maze, sitl (no Gazebo)
Available vehicles: copter, rover
        """
    )
    
    parser.add_argument(
        '--world', '-w', type=str, default='runway',
        choices=['runway', 'maze', 'sitl'],
        help='Simulation world (default: runway)'
    )
    parser.add_argument(
        '--vehicle', '-v', type=str, default='copter',
        choices=['copter', 'rover'],
        help='Vehicle type (default: copter)'
    )
    parser.add_argument(
        '--mavproxy', '-m', action='store_true',
        help='Also launch MAVProxy in a new terminal'
    )
    parser.add_argument(
        '--mavproxy-port', type=int, default=14550,
        help='MAVProxy master port (default: 14550)'
    )
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    print("=" * 60)
    print("  ArduPilot SITL + Gazebo (Official ROS 2 DDS)")
    print("=" * 60)
    print()
    
    # Check ROS 2
    if not check_ros2():
        print("[ERROR] ROS 2 not found!")
        print("Please source ROS 2:")
        print("  source /opt/ros/humble/setup.bash")
        return 1
    
    print("[OK] ROS 2 available")
    
    # Check ArduPilot packages
    if not check_ardupilot_packages():
        print("[ERROR] ArduPilot ROS 2 packages not found!")
        print()
        print("Please install ArduPilot ROS 2:")
        print("  ./setup/install_ardupilot.sh")
        print()
        print("Then source the workspace:")
        print("  source ~/ardu_ws/install/setup.bash")
        return 1
    
    print("[OK] ArduPilot ROS 2 packages found")
    
    # Get launch command
    launch_cmd = get_launch_command(args.world, args.vehicle)
    
    print()
    print(f"World: {args.world}")
    print(f"Vehicle: {args.vehicle}")
    print(f"Launch: {' '.join(launch_cmd)}")
    print()
    
    # Launch MAVProxy if requested
    if args.mavproxy:
        print("[INFO] Starting MAVProxy...")
        # Delay to let SITL start first
        time.sleep(2)
        launch_mavproxy(args.mavproxy_port)
    
    print("Starting simulation...")
    print("Press Ctrl+C to stop.")
    print()
    print("-" * 60)
    
    # Handle Ctrl+C gracefully
    def signal_handler(sig, frame):
        print("\nShutting down...")
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    
    # Run the launch command
    try:
        subprocess.run(launch_cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] Launch failed: {e}")
        return 1
    except KeyboardInterrupt:
        print("\nShutdown complete.")
    
    return 0


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