#!/bin/bash
# =============================================================================
# ArduPilot ROS 2 + Gazebo Installation Script
# =============================================================================
# Installs the official ArduPilot ROS 2 packages with DDS and Gazebo support.
# This is the recommended way to run ArduPilot SITL with ROS 2.
#
# Prerequisites:
#   - Ubuntu 22.04 (ROS 2 Humble)
#   - ROS 2 Humble installed
#
# Usage: ./install_ardupilot.sh [--skip-sitl] [--skip-gazebo]
#
# This installs:
#   1. ArduPilot SITL with DDS support
#   2. ardupilot_ros - ArduPilot ROS 2 packages
#   3. ardupilot_gz - ArduPilot Gazebo integration
#   4. MAVProxy ground control station
# =============================================================================

set -e

echo "=============================================="
echo "  ArduPilot ROS 2 + Gazebo Installation"
echo "=============================================="
echo ""

# Parse arguments
SKIP_SITL=false
SKIP_GAZEBO=false
for arg in "$@"; do
    case $arg in
        --skip-sitl)
            SKIP_SITL=true
            ;;
        --skip-gazebo)
            SKIP_GAZEBO=true
            ;;
    esac
done

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Directories
ARDUPILOT_WS="$HOME/ardu_ws"
ARDUPILOT_DIR="$HOME/ardupilot"

# Check for ROS 2
if [ ! -f "/opt/ros/humble/setup.bash" ] && [ ! -f "/opt/ros/jazzy/setup.bash" ]; then
    echo "[ERROR] ROS 2 not found!"
    echo "Please install ROS 2 first:"
    echo "  ./setup/install_ubuntu.sh"
    exit 1
fi

# Source ROS 2
if [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    ROS_DISTRO="humble"
    GZ_VERSION="harmonic"
elif [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    ROS_DISTRO="jazzy"
    GZ_VERSION="harmonic"
fi

echo "[INFO] ROS 2 Distro: $ROS_DISTRO"
echo "[INFO] Gazebo Version: $GZ_VERSION"
echo "[INFO] Workspace: $ARDUPILOT_WS"
echo ""

# -----------------------------------------------------------------------------
# Step 1: Install System Dependencies
# -----------------------------------------------------------------------------
echo "[STEP 1/7] Installing system dependencies..."

sudo apt-get update
sudo apt-get install -y \
    git \
    cmake \
    build-essential \
    python3 \
    python3-pip \
    python3-dev \
    python3-venv \
    python3-vcstool \
    python3-rosdep \
    python3-colcon-common-extensions \
    wget \
    curl

# Install additional dependencies
sudo apt-get install -y \
    default-jre \
    libxml2-dev \
    libxslt1-dev \
    libtool \
    automake \
    autoconf \
    libexpat1-dev \
    ccache || true

echo "[INFO] System dependencies installed"

# -----------------------------------------------------------------------------
# Step 2: Install Gazebo Harmonic
# -----------------------------------------------------------------------------
if [ "$SKIP_GAZEBO" = false ]; then
    echo ""
    echo "[STEP 2/7] Installing Gazebo $GZ_VERSION..."
    
    # Add Gazebo APT sources
    sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null
    sudo apt-get update
    
    # Install Gazebo
    sudo apt-get install -y gz-$GZ_VERSION || {
        echo "[WARN] Could not install gz-$GZ_VERSION, trying gz-garden..."
        sudo apt-get install -y gz-garden || true
    }
    
    echo "[INFO] Gazebo installed"
else
    echo ""
    echo "[STEP 2/7] Skipping Gazebo installation (--skip-gazebo)"
fi

# -----------------------------------------------------------------------------
# Step 3: Create ArduPilot ROS 2 Workspace
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 3/7] Setting up ArduPilot ROS 2 workspace..."

mkdir -p "$ARDUPILOT_WS/src"
cd "$ARDUPILOT_WS"

# Import ArduPilot ROS 2 repositories using vcstool
echo "[INFO] Importing ArduPilot ROS 2 repositories..."

# Create ros2.repos file for ArduPilot packages
cat > /tmp/ardupilot_ros2.repos << 'EOF'
repositories:
  ardupilot:
    type: git
    url: https://github.com/ArduPilot/ardupilot.git
    version: master
  ardupilot_gz:
    type: git
    url: https://github.com/ArduPilot/ardupilot_gz.git
    version: main
EOF

vcs import --recursive src < /tmp/ardupilot_ros2.repos

# Also import Gazebo-related packages
if [ "$SKIP_GAZEBO" = false ]; then
    echo "[INFO] Importing Gazebo packages..."
    vcs import --input https://raw.githubusercontent.com/ArduPilot/ardupilot_gz/main/ros2_gz.repos --recursive src || true
fi

echo "[INFO] ArduPilot ROS 2 repositories imported"

# -----------------------------------------------------------------------------
# Step 4: Install ArduPilot SITL Prerequisites
# -----------------------------------------------------------------------------
if [ "$SKIP_SITL" = false ]; then
    echo ""
    echo "[STEP 4/7] Installing ArduPilot SITL prerequisites..."
    
    cd "$ARDUPILOT_WS/src/ardupilot"
    
    # Run ArduPilot prerequisites installer
    Tools/environment_install/install-prereqs-ubuntu.sh -y
    
    # Reload profile
    . ~/.profile || true
    
    echo "[INFO] ArduPilot SITL prerequisites installed"
else
    echo ""
    echo "[STEP 4/7] Skipping SITL prerequisites (--skip-sitl)"
fi

# -----------------------------------------------------------------------------
# Step 5: Configure rosdep for Gazebo
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 5/7] Configuring rosdep..."

# Initialize rosdep if needed
if [ ! -f /etc/ros/rosdep/sources.list.d/20-default.list ]; then
    sudo rosdep init || true
fi
rosdep update

# Add Gazebo sources to rosdep (for non-default ROS 2 Humble + Gazebo Harmonic pairing)
if [ "$ROS_DISTRO" = "humble" ] && [ "$GZ_VERSION" = "harmonic" ]; then
    sudo wget https://raw.githubusercontent.com/osrf/osrf-rosdep/master/gz/00-gazebo.list \
        -O /etc/ros/rosdep/sources.list.d/00-gazebo.list || true
    rosdep update || true
fi

# Install ROS dependencies
cd "$ARDUPILOT_WS"
rosdep install --from-paths src --ignore-src -y || true

echo "[INFO] rosdep configured"

# -----------------------------------------------------------------------------
# Step 6: Build ArduPilot ROS 2 Packages
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 6/7] Building ArduPilot ROS 2 packages..."

cd "$ARDUPILOT_WS"
source /opt/ros/$ROS_DISTRO/setup.bash

# Set Gazebo version environment variable
export GZ_VERSION=$GZ_VERSION

# Build packages
if [ "$SKIP_GAZEBO" = false ]; then
    echo "[INFO] Building with Gazebo support..."
    colcon build --packages-up-to ardupilot_gz_bringup --symlink-install || {
        echo "[WARN] Full build failed, trying core packages only..."
        colcon build --packages-up-to ardupilot_sitl --symlink-install
    }
else
    echo "[INFO] Building SITL only..."
    colcon build --packages-up-to ardupilot_sitl --symlink-install
fi

echo "[INFO] ArduPilot ROS 2 packages built"

# -----------------------------------------------------------------------------
# Step 7: Configure Environment Variables
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 7/7] Configuring environment variables..."

BASHRC_MARKER="# === ArduPilot ROS 2 Configuration ==="

if ! grep -q "$BASHRC_MARKER" ~/.bashrc; then
    echo "" >> ~/.bashrc
    echo "$BASHRC_MARKER" >> ~/.bashrc
    echo "" >> ~/.bashrc
    
    # Gazebo version
    echo "# Gazebo Version" >> ~/.bashrc
    echo "export GZ_VERSION=$GZ_VERSION" >> ~/.bashrc
    echo "" >> ~/.bashrc
    
    # ArduPilot workspace
    echo "# ArduPilot ROS 2 Workspace" >> ~/.bashrc
    echo "export ARDUPILOT_WS=$ARDUPILOT_WS" >> ~/.bashrc
    echo "" >> ~/.bashrc
    
    # ArduPilot SITL paths
    echo "# ArduPilot SITL" >> ~/.bashrc
    echo "export PATH=\$PATH:$ARDUPILOT_WS/src/ardupilot/Tools/autotest" >> ~/.bashrc
    echo "" >> ~/.bashrc
    
    # Source ROS 2 and workspace
    echo "# Source ROS 2 and ArduPilot workspace" >> ~/.bashrc
    echo "source /opt/ros/$ROS_DISTRO/setup.bash" >> ~/.bashrc
    echo "if [ -f $ARDUPILOT_WS/install/setup.bash ]; then" >> ~/.bashrc
    echo "    source $ARDUPILOT_WS/install/setup.bash" >> ~/.bashrc
    echo "fi" >> ~/.bashrc
    echo "" >> ~/.bashrc
    
    echo "[INFO] Environment variables added to ~/.bashrc"
else
    echo "[INFO] Environment variables already configured in ~/.bashrc"
fi

# Install pymavlink in venv if it exists
if [ -d "$PROJECT_ROOT/venv" ]; then
    echo "[INFO] Installing pymavlink in project venv..."
    source "$PROJECT_ROOT/venv/bin/activate"
    pip install pymavlink mavproxy || true
    deactivate
fi

# Also install pymavlink and mavproxy globally
pip3 install --user pymavlink mavproxy || true

# -----------------------------------------------------------------------------
# Verification
# -----------------------------------------------------------------------------
echo ""
echo "=============================================="
echo "  Verifying Installation"
echo "=============================================="
echo ""

# Source the workspace
source "$ARDUPILOT_WS/install/setup.bash" || true

# Check ROS 2 packages
ros2 pkg list | grep -q "ardupilot_sitl" && echo "[OK] ardupilot_sitl package" || echo "[WARN] ardupilot_sitl not found"

if [ "$SKIP_GAZEBO" = false ]; then
    ros2 pkg list | grep -q "ardupilot_gz_bringup" && echo "[OK] ardupilot_gz_bringup package" || echo "[WARN] ardupilot_gz_bringup not found"
fi

# Check Gazebo
if command -v gz &> /dev/null; then
    echo "[OK] Gazebo (gz command)"
elif command -v ign &> /dev/null; then
    echo "[OK] Gazebo Fortress (ign command)"
else
    echo "[WARN] Gazebo command not found"
fi

# Check MAVProxy
python3 -c "from pymavlink import mavutil" &> /dev/null && echo "[OK] pymavlink" || echo "[WARN] pymavlink not found"

# -----------------------------------------------------------------------------
# Complete
# -----------------------------------------------------------------------------
echo ""
echo "=============================================="
echo "  ArduPilot ROS 2 Installation Complete!"
echo "=============================================="
echo ""
echo "IMPORTANT: Run the following to apply changes:"
echo "  source ~/.bashrc"
echo ""
echo "Quick Start - Run SITL with ROS 2:"
echo "  source ~/ardu_ws/install/setup.bash"
echo "  ros2 launch ardupilot_sitl sitl_dds_udp.launch.py \\"
echo "      transport:=udp4 \\"
echo "      synthetic_clock:=True \\"
echo "      model:=quad"
echo ""
echo "Quick Start - Run with Gazebo:"
echo "  source ~/ardu_ws/install/setup.bash"
echo "  ros2 launch ardupilot_gz_bringup iris_runway.launch.py"
echo ""
echo "MAVProxy (in another terminal):"
echo "  mavproxy.py --console --map --master=:14550"
echo ""
echo "ROS 2 Topics:"
echo "  ros2 topic list"
echo "  ros2 topic echo /ap/geopose/filtered"
echo ""
echo "For more info, see: docs/ardupilot.md"
echo ""