#!/bin/bash
# =============================================================================
# ArduPilot SITL + Gazebo Installation Script (MAVLink Mode)
# =============================================================================
# Installs ArduPilot SITL with the ardupilot_gazebo plugin.
# This uses MAVLink instead of DDS - simpler and more reliable.
#
# Usage: ./install_ardupilot.sh
# =============================================================================

set -e

echo "=============================================="
echo "  ArduPilot SITL + Gazebo Installation"
echo "=============================================="
echo ""

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

# Directories
ARDUPILOT_HOME="$HOME/ardupilot"
ARDUPILOT_GZ="$HOME/ardupilot_gazebo"

# Check for 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"
else
    echo "[WARN] ROS 2 not found, will install in standalone mode"
    ROS_DISTRO=""
fi

echo "[INFO] ROS 2: ${ROS_DISTRO:-not installed}"
echo ""

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

sudo apt-get update
sudo apt-get install -y \
    git cmake build-essential \
    python3 python3-pip python3-dev \
    wget curl

echo "[OK] System dependencies"

# -----------------------------------------------------------------------------
# Step 2: Clone and Setup ArduPilot
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/5] Setting up ArduPilot SITL..."

if [ ! -d "$ARDUPILOT_HOME" ]; then
    git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git "$ARDUPILOT_HOME"
fi

cd "$ARDUPILOT_HOME"

# Install ArduPilot prerequisites
Tools/environment_install/install-prereqs-ubuntu.sh -y
. ~/.profile || true

# Source ArduPilot environment (includes venv with empy)
if [ -f "$HOME/.ardupilot_env" ]; then
    source "$HOME/.ardupilot_env"
fi
if [ -f "$HOME/venv-ardupilot/bin/activate" ]; then
    source "$HOME/venv-ardupilot/bin/activate"
fi

# Build ArduCopter SITL (without DDS to avoid complexity)
./waf configure --board sitl
./waf copter

echo "[OK] ArduPilot SITL built"

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

# Add Gazebo repo
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

# Install Gazebo
sudo apt-get install -y gz-harmonic || sudo apt-get install -y gz-garden || {
    echo "[WARN] Could not install Gazebo Harmonic/Garden"
}

echo "[OK] Gazebo installed"

# -----------------------------------------------------------------------------
# Step 4: Build ArduPilot Gazebo Plugin
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 4/5] Building ArduPilot Gazebo plugin..."

if [ ! -d "$ARDUPILOT_GZ" ]; then
    git clone https://github.com/ArduPilot/ardupilot_gazebo.git "$ARDUPILOT_GZ"
fi

cd "$ARDUPILOT_GZ"
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(nproc)

echo "[OK] ArduPilot Gazebo plugin built"

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

BASHRC_MARKER="# === ArduPilot SITL ==="

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

$BASHRC_MARKER
# Source ArduPilot environment (venv + tools)
if [ -f ~/.ardupilot_env ]; then
    source ~/.ardupilot_env
fi
export ARDUPILOT_HOME=$ARDUPILOT_HOME
export PATH=\$PATH:$ARDUPILOT_HOME/Tools/autotest
export PATH=\$PATH:\$HOME/.local/bin
export GZ_SIM_SYSTEM_PLUGIN_PATH=$ARDUPILOT_GZ/build:\$GZ_SIM_SYSTEM_PLUGIN_PATH
export GZ_SIM_RESOURCE_PATH=$ARDUPILOT_GZ/models:$ARDUPILOT_GZ/worlds:\$GZ_SIM_RESOURCE_PATH
EOF
    echo "[OK] Environment configured"
else
    echo "[OK] Environment already configured"
fi

# Install MAVProxy and Python dependencies
pip3 install --user pymavlink mavproxy pexpect

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

source ~/.bashrc 2>/dev/null || true
export PATH=$PATH:$ARDUPILOT_HOME/Tools/autotest:$HOME/.local/bin

command -v sim_vehicle.py &> /dev/null && echo "[OK] sim_vehicle.py" || echo "[WARN] sim_vehicle.py not found"
command -v gz &> /dev/null && echo "[OK] Gazebo (gz)" || echo "[WARN] Gazebo not found"
command -v mavproxy.py &> /dev/null && echo "[OK] MAVProxy" || echo "[WARN] MAVProxy not in PATH"
[ -f "$ARDUPILOT_GZ/build/libArduPilotPlugin.so" ] && echo "[OK] ArduPilot Gazebo plugin" || echo "[WARN] Plugin not built"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Run: source ~/.bashrc"
echo ""
echo "Quick Start (2 terminals):"
echo ""
echo "Terminal 1 - Start Gazebo:"
echo "  gz sim -v4 -r $ARDUPILOT_GZ/worlds/iris_runway.sdf"
echo ""
echo "Terminal 2 - Start SITL + MAVProxy:"
echo "  sim_vehicle.py -v ArduCopter -f gazebo-iris --model JSON --console"
echo ""
echo "Or use sim_vehicle.py standalone (no Gazebo):"
echo "  sim_vehicle.py -v ArduCopter --console --map"
echo ""