#!/bin/bash
# =============================================================================
# UAV-UGV Simulation - Run Script
# =============================================================================
# Launches Gazebo with ArduPilot SITL for GPS-denied navigation testing
#
# Usage:
#   ./scripts/run_simulation.sh                     # Default world
#   ./scripts/run_simulation.sh --world runway      # Runway world
#   ./scripts/run_simulation.sh --software-render   # For WSL/no GPU
# =============================================================================

set -e

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 jazzy humble iron galactic; do
    if [ -d "/opt/ros/$distro" ]; then
        ROS_DISTRO="$distro"
        break
    fi
done

if [ -n "$ROS_DISTRO" ]; then
    source /opt/ros/${ROS_DISTRO}/setup.bash
    echo -e "${GREEN}ROS 2: $ROS_DISTRO${NC}"
fi

# ArduPilot paths
ARDUPILOT_HOME="${ARDUPILOT_HOME:-$HOME/ardupilot}"
ARDUPILOT_GZ="${ARDUPILOT_GZ:-$HOME/ardupilot_gazebo}"

# Parse arguments
WORLD_NAME="iris_runway"
VEHICLE="ArduCopter"
USE_SOFTWARE_RENDER=false
INSTANCE=0

while [[ $# -gt 0 ]]; do
    case $1 in
        --world)
            WORLD_NAME="$2"
            shift 2
            ;;
        --vehicle)
            VEHICLE="$2"
            shift 2
            ;;
        --software-render)
            USE_SOFTWARE_RENDER=true
            shift
            ;;
        --instance)
            INSTANCE="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --world NAME        World name (default: iris_runway)"
            echo "  --vehicle TYPE      ArduCopter or Rover (default: ArduCopter)"
            echo "  --software-render   Force software rendering (for WSL)"
            echo "  --instance N        SITL instance number (default: 0)"
            echo "  --help              Show this help"
            echo ""
            echo "Available worlds in ardupilot_gazebo:"
            ls "$ARDUPILOT_GZ/worlds/"*.sdf 2>/dev/null | xargs -I {} basename {} .sdf || echo "  (none found)"
            exit 0
            ;;
        *)
            WORLD_NAME="$1"
            shift
            ;;
    esac
done

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

if [ -d "$ARDUPILOT_GZ" ]; then
    export GZ_SIM_SYSTEM_PLUGIN_PATH="$ARDUPILOT_GZ/build:${GZ_SIM_SYSTEM_PLUGIN_PATH:-}"
    export GZ_SIM_RESOURCE_PATH="$ARDUPILOT_GZ/models:$ARDUPILOT_GZ/worlds:$GZ_SIM_RESOURCE_PATH"
fi

# Find world file
WORLD_FILE=""
if [ -f "$PROJECT_DIR/worlds/${WORLD_NAME}.sdf" ]; then
    WORLD_FILE="$PROJECT_DIR/worlds/${WORLD_NAME}.sdf"
elif [ -f "$ARDUPILOT_GZ/worlds/${WORLD_NAME}.sdf" ]; then
    WORLD_FILE="$ARDUPILOT_GZ/worlds/${WORLD_NAME}.sdf"
elif [ -f "$WORLD_NAME" ]; then
    WORLD_FILE="$WORLD_NAME"
else
    echo -e "${RED}ERROR: World file not found: $WORLD_NAME${NC}"
    echo ""
    echo "Available worlds:"
    echo "  Project worlds:"
    ls "$PROJECT_DIR/worlds/"*.sdf 2>/dev/null | xargs -I {} basename {} .sdf || echo "    (none)"
    echo "  ArduPilot Gazebo worlds:"
    ls "$ARDUPILOT_GZ/worlds/"*.sdf 2>/dev/null | xargs -I {} basename {} .sdf || echo "    (none)"
    exit 1
fi

# WSL/Display setup
if $IS_WSL; then
    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}"
fi

if $USE_SOFTWARE_RENDER; then
    echo -e "${YELLOW}Using software rendering${NC}"
    export LIBGL_ALWAYS_SOFTWARE=1
    export MESA_GL_VERSION_OVERRIDE=3.3
fi

# Cleanup function
cleanup() {
    echo ""
    echo -e "${YELLOW}Shutting down simulation...${NC}"
    
    pkill -f "sim_vehicle.py" 2>/dev/null || true
    pkill -f "mavproxy" 2>/dev/null || true
    pkill -f "arducopter" 2>/dev/null || true
    pkill -f "ardurover" 2>/dev/null || true
    pkill -f "ruby" 2>/dev/null || true  # gz sim uses ruby
    pkill -f "gz sim" 2>/dev/null || true
    
    sleep 2
    echo -e "${GREEN}Cleanup complete.${NC}"
}

trap cleanup EXIT INT TERM

# Check ArduPilot installation
if [ ! -d "$ARDUPILOT_HOME" ]; then
    echo -e "${RED}ERROR: ArduPilot not found at $ARDUPILOT_HOME${NC}"
    echo "Run setup.sh first to install ArduPilot"
    exit 1
fi

# Check Gazebo plugin
if [ ! -f "$ARDUPILOT_GZ/build/libArduPilotPlugin.so" ]; then
    echo -e "${RED}ERROR: ArduPilot Gazebo plugin not found${NC}"
    echo "Run setup.sh first to build the plugin"
    exit 1
fi

echo ""
echo -e "${GREEN}Configuration:${NC}"
echo "  Vehicle: $VEHICLE"
echo "  World: $(basename $WORLD_FILE .sdf)"
echo "  Instance: $INSTANCE"
echo ""

# Start Gazebo in background
echo -e "${GREEN}Starting Gazebo...${NC}"
gz sim -v4 -r "$WORLD_FILE" &
GZ_PID=$!

sleep 5

# Check if Gazebo started
if ! kill -0 $GZ_PID 2>/dev/null; then
    echo -e "${RED}ERROR: Gazebo failed to start${NC}"
    exit 1
fi

echo -e "${GREEN}Gazebo running (PID: $GZ_PID)${NC}"

# Start ArduPilot SITL
echo -e "${GREEN}Starting ArduPilot SITL...${NC}"
cd "$ARDUPILOT_HOME"

export PATH=$PATH:$ARDUPILOT_HOME/Tools/autotest:$HOME/.local/bin

sim_vehicle.py \
    -v $VEHICLE \
    -f gazebo-iris \
    --model JSON \
    -I $INSTANCE \
    --console \
    &
SITL_PID=$!

sleep 10

echo ""
echo -e "${GREEN}==========================================${NC}"
echo -e "${GREEN}  Simulation Running${NC}"
echo -e "${GREEN}==========================================${NC}"
echo ""
echo -e "${BLUE}Components:${NC}"
echo "  - Gazebo: PID $GZ_PID"
echo "  - ArduPilot SITL: PID $SITL_PID"
echo ""
echo -e "${BLUE}MAVProxy Commands (in console):${NC}"
echo "  mode guided      # Switch to GUIDED mode"
echo "  arm throttle     # Arm the drone"
echo "  takeoff 5        # Takeoff to 5 meters"
echo "  guided 10 5 -10  # Fly to position (N, E, D)"
echo "  rtl              # Return to launch"
echo "  land             # Land"
echo ""
echo -e "${YELLOW}Press Ctrl+C to stop simulation${NC}"
echo ""

# Wait for processes
wait