#!/bin/bash
# =============================================================================
# RDC Simulation - Ubuntu/WSL2 Installation Script
# =============================================================================
# Complete installation for GPS-Denied Drone Landing Simulation
# Installs: ROS 2, Gazebo, ArduPilot SITL, and all dependencies
#
# Usage: 
#   ./setup/install_ubuntu.sh                  # Full installation (recommended)
#   ./setup/install_ubuntu.sh --skip-ardupilot # Skip ArduPilot (basic only)
# =============================================================================

set -e

echo "=============================================="
echo "  RDC Simulation - Ubuntu Installation"
echo "=============================================="
echo ""

# Get script directory and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
VENV_DIR="$PROJECT_ROOT/venv"

# ArduPilot directories
ARDUPILOT_HOME="$HOME/ardupilot"
ARDUPILOT_GZ="$HOME/ardupilot_gazebo"

echo "[INFO] Project root: $PROJECT_ROOT"
echo "[INFO] Virtual environment: $VENV_DIR"

# Detect Ubuntu version
if [ -f /etc/os-release ]; then
    . /etc/os-release
    UBUNTU_VERSION=$VERSION_ID
    echo "[INFO] Detected: $NAME $VERSION_ID"
else
    echo "[WARN] Could not detect Ubuntu version, assuming 22.04"
    UBUNTU_VERSION="22.04"
fi

# Check for skip ArduPilot option (default is to install ArduPilot)
INSTALL_ARDUPILOT=true
for arg in "$@"; do
    if [ "$arg" = "--skip-ardupilot" ]; then
        INSTALL_ARDUPILOT=false
    fi
done

if [ "$INSTALL_ARDUPILOT" = true ]; then
    TOTAL_STEPS=9
    echo "[INFO] Full installation with ArduPilot SITL"
else
    TOTAL_STEPS=6
    echo "[INFO] Basic installation (no ArduPilot)"
fi
echo ""

# =============================================================================
# STEP 1: System Dependencies
# =============================================================================
echo "[STEP 1/$TOTAL_STEPS] Installing system dependencies..."

sudo apt-get update
sudo apt-get install -y \
    curl \
    gnupg \
    lsb-release \
    software-properties-common \
    python3 \
    python3-pip \
    python3-venv \
    python3-dev \
    git \
    cmake \
    build-essential \
    wget \
    netcat-openbsd \
    libgstreamer1.0-dev \
    libgstreamer-plugins-base1.0-dev \
    gstreamer1.0-plugins-bad \
    gstreamer1.0-libav \
    gstreamer1.0-gl \
    libopencv-dev

echo "[OK] System dependencies installed"

# =============================================================================
# STEP 2: ROS 2 Installation
# =============================================================================
echo ""
echo "[STEP 2/$TOTAL_STEPS] Installing ROS 2..."

# Add ROS 2 GPG key
sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg

# Determine ROS 2 distro and repository codename based on Ubuntu version
# ROS 2 only has repositories for specific Ubuntu LTS versions
# For newer/unsupported versions, we use the closest compatible LTS repository
if [ "$UBUNTU_VERSION" = "24.04" ]; then
    ROS_DISTRO="jazzy"
    ROS_UBUNTU_CODENAME="noble"
elif [ "$UBUNTU_VERSION" = "22.04" ]; then
    ROS_DISTRO="humble"
    ROS_UBUNTU_CODENAME="jammy"
else
    # For Ubuntu versions newer than 24.04, use Jazzy with noble repositories
    # This provides best compatibility for cutting-edge Ubuntu versions
    UBUNTU_MAJOR=$(echo "$UBUNTU_VERSION" | cut -d. -f1)
    if [ "$UBUNTU_MAJOR" -ge 24 ] 2>/dev/null; then
        echo "[WARN] Ubuntu $UBUNTU_VERSION not officially supported by ROS 2"
        echo "[INFO] Using ROS 2 Jazzy with Ubuntu 24.04 (noble) repositories"
        ROS_DISTRO="jazzy"
        ROS_UBUNTU_CODENAME="noble"
    else
        echo "[WARN] Ubuntu $UBUNTU_VERSION not officially supported, trying humble"
        ROS_DISTRO="humble"
        ROS_UBUNTU_CODENAME="jammy"
    fi
fi

echo "[INFO] Using ROS 2 $ROS_DISTRO (repository: $ROS_UBUNTU_CODENAME)"

# Add repository using the compatible Ubuntu codename
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $ROS_UBUNTU_CODENAME main" | sudo tee /etc/apt/sources.list.d/ros2.list > /dev/null

sudo apt-get update
sudo apt-get install -y ros-${ROS_DISTRO}-ros-base ros-${ROS_DISTRO}-geometry-msgs ros-${ROS_DISTRO}-std-msgs ros-${ROS_DISTRO}-nav-msgs ros-${ROS_DISTRO}-sensor-msgs || {
    echo "[ERROR] Failed to install ROS 2 $ROS_DISTRO"
    echo "[INFO] This may be due to library incompatibilities with Ubuntu $UBUNTU_VERSION"
    echo "[INFO] Consider using Ubuntu 22.04 or 24.04 for full ROS 2 support"
}

# Install ros-gz bridge
sudo apt-get install -y ros-${ROS_DISTRO}-ros-gz || {
    echo "[WARN] Could not install ros-gz"
}

echo "[OK] ROS 2 $ROS_DISTRO installed"

# =============================================================================
# STEP 3: Gazebo Installation
# =============================================================================
echo ""
echo "[STEP 3/$TOTAL_STEPS] Installing Gazebo..."

# Add Gazebo repository (use compatible Ubuntu codename if needed)
GZ_UBUNTU_CODENAME="${ROS_UBUNTU_CODENAME:-$(lsb_release -cs)}"
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 $GZ_UBUNTU_CODENAME main" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null
sudo apt-get update

# Install Gazebo Harmonic (preferred) or Garden as fallback
if sudo apt-get install -y gz-harmonic; then
    GZ_VERSION="harmonic"
    # Install Harmonic development packages for building plugins
    sudo apt-get install -y \
        libgz-cmake3-dev \
        libgz-sim8-dev \
        libgz-plugin2-dev \
        libgz-common5-dev \
        libgz-physics7-dev \
        libgz-sensors8-dev \
        libgz-rendering8-dev \
        libgz-transport13-dev \
        libgz-msgs10-dev \
        rapidjson-dev \
        2>/dev/null || echo "[WARN] Some Gazebo dev packages may be missing"
elif sudo apt-get install -y gz-garden; then
    GZ_VERSION="garden"
    # Install Garden development packages
    sudo apt-get install -y \
        libgz-cmake3-dev \
        libgz-sim7-dev \
        libgz-plugin2-dev \
        libgz-common5-dev \
        libgz-physics6-dev \
        libgz-sensors7-dev \
        libgz-rendering7-dev \
        libgz-transport12-dev \
        libgz-msgs9-dev \
        rapidjson-dev \
        2>/dev/null || echo "[WARN] Some Gazebo dev packages may be missing"
else
    GZ_VERSION="none"
    echo "[WARN] Could not install Gazebo Harmonic/Garden"
fi

# Verify installation
if command -v gz &> /dev/null; then
    echo "[OK] Gazebo $GZ_VERSION installed ($(gz sim --version 2>/dev/null | head -1 || echo 'version unknown'))"
else
    echo "[WARN] Gazebo command not found"
fi

# =============================================================================
# STEP 4: Python Virtual Environment
# =============================================================================
echo ""
echo "[STEP 4/$TOTAL_STEPS] Creating Python virtual environment..."

if [ -d "$VENV_DIR" ]; then
    rm -rf "$VENV_DIR"
fi

python3 -m venv "$VENV_DIR"
source "$VENV_DIR/bin/activate"
pip install --upgrade pip

if [ -f "$PROJECT_ROOT/requirements.txt" ]; then
    pip install -r "$PROJECT_ROOT/requirements.txt"
else
    pip install pybullet numpy pillow opencv-python pymavlink pexpect
fi

echo "[OK] Python virtual environment created"

# =============================================================================
# STEP 5: Create Activation Script
# =============================================================================
echo ""
echo "[STEP 5/$TOTAL_STEPS] Creating activation script..."

cat > "$PROJECT_ROOT/activate.sh" << 'EOF'
#!/bin/bash
# RDC Simulation - Environment Activation
# Usage: source activate.sh

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# =============================================================================
# 1. ROS 2 Setup
# =============================================================================
if [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    echo "[OK] ROS 2 jazzy"
elif [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    echo "[OK] ROS 2 humble"
else
    echo "[WARN] ROS 2 installation not found in /opt/ros/"
fi

# =============================================================================
# 2. ArduPilot Environment
# =============================================================================
export ARDUPILOT_HOME="$HOME/ardupilot"

# Add ArduPilot tools to PATH
export PATH="$ARDUPILOT_HOME/Tools/autotest:$PATH"
export PATH="$ARDUPILOT_HOME/Tools/scripts:$PATH"
export PATH="$HOME/.local/bin:$PATH"

# =============================================================================
# 3. Python Virtual Environment
# =============================================================================
# Deactivate any existing venv
if [ -n "$VIRTUAL_ENV" ]; then
    deactivate 2>/dev/null || true
fi

# Activate project venv
if [ -f "$SCRIPT_DIR/venv/bin/activate" ]; then
    source "$SCRIPT_DIR/venv/bin/activate"
    echo "[OK] Python venv active"
else
    echo "[WARN] Python venv not found at $SCRIPT_DIR/venv"
fi

# =============================================================================
# 4. Gazebo & Simulation Paths
# =============================================================================
# Add project models to Gazebo path
export GZ_SIM_RESOURCE_PATH="$SCRIPT_DIR/gazebo/models:$GZ_SIM_RESOURCE_PATH"

# ArduPilot Gazebo Plugin paths
if [ -d "$HOME/ardupilot_gazebo/build" ]; then
    export GZ_SIM_SYSTEM_PLUGIN_PATH="$HOME/ardupilot_gazebo/build:$GZ_SIM_SYSTEM_PLUGIN_PATH"
    export GZ_SIM_RESOURCE_PATH="$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds:$GZ_SIM_RESOURCE_PATH"
else
    echo "[WARN] ArduPilot Gazebo plugin not found at $HOME/ardupilot_gazebo"
fi

echo ""
echo "Environment Configured!"
echo "-----------------------"
echo "Project Root:   $SCRIPT_DIR"
echo "ArduPilot Home: $ARDUPILOT_HOME"
echo "-----------------------"
EOF

chmod +x "$PROJECT_ROOT/activate.sh"
echo "[OK] Created: $PROJECT_ROOT/activate.sh"

# =============================================================================
# STEP 6: Verify Base Installation
# =============================================================================
echo ""
echo "[STEP 6/$TOTAL_STEPS] Verifying base installation..."

source "$PROJECT_ROOT/activate.sh"

python3 -c "import pybullet; print('[OK] PyBullet')" || echo "[FAIL] PyBullet"
python3 -c "import numpy; print('[OK] NumPy')" || echo "[FAIL] NumPy"
python3 -c "import cv2; print('[OK] OpenCV')" || echo "[WARN] OpenCV not installed"
python3 -c "from pymavlink import mavutil; print('[OK] pymavlink')" || echo "[WARN] pymavlink not installed"

# =============================================================================
# ARDUPILOT INSTALLATION (Steps 7-10, if not skipped)
# =============================================================================
if [ "$INSTALL_ARDUPILOT" = true ]; then

    # =========================================================================
    # STEP 7: Clone and Setup ArduPilot
    # =========================================================================
    echo ""
    echo "[STEP 7/$TOTAL_STEPS] Setting up ArduPilot SITL..."
    
    if [ ! -d "$ARDUPILOT_HOME" ]; then
        echo "[INFO] Cloning ArduPilot repository..."
        git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git "$ARDUPILOT_HOME"
    else
        echo "[INFO] ArduPilot directory already exists"
    fi
    
    cd "$ARDUPILOT_HOME"
    
    # Install ArduPilot prerequisites (this creates ~/.ardupilot_env)
    echo "[INFO] Installing ArduPilot prerequisites (this may take a while)..."
    Tools/environment_install/install-prereqs-ubuntu.sh -y
    . ~/.profile || true
    
    # Source ArduPilot environment
    if [ -f "$HOME/.ardupilot_env" ]; then
        source "$HOME/.ardupilot_env"
    fi
    
    echo "[OK] ArduPilot prerequisites installed"

    # =========================================================================
    # STEP 8: Build ArduCopter SITL
    # =========================================================================
    echo ""
    echo "[STEP 8/$TOTAL_STEPS] Building ArduCopter SITL..."
    
    cd "$ARDUPILOT_HOME"
    ./waf configure --board sitl
    ./waf copter
    
    echo "[OK] ArduCopter SITL built"

    # =========================================================================
    # STEP 9: Build ArduPilot Gazebo Plugin
    # =========================================================================
    echo ""
    echo "[STEP 9/$TOTAL_STEPS] Building ArduPilot Gazebo plugin..."
    
    if [ ! -d "$ARDUPILOT_GZ" ]; then
        echo "[INFO] Cloning ardupilot_gazebo repository..."
        git clone https://github.com/ArduPilot/ardupilot_gazebo.git "$ARDUPILOT_GZ"
    else
        echo "[INFO] ardupilot_gazebo directory already exists"
    fi
    
    cd "$ARDUPILOT_GZ"
    mkdir -p build && cd build
    cmake .. -DCMAKE_BUILD_TYPE=Release
    make -j$(nproc)
    
    echo "[OK] ArduPilot Gazebo plugin built"
    
    # Install MAVProxy
    pip3 install --user pymavlink mavproxy pexpect
    
    # Verify ArduPilot installation
    echo ""
    echo "Verifying ArduPilot installation..."
    
    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"
fi

# =============================================================================
# INSTALLATION COMPLETE
# =============================================================================
echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""

if [ "$INSTALL_ARDUPILOT" = true ]; then
    echo "Quick Start (3 terminals):"
    echo ""
    echo "Terminal 1 - Start Gazebo:"
    echo "  cd ~/RDC_Simulation"
    echo "  source activate.sh"
    echo "  ./scripts/run_ardupilot_sim.sh runway"
    echo ""
    echo "Terminal 2 - Start ArduCopter SITL:"
    echo "  cd ~/RDC_Simulation"
    echo "  source activate.sh"
    echo "  sim_vehicle.py -v ArduCopter -f gazebo-iris --model JSON --console"
    echo ""
    echo "Terminal 3 - Run Controller:"
    echo "  cd ~/RDC_Simulation"
    echo "  source activate.sh"
    echo "  python scripts/run_ardupilot.py --pattern square"
else
    echo "Basic installation complete."
    echo ""
    echo "To install ArduPilot SITL later, run:"
    echo "  ./setup/install_ubuntu.sh"
    echo ""
    echo "Or run manually:"
    echo "  cd ~/RDC_Simulation"
    echo "  source activate.sh"
fi
echo ""