#!/bin/bash
# UAV-UGV Simulation - Complete Installation Script
# Installs everything needed for GPS-denied navigation simulation
# Compatible with Ubuntu 22.04/24.04 and WSL2

set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

print_header() {
    echo ""
    echo -e "${BLUE}==========================================${NC}"
    echo -e "${BLUE}  $1${NC}"
    echo -e "${BLUE}==========================================${NC}"
    echo ""
}

print_step() {
    echo -e "${GREEN}[$1/$TOTAL_STEPS] $2${NC}"
}

print_info() {
    echo -e "${CYAN}INFO: $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}WARNING: $1${NC}"
}

print_error() {
    echo -e "${RED}ERROR: $1${NC}"
}

# Detect environment
detect_environment() {
    IS_WSL=false
    IS_WSL2=false
    
    if grep -qEi "(microsoft|wsl)" /proc/version 2>/dev/null; then
        IS_WSL=true
        if [ -f /run/WSL ] || grep -qi "wsl2" /proc/version 2>/dev/null; then
            IS_WSL2=true
        fi
    fi
    
    # Detect Ubuntu version
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        UBUNTU_VERSION="$VERSION_ID"
        UBUNTU_CODENAME="$VERSION_CODENAME"
    fi
    
    # Determine ROS distro
    case "$UBUNTU_VERSION" in
        "22.04") ROS_DISTRO="humble" ;;
        "24.04") ROS_DISTRO="jazzy" ;;
        "20.04") ROS_DISTRO="galactic" ;;
        *) ROS_DISTRO="humble" ;;
    esac
}

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

print_header "UAV-UGV Simulation - Complete Setup"
echo "GPS-Denied Navigation with Geofencing"
echo ""

detect_environment

echo -e "${CYAN}Detected Environment:${NC}"
echo "  Ubuntu: $UBUNTU_VERSION ($UBUNTU_CODENAME)"
echo "  WSL: $IS_WSL | WSL2: $IS_WSL2"
echo "  ROS 2 Target: $ROS_DISTRO"
echo ""

TOTAL_STEPS=10
STEP=1

# ============================================================================
# STEP 1: System Update
# ============================================================================
print_step $STEP "Updating system packages"
sudo apt-get update
sudo apt-get upgrade -y
((STEP++))

# ============================================================================
# STEP 2: Install Base Dependencies
# ============================================================================
print_step $STEP "Installing base dependencies"
sudo apt-get install -y \
    curl \
    gnupg \
    lsb-release \
    software-properties-common \
    wget \
    git \
    gitk \
    build-essential \
    cmake \
    python3-dev \
    python3-pip \
    python3-venv \
    python3-opencv \
    python3-matplotlib \
    python3-lxml \
    python3-yaml \
    python3-scipy \
    python3-future \
    libopencv-dev \
    libxml2-dev \
    libxslt1-dev \
    ccache \
    gawk \
    libtool-bin

# WSL-specific packages
if $IS_WSL; then
    print_info "Installing WSL GUI support packages"
    sudo apt-get install -y \
        x11-apps \
        x11-xserver-utils \
        dbus-x11 \
        mesa-utils \
        libgl1-mesa-glx
fi
((STEP++))

# ============================================================================
# STEP 3: Install ROS 2
# ============================================================================
print_step $STEP "Installing ROS 2 $ROS_DISTRO"

if [ ! -d "/opt/ros/$ROS_DISTRO" ]; then
    # Add ROS 2 repository
    if [ ! -f /usr/share/keyrings/ros-archive-keyring.gpg ]; then
        sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key \
            -o /usr/share/keyrings/ros-archive-keyring.gpg
    fi
    
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] \
        http://packages.ros.org/ros2/ubuntu $(lsb_release -cs) main" | \
        sudo tee /etc/apt/sources.list.d/ros2.list > /dev/null
    
    sudo apt-get update
    sudo apt-get install -y ros-${ROS_DISTRO}-desktop python3-colcon-common-extensions
else
    print_info "ROS 2 $ROS_DISTRO already installed"
fi
((STEP++))

# ============================================================================
# STEP 4: Install ROS 2 Packages
# ============================================================================
print_step $STEP "Installing ROS 2 packages"
sudo apt-get install -y \
    ros-${ROS_DISTRO}-mavros \
    ros-${ROS_DISTRO}-mavros-extras \
    ros-${ROS_DISTRO}-cv-bridge \
    ros-${ROS_DISTRO}-image-transport \
    ros-${ROS_DISTRO}-tf2 \
    ros-${ROS_DISTRO}-tf2-ros \
    ros-${ROS_DISTRO}-tf2-geometry-msgs \
    ros-${ROS_DISTRO}-gazebo-ros-pkgs 2>/dev/null || \
    print_warning "Some ROS packages not available for $ROS_DISTRO"

# Install GeographicLib datasets for MAVROS
GEOGRAPHICLIB_SCRIPT="/opt/ros/${ROS_DISTRO}/lib/mavros/install_geographiclib_datasets.sh"
if [ -f "$GEOGRAPHICLIB_SCRIPT" ]; then
    if [ ! -f /usr/share/GeographicLib/geoids/egm96-5.pgm ]; then
        print_info "Installing GeographicLib datasets (this may take a minute)"
        sudo "$GEOGRAPHICLIB_SCRIPT" || print_warning "GeographicLib installation failed"
    fi
fi
((STEP++))

# ============================================================================
# STEP 5: Install Gazebo
# ============================================================================
print_step $STEP "Installing Gazebo"
sudo apt-get install -y gazebo libgazebo-dev || \
sudo apt-get install -y gazebo11 libgazebo11-dev || \
print_warning "Gazebo installation may require manual setup"
((STEP++))

# ============================================================================
# STEP 6: Install ArduPilot SITL
# ============================================================================
print_step $STEP "Installing ArduPilot SITL"

ARDUPILOT_HOME="$HOME/ardupilot"
if [ ! -d "$ARDUPILOT_HOME" ]; then
    print_info "Cloning ArduPilot repository..."
    cd "$HOME"
    git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git
    cd ardupilot
    git submodule update --init --recursive
else
    print_info "ArduPilot already exists, updating..."
    cd "$ARDUPILOT_HOME"
    git fetch origin
    git submodule update --init --recursive
fi

# Install ArduPilot prerequisites
print_info "Installing ArduPilot prerequisites..."
cd "$ARDUPILOT_HOME"
USER_NONINTERACTIVE=1 Tools/environment_install/install-prereqs-ubuntu.sh -y || true

# Reload profile
. ~/.profile 2>/dev/null || true

# Build ArduPilot SITL
print_info "Building ArduPilot SITL (this may take several minutes)..."
cd "$ARDUPILOT_HOME"
./waf configure --board sitl
./waf copter
./waf rover
((STEP++))

# ============================================================================
# STEP 7: Install ardupilot_gazebo Plugin
# ============================================================================
print_step $STEP "Installing ardupilot_gazebo plugin"

ARDUPILOT_GAZEBO_HOME="$HOME/ardupilot_gazebo"
if [ ! -d "$ARDUPILOT_GAZEBO_HOME" ]; then
    print_info "Cloning ardupilot_gazebo..."
    cd "$HOME"
    git clone https://github.com/ArduPilot/ardupilot_gazebo.git
else
    print_info "ardupilot_gazebo already exists, updating..."
    cd "$ARDUPILOT_GAZEBO_HOME"
    git pull origin main || true
fi

cd "$ARDUPILOT_GAZEBO_HOME"
mkdir -p build && cd build
cmake ..
make -j$(nproc)
sudo make install
((STEP++))

# ============================================================================
# STEP 8: Setup Python Virtual Environment
# ============================================================================
print_step $STEP "Setting up Python environment"

cd "$SCRIPT_DIR"

if [ ! -d "venv" ]; then
    python3 -m venv venv
fi

source venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt || print_warning "Some Python packages failed"
deactivate
((STEP++))

# ============================================================================
# STEP 9: Configure Environment
# ============================================================================
print_step $STEP "Configuring environment"

# Create WSL environment file
if $IS_WSL; then
    cat > "$SCRIPT_DIR/wsl_env.sh" << 'WSLEOF'
#!/bin/bash
# WSL Environment for UAV-UGV Simulation

if [ -d "/mnt/wslg" ]; then
    export DISPLAY=:0
else
    export DISPLAY=$(cat /etc/resolv.conf 2>/dev/null | grep nameserver | awk '{print $2}'):0
fi

export LIBGL_ALWAYS_INDIRECT=0
export MESA_GL_VERSION_OVERRIDE=3.3
export MESA_GLSL_VERSION_OVERRIDE=330
export OGRE_RTT_MODE=Copy

# Source Gazebo setup
[ -f /usr/share/gazebo/setup.bash ] && source /usr/share/gazebo/setup.bash
[ -f /usr/share/gazebo-11/setup.bash ] && source /usr/share/gazebo-11/setup.bash
WSLEOF
    chmod +x "$SCRIPT_DIR/wsl_env.sh"
fi

# Create activation script
cat > "$SCRIPT_DIR/activate_venv.sh" << EOF
#!/bin/bash
SCRIPT_DIR="\$(cd "\$(dirname "\${BASH_SOURCE[0]}")" && pwd)"

# ROS 2
source /opt/ros/${ROS_DISTRO}/setup.bash

# Gazebo
[ -f /usr/share/gazebo/setup.bash ] && source /usr/share/gazebo/setup.bash
[ -f /usr/share/gazebo-11/setup.bash ] && source /usr/share/gazebo-11/setup.bash

# Python venv
source "\$SCRIPT_DIR/venv/bin/activate"

# Gazebo paths
export GAZEBO_MODEL_PATH="\$SCRIPT_DIR/models:\$HOME/ardupilot_gazebo/models:\${GAZEBO_MODEL_PATH:-}"
export GAZEBO_RESOURCE_PATH="\$SCRIPT_DIR/worlds:\$HOME/ardupilot_gazebo/worlds:\${GAZEBO_RESOURCE_PATH:-}"

# ArduPilot
export PATH="\$PATH:\$HOME/ardupilot/Tools/autotest"
export ARDUPILOT_HOME="\$HOME/ardupilot"

# WSL
if grep -qEi "(microsoft|wsl)" /proc/version 2>/dev/null; then
    [ -f "\$SCRIPT_DIR/wsl_env.sh" ] && source "\$SCRIPT_DIR/wsl_env.sh"
fi

echo -e "\033[0;32mEnvironment activated (ROS 2 ${ROS_DISTRO})\033[0m"
echo ""
echo "Run simulation:  bash scripts/run_simulation.sh"
echo "With rendering:  bash scripts/run_simulation.sh --software-render"
EOF
chmod +x "$SCRIPT_DIR/activate_venv.sh"

((STEP++))

# ============================================================================
# STEP 10: Make Scripts Executable
# ============================================================================
print_step $STEP "Finalizing installation"
chmod +x "$SCRIPT_DIR/scripts/"*.sh 2>/dev/null || true
chmod +x "$SCRIPT_DIR/activate_venv.sh" 2>/dev/null || true
((STEP++))

# ============================================================================
# COMPLETE
# ============================================================================
print_header "Installation Complete!"

echo -e "${GREEN}All components installed:${NC}"
echo "  - ROS 2 $ROS_DISTRO"
echo "  - Gazebo"
echo "  - ArduPilot SITL ($ARDUPILOT_HOME)"
echo "  - ardupilot_gazebo ($ARDUPILOT_GAZEBO_HOME)"
echo "  - MAVROS"
echo "  - Python dependencies"
echo ""

if $IS_WSL; then
    echo -e "${YELLOW}WSL Setup Notes:${NC}"
    echo "  - GUI apps require WSLg (Windows 11) or VcXsrv (Windows 10)"
    echo "  - Use --software-render flag if graphics are slow"
    echo ""
fi

echo -e "${CYAN}To run the simulation:${NC}"
echo ""
echo "  cd $SCRIPT_DIR"
echo "  source activate_venv.sh"
echo "  bash scripts/run_simulation.sh"
if $IS_WSL; then
echo ""
echo "  # Or with software rendering for WSL:"
echo "  bash scripts/run_simulation.sh --software-render"
fi
echo ""
echo -e "${GREEN}==========================================${NC}"