#!/bin/bash
set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

print_header() {
    echo -e "${BLUE}==========================================${NC}"
    echo -e "${BLUE}  UAV-UGV Simulation Setup${NC}"
    echo -e "${BLUE}  GPS-Denied Navigation with Geofencing${NC}"
    echo -e "${BLUE}==========================================${NC}"
    echo ""
}

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

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

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

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

# Detect environment
detect_environment() {
    IS_WSL=false
    IS_WSL2=false
    UBUNTU_VERSION=""
    ROS_DISTRO=""
    
    # Detect WSL
    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"
    fi
    
    # Determine ROS distro based on Ubuntu version
    case "$UBUNTU_VERSION" in
        "22.04")
            ROS_DISTRO="humble"
            ;;
        "24.04")
            ROS_DISTRO="jazzy"
            ;;
        "20.04")
            ROS_DISTRO="galactic"
            ;;
        *)
            ROS_DISTRO="humble"  # Default fallback
            ;;
    esac
    
    echo -e "${BLUE}Detected Environment:${NC}"
    echo "  - Ubuntu: $UBUNTU_VERSION ($UBUNTU_CODENAME)"
    echo "  - WSL: $IS_WSL"
    echo "  - WSL2: $IS_WSL2"
    echo "  - ROS 2 Distro: $ROS_DISTRO"
    echo ""
}

# Check if ROS 2 is installed
check_ros2_installed() {
    if [ -d "/opt/ros/$ROS_DISTRO" ]; then
        return 0
    else
        return 1
    fi
}

# Install ROS 2 repository
install_ros2_repo() {
    print_info "Setting up ROS 2 repository..."
    
    sudo apt-get install -y software-properties-common curl gnupg lsb-release
    
    # 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
    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
}

# Setup WSL-specific configurations
setup_wsl() {
    print_info "Configuring WSL-specific settings..."
    
    # Create WSL environment setup
    WSL_ENV_FILE="$SCRIPT_DIR/wsl_env.sh"
    
    cat > "$WSL_ENV_FILE" << 'WSLEOF'
#!/bin/bash
# WSL-specific environment variables

# Detect WSL version and set DISPLAY
if grep -qi "wsl2" /proc/version 2>/dev/null || [ -f /run/WSL ]; then
    # WSL2 with WSLg (Windows 11)
    if [ -d "/mnt/wslg" ]; then
        export DISPLAY=:0
    else
        # WSL2 without WSLg (Windows 10) - use X server
        export DISPLAY=$(cat /etc/resolv.conf | grep nameserver | awk '{print $2}'):0
    fi
else
    # WSL1
    export DISPLAY=localhost:0
fi

# Performance settings for Gazebo in WSL
export LIBGL_ALWAYS_INDIRECT=0

# If GPU acceleration isn't working, uncomment this:
# export LIBGL_ALWAYS_SOFTWARE=1

# Mesa driver settings (helps with some rendering issues)
export MESA_GL_VERSION_OVERRIDE=3.3
export MESA_GLSL_VERSION_OVERRIDE=330

# Gazebo specific
export OGRE_RTT_MODE=Copy
WSLEOF

    chmod +x "$WSL_ENV_FILE"
    
    # Add to bashrc if not already present
    if ! grep -q "wsl_env.sh" ~/.bashrc 2>/dev/null; then
        echo "" >> ~/.bashrc
        echo "# WSL environment for UAV-UGV simulation" >> ~/.bashrc
        echo "if [ -f \"$WSL_ENV_FILE\" ]; then source \"$WSL_ENV_FILE\"; fi" >> ~/.bashrc
    fi
    
    print_info "WSL environment configured. Source ~/.bashrc to apply."
}

# Setup X11 for GUI applications
setup_x11_wsl() {
    print_info "Installing X11 utilities for GUI support..."
    
    sudo apt-get install -y \
        x11-apps \
        x11-xserver-utils \
        dbus-x11 \
        libgl1-mesa-glx \
        mesa-utils 2>/dev/null || true
}

# Main setup
print_header

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

# Detect environment
detect_environment

# Determine total steps based on what's needed
if check_ros2_installed; then
    TOTAL_STEPS=7
else
    TOTAL_STEPS=9
    print_warning "ROS 2 $ROS_DISTRO not found. Will attempt to install."
fi

if $IS_WSL; then
    TOTAL_STEPS=$((TOTAL_STEPS + 1))
fi

STEP=1

# Step: Update system
print_step $STEP "Updating system packages..."
sudo apt-get update
((STEP++))

# Step: WSL-specific setup
if $IS_WSL; then
    print_step $STEP "Setting up WSL environment..."
    setup_wsl
    setup_x11_wsl
    ((STEP++))
fi

# Step: Install ROS 2 if not present
if ! check_ros2_installed; then
    print_step $STEP "Installing ROS 2 repository..."
    install_ros2_repo
    ((STEP++))
    
    print_step $STEP "Installing ROS 2 $ROS_DISTRO..."
    sudo apt-get install -y ros-${ROS_DISTRO}-desktop || {
        print_warning "Failed to install ros-${ROS_DISTRO}-desktop, trying base..."
        sudo apt-get install -y ros-${ROS_DISTRO}-ros-base
    }
    ((STEP++))
fi

# Step: Install system dependencies
print_step $STEP "Installing system dependencies..."
sudo apt-get install -y \
    python3-pip \
    python3-venv \
    python3-opencv \
    libopencv-dev \
    python3-colcon-common-extensions \
    build-essential \
    cmake \
    git || true
((STEP++))

# Step: Install ROS 2 packages
print_step $STEP "Installing ROS 2 packages..."
ROS_PACKAGES=(
    "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"
)

# Gazebo packages differ by distro
if [ "$ROS_DISTRO" = "humble" ]; then
    ROS_PACKAGES+=("ros-${ROS_DISTRO}-gazebo-ros-pkgs")
elif [ "$ROS_DISTRO" = "jazzy" ]; then
    # Jazzy uses Gazebo Harmonic (gz-sim)
    ROS_PACKAGES+=("ros-${ROS_DISTRO}-ros-gz")
fi

# Install available packages (some may not exist for all distros)
for pkg in "${ROS_PACKAGES[@]}"; do
    sudo apt-get install -y "$pkg" 2>/dev/null || {
        print_warning "Package $pkg not available, skipping..."
    }
done
((STEP++))

# Step: Install MAVROS GeographicLib datasets
print_step $STEP "Installing MAVROS GeographicLib datasets..."
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
        sudo "$GEOGRAPHICLIB_SCRIPT" || print_warning "GeographicLib datasets installation failed"
    else
        print_info "GeographicLib datasets already installed"
    fi
else
    print_warning "MAVROS not installed, skipping GeographicLib datasets"
fi
((STEP++))

# Step: Create Python virtual environment
print_step $STEP "Creating Python virtual environment..."
if [ ! -d "venv" ]; then
    python3 -m venv venv
fi
source venv/bin/activate

# Upgrade pip and install dependencies
pip install --upgrade pip
pip install -r requirements.txt || {
    print_warning "Some Python packages failed to install. Check requirements.txt"
}
((STEP++))

# Step: Build ROS 2 package
print_step $STEP "Building ROS 2 package..."
if [ -f "/opt/ros/${ROS_DISTRO}/setup.bash" ]; then
    source /opt/ros/${ROS_DISTRO}/setup.bash
    
    # Try to build if in a ROS workspace
    if [ -d "$SCRIPT_DIR/../src" ] || [ -f "$SCRIPT_DIR/package.xml" ]; then
        cd "$SCRIPT_DIR/.."
        colcon build --packages-select uav_ugv_simulation --symlink-install 2>/dev/null || {
            print_warning "Colcon build skipped. To build manually:"
            print_info "  cd ~/ros2_ws && colcon build --packages-select uav_ugv_simulation"
        }
        cd "$SCRIPT_DIR"
    fi
else
    print_warning "ROS 2 not found, skipping package build"
fi
((STEP++))

# Step: Make scripts executable and create activation script
print_step $STEP "Finalizing setup..."
chmod +x scripts/*.sh 2>/dev/null || true

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

# Activate Python virtual environment
source "\$SCRIPT_DIR/venv/bin/activate"

# Source ROS 2
if [ -f "/opt/ros/${ROS_DISTRO}/setup.bash" ]; then
    source /opt/ros/${ROS_DISTRO}/setup.bash
fi

# Source workspace if built
if [ -f "\$SCRIPT_DIR/../install/setup.bash" ]; then
    source "\$SCRIPT_DIR/../install/setup.bash"
elif [ -f "\$SCRIPT_DIR/install/setup.bash" ]; then
    source "\$SCRIPT_DIR/install/setup.bash"
fi

# Gazebo paths
export GAZEBO_MODEL_PATH="\$SCRIPT_DIR/models:\$GAZEBO_MODEL_PATH"
export GAZEBO_RESOURCE_PATH="\$SCRIPT_DIR/worlds:\$GAZEBO_RESOURCE_PATH"

# ArduPilot Gazebo (if installed)
if [ -d "\$HOME/ardupilot_gazebo" ]; then
    export GAZEBO_MODEL_PATH="\$HOME/ardupilot_gazebo/models:\$GAZEBO_MODEL_PATH"
    export GAZEBO_RESOURCE_PATH="\$HOME/ardupilot_gazebo/worlds:\$GAZEBO_RESOURCE_PATH"
fi

# WSL environment (if applicable)
if [ -f "\$SCRIPT_DIR/wsl_env.sh" ]; then
    source "\$SCRIPT_DIR/wsl_env.sh"
fi

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

# Summary
echo ""
echo -e "${GREEN}==========================================${NC}"
echo -e "${GREEN}  Setup Complete!${NC}"
echo -e "${GREEN}==========================================${NC}"
echo ""

if $IS_WSL; then
    echo -e "${YELLOW}WSL Setup Notes:${NC}"
    echo "  - WSL environment file created: wsl_env.sh"
    echo "  - For GUI apps, ensure X server is running (Windows 11 has WSLg built-in)"
    echo "  - See docs/wsl_setup_guide.md for detailed instructions"
    echo ""
fi

echo -e "${BLUE}Next steps:${NC}"
echo "  1. source ~/.bashrc"
echo "  2. source activate_venv.sh"
echo "  3. bash scripts/run_simulation.sh"
echo ""
echo -e "${BLUE}GPS-Denied Navigation:${NC}"
echo "  - All navigation uses LOCAL coordinates"
echo "  - GPS is ONLY used for geofencing"
echo ""

if ! check_ros2_installed && [ ! -f "/opt/ros/${ROS_DISTRO}/setup.bash" ]; then
    echo -e "${YELLOW}IMPORTANT:${NC}"
    echo "  ROS 2 installation may have failed. Please install manually:"
    echo "  See: https://docs.ros.org/en/${ROS_DISTRO}/Installation.html"
    echo ""
fi