#!/bin/bash
set -e

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

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

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

# Detect WSL
IS_WSL=false
if grep -qEi "(microsoft|wsl)" /proc/version 2>/dev/null; then
    IS_WSL=true
    echo -e "${YELLOW}Running in WSL environment${NC}"
fi

# Detect ROS distro
ROS_DISTRO=""
for distro in humble jazzy iron galactic; do
    if [ -d "/opt/ros/$distro" ]; then
        ROS_DISTRO="$distro"
        break
    fi
done

if [ -z "$ROS_DISTRO" ]; then
    echo -e "${RED}ERROR: No ROS 2 installation found!${NC}"
    echo "Please install ROS 2 first. See docs/setup_guide.md or docs/wsl_setup_guide.md"
    exit 1
fi

echo -e "${GREEN}Using ROS 2: $ROS_DISTRO${NC}"

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

# Source Gazebo setup (CRITICAL for shader paths)
if [ -f /usr/share/gazebo/setup.bash ]; then
    source /usr/share/gazebo/setup.bash
elif [ -f /usr/share/gazebo-11/setup.bash ]; then
    source /usr/share/gazebo-11/setup.bash
fi

# Parse arguments
WORLD=""
USE_SOFTWARE_RENDER=false
HEADLESS=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --world)
            WORLD="$2"
            shift 2
            ;;
        --headless)
            HEADLESS=true
            shift
            ;;
        --no-gpu|--software-render)
            USE_SOFTWARE_RENDER=true
            shift
            ;;
        --help)
            echo "Usage: $0 [OPTIONS] [WORLD_FILE]"
            echo ""
            echo "Options:"
            echo "  --world FILE       Specify world file (default: empty_custom.world)"
            echo "  --headless         Run without GUI (gzserver only)"
            echo "  --software-render  Force software rendering (recommended for WSL)"
            echo "  --help             Show this help"
            echo ""
            echo "Examples:"
            echo "  $0                                    # Default world"
            echo "  $0 worlds/indoor_warehouse.world     # Indoor environment"
            echo "  $0 --software-render                 # For WSL without GPU"
            exit 0
            ;;
        *)
            if [ -f "$1" ]; then
                WORLD="$1"
            elif [ -f "$PROJECT_DIR/$1" ]; then
                WORLD="$PROJECT_DIR/$1"
            elif [ -f "$PROJECT_DIR/worlds/$1" ]; then
                WORLD="$PROJECT_DIR/worlds/$1"
            fi
            shift
            ;;
    esac
done

# Default world
if [ -z "$WORLD" ]; then
    WORLD="$PROJECT_DIR/worlds/empty_custom.world"
fi

# Validate world file
if [ ! -f "$WORLD" ]; then
    echo -e "${RED}ERROR: World file not found: $WORLD${NC}"
    echo "Available worlds:"
    ls -1 "$PROJECT_DIR/worlds/"*.world 2>/dev/null || echo "  (none found)"
    exit 1
fi

# Setup Gazebo paths
export GAZEBO_MODEL_PATH="$PROJECT_DIR/models:${GAZEBO_MODEL_PATH:-}"
export GAZEBO_RESOURCE_PATH="$PROJECT_DIR/worlds:${GAZEBO_RESOURCE_PATH:-}"

# ArduPilot Gazebo paths (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-specific setup
if $IS_WSL; then
    # Source WSL environment if exists
    if [ -f "$PROJECT_DIR/wsl_env.sh" ]; then
        source "$PROJECT_DIR/wsl_env.sh"
    fi
    
    # Set DISPLAY if not set
    if [ -z "$DISPLAY" ]; then
        if [ -d "/mnt/wslg" ]; then
            export DISPLAY=:0
        else
            export DISPLAY=$(cat /etc/resolv.conf | grep nameserver | awk '{print $2}'):0
        fi
    fi
    
    echo -e "${BLUE}DISPLAY: $DISPLAY${NC}"
    
    # Force software rendering in WSL by default if having issues
    if ! $USE_SOFTWARE_RENDER; then
        echo -e "${YELLOW}TIP: If Gazebo crashes or has rendering issues, use --software-render${NC}"
    fi
fi

# Apply software rendering if requested or in WSL with issues
if $USE_SOFTWARE_RENDER; then
    echo -e "${YELLOW}Using software rendering${NC}"
    export LIBGL_ALWAYS_SOFTWARE=1
    export MESA_GL_VERSION_OVERRIDE=3.3
    export MESA_GLSL_VERSION_OVERRIDE=330
fi

# Cleanup function
cleanup() {
    echo ""
    echo -e "${YELLOW}Shutting down simulation...${NC}"
    pkill -f "gzserver" 2>/dev/null || true
    pkill -f "gzclient" 2>/dev/null || true
    pkill -f "gazebo" 2>/dev/null || true
    sleep 1
    echo -e "${GREEN}Cleanup complete.${NC}"
}

trap cleanup EXIT INT TERM

# Activate virtual environment if exists
if [ -f "$PROJECT_DIR/venv/bin/activate" ]; then
    source "$PROJECT_DIR/venv/bin/activate"
fi

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

echo ""
echo -e "${GREEN}Starting simulation with world: $(basename $WORLD)${NC}"
echo -e "${BLUE}World path: $WORLD${NC}"
echo ""

# Check if ROS package exists
ROS_PKG_EXISTS=false
if ros2 pkg list 2>/dev/null | grep -q "uav_ugv_simulation"; then
    ROS_PKG_EXISTS=true
fi

# Launch simulation
if $ROS_PKG_EXISTS; then
    echo -e "${GREEN}Launching via ROS 2...${NC}"
    ros2 launch uav_ugv_simulation full_simulation.launch.py world:="$WORLD"
else
    echo -e "${YELLOW}ROS package not built. Launching Gazebo directly...${NC}"
    echo -e "${BLUE}To build the package: cd ~/ros2_ws && colcon build --packages-select uav_ugv_simulation${NC}"
    echo ""
    
    if $HEADLESS; then
        gzserver --verbose "$WORLD"
    else
        gazebo --verbose "$WORLD"
    fi
fi