#!/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
#
# Usage:
#   ./setup.sh                    # Full installation
#   ./setup.sh --skip-ardupilot   # Skip ArduPilot (Gazebo only)
# =============================================================================

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}"
}

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

# Deactivate existing venv if active to avoid conflicts
if [ -n "$VIRTUAL_ENV" ]; then
    deactivate 2>/dev/null || true
fi

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

# 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 grep -qi "wsl2" /proc/version 2>/dev/null || [ -f /run/WSL ]; 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"
    else
        UBUNTU_VERSION="22.04"
        UBUNTU_CODENAME="jammy"
    fi
    
    # Determine ROS distro and compatible repository codename
    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 newer versions, use latest available
        UBUNTU_MAJOR=$(echo "$UBUNTU_VERSION" | cut -d. -f1)
        if [ "$UBUNTU_MAJOR" -ge 24 ] 2>/dev/null; then
            ROS_DISTRO="jazzy"
            ROS_UBUNTU_CODENAME="noble"
        else
            ROS_DISTRO="humble"
            ROS_UBUNTU_CODENAME="jammy"
        fi
    fi
}

# Parse arguments
INSTALL_ARDUPILOT=true
for arg in "$@"; do
    if [ "$arg" = "--skip-ardupilot" ]; then
        INSTALL_ARDUPILOT=false
    fi
done

detect_environment

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

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 "  ArduPilot: $INSTALL_ARDUPILOT"
echo ""

if [ "$INSTALL_ARDUPILOT" = true ]; then
    TOTAL_STEPS=10
else
    TOTAL_STEPS=7
fi
STEP=1

# ============================================================================
# STEP 1: System Update & Dependencies
# ============================================================================
print_step $STEP "Installing system dependencies"

sudo apt-get update
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-full \
    python3-pexpect \
    python3-opencv \
    python3-matplotlib \
    python3-lxml \
    python3-yaml \
    python3-scipy \
    python3-future \
    pipx \
    libopencv-dev \
    libxml2-dev \
    libxslt1-dev \
    libgstreamer1.0-dev \
    libgstreamer-plugins-base1.0-dev \
    gstreamer1.0-plugins-bad \
    gstreamer1.0-libav \
    gstreamer1.0-gl \
    ccache \
    gawk \
    libtool-bin \
    netcat-openbsd

# 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 2: Install ROS 2
# ============================================================================
print_step $STEP "Installing ROS 2 $ROS_DISTRO"

# 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

# Add repository using 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

# Install ROS 2
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 \
    python3-colcon-common-extensions || {
    print_error "Failed to install ROS 2 $ROS_DISTRO"
    exit 1
}

# Install ros-gz bridge
sudo apt-get install -y ros-${ROS_DISTRO}-ros-gz || \
    print_warning "Could not install ros-gz bridge"

print_info "ROS 2 $ROS_DISTRO installed"
((STEP++))

# ============================================================================
# STEP 3: Install Gazebo (Harmonic or Garden)
# ============================================================================
print_step $STEP "Installing Gazebo"

# Add Gazebo repository
GZ_UBUNTU_CODENAME="$ROS_UBUNTU_CODENAME"
sudo wget -q https://packages.osrfoundation.org/gazebo.gpg \
    -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg 2>/dev/null || true

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

# Try Gazebo Harmonic first, then Garden
GZ_VERSION=""
if sudo apt-get install -y gz-harmonic 2>/dev/null; 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 || print_warning "Some Gazebo dev packages may be missing"
elif sudo apt-get install -y gz-garden 2>/dev/null; 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 || print_warning "Some Gazebo dev packages may be missing"
else
    print_warning "Could not install Gazebo Harmonic/Garden"
    GZ_VERSION="none"
fi

if command -v gz &> /dev/null; then
    print_info "Gazebo $GZ_VERSION installed"
else
    print_warning "Gazebo command not found"
fi
((STEP++))

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

if [ -d "$SCRIPT_DIR/venv" ]; then
    rm -rf "$SCRIPT_DIR/venv"
fi

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

if [ -f "$SCRIPT_DIR/requirements.txt" ]; then
    pip install -r "$SCRIPT_DIR/requirements.txt" || print_warning "Some Python packages failed"
else
    pip install numpy opencv-python scipy shapely filterpy transforms3d pymavlink pexpect future
fi

deactivate
((STEP++))

# ============================================================================
# STEP 5: Create Activation Script
# ============================================================================
print_step $STEP "Creating activation script"

cat > "$SCRIPT_DIR/activate_venv.sh" << 'ACTIVATE_EOF'
#!/bin/bash
# UAV-UGV Simulation - Environment Activation
# Usage: source activate_venv.sh

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

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

# ArduPilot Environment
export ARDUPILOT_HOME="$HOME/ardupilot"
export PATH="$ARDUPILOT_HOME/Tools/autotest:$PATH"
export PATH="$HOME/.local/bin:$PATH"

# Deactivate 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"
fi

# Gazebo paths (new Gazebo - Ignition/Harmonic)
export GZ_SIM_RESOURCE_PATH="$SCRIPT_DIR/models:$SCRIPT_DIR/worlds:${GZ_SIM_RESOURCE_PATH:-}"

# ArduPilot Gazebo plugin
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"
fi

# WSL environment
if grep -qEi "(microsoft|wsl)" /proc/version 2>/dev/null; then
    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
fi

echo -e "\033[0;32mEnvironment activated (ROS 2 $ROS_VER)\033[0m"
echo ""
echo "Run simulation:  bash scripts/run_simulation.sh"
ACTIVATE_EOF

chmod +x "$SCRIPT_DIR/activate_venv.sh"
((STEP++))

# ============================================================================
# STEP 6: Create WSL Environment (if applicable)
# ============================================================================
print_step $STEP "Configuring environment files"

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
WSLEOF
    chmod +x "$SCRIPT_DIR/wsl_env.sh"
fi
((STEP++))

# ============================================================================
# ARDUPILOT INSTALLATION (if not skipped)
# ============================================================================
if [ "$INSTALL_ARDUPILOT" = true ]; then

    # ========================================================================
    # STEP 7: Clone and Setup ArduPilot
    # ========================================================================
    print_step $STEP "Setting up ArduPilot SITL"
    
    if [ ! -d "$ARDUPILOT_HOME" ]; then
        print_info "Cloning ArduPilot repository..."
        git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git "$ARDUPILOT_HOME"
    else
        print_info "ArduPilot directory already exists"
    fi
    
    cd "$ARDUPILOT_HOME"
    git submodule update --init --recursive
    
    # Install ArduPilot prerequisites
    print_info "Installing ArduPilot prerequisites (this may take a while)..."
    Tools/environment_install/install-prereqs-ubuntu.sh -y || true
    . ~/.profile 2>/dev/null || true
    
    # Source ArduPilot environment
    [ -f "$HOME/.ardupilot_env" ] && source "$HOME/.ardupilot_env"
    
    print_info "ArduPilot prerequisites installed"
    ((STEP++))
    
    # ========================================================================
    # STEP 8: Build ArduCopter SITL
    # ========================================================================
    print_step $STEP "Building ArduPilot SITL (this takes several minutes)"
    
    cd "$ARDUPILOT_HOME"
    ./waf configure --board sitl
    ./waf copter
    ./waf rover
    
    print_info "ArduPilot SITL built"
    ((STEP++))
    
    # ========================================================================
    # STEP 9: Build ArduPilot Gazebo Plugin
    # ========================================================================
    print_step $STEP "Building ArduPilot Gazebo plugin"
    
    if [ ! -d "$ARDUPILOT_GZ" ]; then
        print_info "Cloning ardupilot_gazebo repository..."
        git clone https://github.com/ArduPilot/ardupilot_gazebo.git "$ARDUPILOT_GZ"
    else
        print_info "ardupilot_gazebo directory already exists"
        cd "$ARDUPILOT_GZ"
        git pull origin main || true
    fi
    
    cd "$ARDUPILOT_GZ"
    mkdir -p build && cd build
    cmake .. -DCMAKE_BUILD_TYPE=Release
    make -j$(nproc)
    
    print_info "ArduPilot Gazebo plugin built"
    
    # Install MAVProxy using pipx (required for Ubuntu 23.04+ PEP 668)
    print_info "Installing MAVProxy..."
    if pipx install MAVProxy --include-deps; then
        print_info "Injecting dependencies into MAVProxy venv..."
        pipx inject mavproxy future pexpect
    else
        # Fallback: try pip with --break-system-packages
        pip3 install --user --break-system-packages pymavlink mavproxy 2>/dev/null || \
        pip3 install --user pymavlink mavproxy 2>/dev/null || \
        print_warning "MAVProxy installation failed - install manually"
    fi
    ((STEP++))
    
    # ========================================================================
    # STEP 10: Verify Installation
    # ========================================================================
    print_step $STEP "Verifying installation"
    
    export PATH=$PATH:$ARDUPILOT_HOME/Tools/autotest:$HOME/.local/bin
    
    echo ""
    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"
    ((STEP++))

fi

# ============================================================================
# Make scripts executable
# ============================================================================
chmod +x "$SCRIPT_DIR/scripts/"*.sh 2>/dev/null || true

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

echo -e "${GREEN}Components installed:${NC}"
echo "  - ROS 2 $ROS_DISTRO"
echo "  - Gazebo $GZ_VERSION"
if [ "$INSTALL_ARDUPILOT" = true ]; then
    echo "  - ArduPilot SITL ($ARDUPILOT_HOME)"
    echo "  - ardupilot_gazebo ($ARDUPILOT_GZ)"
fi
echo "  - Python dependencies"
echo ""

if $IS_WSL; then
    echo -e "${YELLOW}WSL Notes:${NC}"
    echo "  - GUI requires WSLg (Win11) or VcXsrv (Win10)"
    echo "  - Use --software-render 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"
echo ""
echo -e "${GREEN}==========================================${NC}"