#!/bin/bash
# WSL Quick Setup Script
# This script sets up the simulation environment specifically for WSL2

set -e

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

echo -e "${BLUE}==========================================${NC}"
echo -e "${BLUE}  UAV-UGV Simulation - WSL Setup${NC}"
echo -e "${BLUE}==========================================${NC}"
echo ""

# Check if running in WSL
if ! grep -qEi "(microsoft|wsl)" /proc/version 2>/dev/null; then
    echo -e "${YELLOW}This script is designed for WSL. Running anyway...${NC}"
fi

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

# Detect Ubuntu version
. /etc/os-release
echo -e "${BLUE}Detected: Ubuntu $VERSION_ID ($VERSION_CODENAME)${NC}"

# Determine ROS distro
case "$VERSION_ID" in
    "22.04") ROS_DISTRO="humble" ;;
    "24.04") ROS_DISTRO="jazzy" ;;
    *) ROS_DISTRO="humble"; echo -e "${YELLOW}Unknown Ubuntu version, defaulting to Humble${NC}" ;;
esac

echo -e "${BLUE}Target ROS 2 distro: $ROS_DISTRO${NC}"
echo ""

# Step 1: Update and install prerequisites
echo -e "${GREEN}[1/6] Installing prerequisites...${NC}"
sudo apt-get update
sudo apt-get install -y \
    software-properties-common \
    curl \
    gnupg \
    lsb-release \
    x11-apps \
    x11-xserver-utils \
    dbus-x11 \
    mesa-utils \
    libgl1-mesa-glx

# Step 2: Add ROS 2 repository
echo -e "${GREEN}[2/6] Adding ROS 2 repository...${NC}"
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

if [ ! -f /etc/apt/sources.list.d/ros2.list ]; then
    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
fi

sudo apt-get update

# Step 3: Install ROS 2
echo -e "${GREEN}[3/6] Installing ROS 2 $ROS_DISTRO...${NC}"
if [ ! -d "/opt/ros/$ROS_DISTRO" ]; then
    sudo apt-get install -y ros-${ROS_DISTRO}-desktop python3-colcon-common-extensions || {
        echo -e "${YELLOW}Desktop install failed, trying base...${NC}"
        sudo apt-get install -y ros-${ROS_DISTRO}-ros-base
    }
else
    echo -e "${BLUE}ROS 2 $ROS_DISTRO already installed${NC}"
fi

# Step 4: Install ROS packages
echo -e "${GREEN}[4/6] Installing ROS 2 packages...${NC}"
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 2>/dev/null || {
    echo -e "${YELLOW}Some packages unavailable for $ROS_DISTRO${NC}"
}

# Gazebo (different for Humble vs Jazzy)
if [ "$ROS_DISTRO" = "humble" ]; then
    sudo apt-get install -y ros-humble-gazebo-ros-pkgs 2>/dev/null || true
elif [ "$ROS_DISTRO" = "jazzy" ]; then
    sudo apt-get install -y ros-jazzy-ros-gz 2>/dev/null || true
fi

# GeographicLib datasets
echo -e "${GREEN}[5/6] Installing GeographicLib datasets...${NC}"
GEOGRAPHICLIB_SCRIPT="/opt/ros/${ROS_DISTRO}/lib/mavros/install_geographiclib_datasets.sh"
if [ -f "$GEOGRAPHICLIB_SCRIPT" ] && [ ! -f /usr/share/GeographicLib/geoids/egm96-5.pgm ]; then
    sudo "$GEOGRAPHICLIB_SCRIPT" || echo -e "${YELLOW}GeographicLib install failed${NC}"
fi

# Step 5: Setup Python environment
echo -e "${GREEN}[6/6] Setting up Python environment...${NC}"
cd "$PROJECT_DIR"

sudo apt-get install -y python3-pip python3-venv python3-opencv libopencv-dev

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

source venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt 2>/dev/null || echo -e "${YELLOW}Some pip packages failed${NC}"

# Create WSL environment file
cat > "$PROJECT_DIR/wsl_env.sh" << 'EOF'
#!/bin/bash
# WSL Environment Setup

# DISPLAY configuration
if [ -d "/mnt/wslg" ]; then
    # WSLg (Windows 11)
    export DISPLAY=:0
else
    # X server (Windows 10)
    export DISPLAY=$(cat /etc/resolv.conf | grep nameserver | awk '{print $2}'):0
fi

# OpenGL settings for WSL
export LIBGL_ALWAYS_INDIRECT=0
export MESA_GL_VERSION_OVERRIDE=3.3

# Uncomment for software rendering if GPU issues:
# export LIBGL_ALWAYS_SOFTWARE=1

# Gazebo performance
export OGRE_RTT_MODE=Copy
EOF
chmod +x "$PROJECT_DIR/wsl_env.sh"

# Add to bashrc
if ! grep -q "wsl_env.sh" ~/.bashrc 2>/dev/null; then
    echo "" >> ~/.bashrc
    echo "# UAV-UGV Simulation WSL environment" >> ~/.bashrc
    echo "[ -f \"$PROJECT_DIR/wsl_env.sh\" ] && source \"$PROJECT_DIR/wsl_env.sh\"" >> ~/.bashrc
    echo "[ -f /opt/ros/${ROS_DISTRO}/setup.bash ] && source /opt/ros/${ROS_DISTRO}/setup.bash" >> ~/.bashrc
fi

# Make scripts executable
chmod +x "$PROJECT_DIR/scripts/"*.sh 2>/dev/null || true
chmod +x "$PROJECT_DIR/setup.sh" 2>/dev/null || true
chmod +x "$PROJECT_DIR/activate_venv.sh" 2>/dev/null || true

echo ""
echo -e "${GREEN}==========================================${NC}"
echo -e "${GREEN}  WSL Setup Complete!${NC}"
echo -e "${GREEN}==========================================${NC}"
echo ""
echo -e "${BLUE}Next steps:${NC}"
echo "  1. Close and reopen your terminal (or run: source ~/.bashrc)"
echo "  2. Test GUI: xcalc (should open calculator)"
echo "  3. Test Gazebo: gazebo --verbose"
echo "  4. Run simulation: source activate_venv.sh && bash scripts/run_simulation.sh"
echo ""
echo -e "${YELLOW}WSL Tips:${NC}"
echo "  - Windows 11: GUI works out of the box (WSLg)"
echo "  - Windows 10: Install VcXsrv and run XLaunch first"
echo "  - Slow graphics? Run with: bash scripts/run_simulation.sh --software-render"
echo "  - See docs/wsl_setup_guide.md for detailed help"
echo ""