#!/bin/bash
# =============================================================================
# ArduPilot ROS 2 + Gazebo Installation Script
# =============================================================================
# Installs the official ArduPilot ROS 2 packages with DDS and Gazebo support.
#
# Usage: ./install_ardupilot.sh
#
# This installs:
#   1. Micro-XRCE-DDS-Gen (required for DDS)
#   2. ArduPilot SITL with DDS support
#   3. ardupilot_gz - ArduPilot Gazebo integration
#   4. MAVProxy ground control station
# =============================================================================

set -e

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

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

# Directories
ARDUPILOT_WS="$HOME/ardu_ws"

# 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 "Install with: ./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"
elif [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    ROS_DISTRO="jazzy"
fi

GZ_VERSION="harmonic"

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

# -----------------------------------------------------------------------------
# Step 1: 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 default-jre gradle

echo "[OK] System dependencies"

# -----------------------------------------------------------------------------
# Step 2: Install Micro-XRCE-DDS-Gen (required for DDS)
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/7] Installing Micro-XRCE-DDS-Gen..."

XRCEDDSGEN_DIR="$HOME/Micro-XRCE-DDS-Gen"

if ! command -v microxrceddsgen &> /dev/null; then
    if [ ! -d "$XRCEDDSGEN_DIR" ]; then
        git clone --recurse-submodules https://github.com/eProsima/Micro-XRCE-DDS-Gen.git "$XRCEDDSGEN_DIR"
    fi
    
    cd "$XRCEDDSGEN_DIR"
    ./gradlew assemble
    
    # Add to PATH
    if ! grep -q "Micro-XRCE-DDS-Gen" ~/.bashrc; then
        echo "" >> ~/.bashrc
        echo "# Micro-XRCE-DDS-Gen" >> ~/.bashrc
        echo "export PATH=\$PATH:$XRCEDDSGEN_DIR/scripts" >> ~/.bashrc
    fi
    export PATH=$PATH:$XRCEDDSGEN_DIR/scripts
    
    echo "[OK] Micro-XRCE-DDS-Gen installed"
else
    echo "[OK] Micro-XRCE-DDS-Gen already installed"
fi

# -----------------------------------------------------------------------------
# Step 3: Install Gazebo Harmonic
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 3/7] Installing Gazebo Harmonic..."

sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg 2>/dev/null
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
sudo apt-get install -y gz-harmonic || echo "[WARN] Could not install gz-harmonic"

echo "[OK] Gazebo"

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

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

# Import ArduPilot repos
vcs import --recursive src < <(cat << '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
) || true

# Import Gazebo packages
vcs import --input https://raw.githubusercontent.com/ArduPilot/ardupilot_gz/main/ros2_gz.repos --recursive src 2>/dev/null || true

echo "[OK] Repositories imported"

# -----------------------------------------------------------------------------
# Step 5: Install ArduPilot Prerequisites
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 5/7] Installing ArduPilot prerequisites..."

cd "$ARDUPILOT_WS/src/ardupilot"
Tools/environment_install/install-prereqs-ubuntu.sh -y
. ~/.profile || true

echo "[OK] ArduPilot prerequisites"

# -----------------------------------------------------------------------------
# Step 6: Configure rosdep and Build
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 6/7] Building packages..."

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

# Add Gazebo Harmonic sources
sudo wget https://raw.githubusercontent.com/osrf/osrf-rosdep/master/gz/00-gazebo.list \
    -O /etc/ros/rosdep/sources.list.d/00-gazebo.list 2>/dev/null || true
rosdep update || true

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

# Build
source /opt/ros/$ROS_DISTRO/setup.bash
export GZ_VERSION=$GZ_VERSION
colcon build --packages-up-to ardupilot_gz_bringup --symlink-install --allow-overriding ros_gz_bridge ros_gz_sim sdformat_urdf || {
    echo "[WARN] Full build failed, trying ardupilot_sitl only..."
    colcon build --packages-up-to ardupilot_sitl --symlink-install
}

echo "[OK] Build complete"

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

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

if ! grep -q "$BASHRC_MARKER" ~/.bashrc; then
    cat >> ~/.bashrc << EOF

$BASHRC_MARKER
export GZ_VERSION=$GZ_VERSION
export PATH=\$PATH:$ARDUPILOT_WS/src/ardupilot/Tools/autotest
export PATH=\$PATH:\$HOME/.local/bin
source /opt/ros/$ROS_DISTRO/setup.bash
[ -f $ARDUPILOT_WS/install/setup.bash ] && source $ARDUPILOT_WS/install/setup.bash
EOF
    echo "[OK] Environment configured"
else
    echo "[OK] Environment already configured"
fi

# Install MAVProxy
pip3 install --user pymavlink mavproxy || true

# Verify
echo ""
echo "=============================================="
echo "  Verifying Installation"
echo "=============================================="

source "$ARDUPILOT_WS/install/setup.bash" 2>/dev/null || true

ros2 pkg list 2>/dev/null | grep -q "ardupilot_sitl" && echo "[OK] ardupilot_sitl" || echo "[WARN] ardupilot_sitl not found"
ros2 pkg list 2>/dev/null | grep -q "ardupilot_gz_bringup" && echo "[OK] ardupilot_gz_bringup" || echo "[WARN] ardupilot_gz_bringup not found"
command -v microxrceddsgen &> /dev/null && echo "[OK] microxrceddsgen" || echo "[WARN] microxrceddsgen not in PATH"
command -v mavproxy.py &> /dev/null && echo "[OK] MAVProxy" || echo "[WARN] MAVProxy not in PATH (add ~/.local/bin to PATH)"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Run: source ~/.bashrc"
echo ""
echo "Quick Start (2 terminals):"
echo ""
echo "Terminal 1:"
echo "  source ~/ardu_ws/install/setup.bash"
echo "  ros2 launch ardupilot_gz_bringup iris_runway.launch.py"
echo ""
echo "Terminal 2:"
echo "  mavproxy.py --console --map --master=:14550"
echo ""