#!/bin/bash

set -e

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

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

print_info() { echo -e "${CYAN}[INFO]${NC} $1"; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }

SOFTWARE_RENDER=auto
WORLD="uav_ugv_search.sdf"
MISSION="hover"
ALTITUDE=5.0
DURATION=30.0

while [[ $# -gt 0 ]]; do
    case $1 in
        --software-render) SOFTWARE_RENDER=true; shift ;;
        --no-software-render) SOFTWARE_RENDER=false; shift ;;
        --world) WORLD="$2"; shift 2 ;;
        --mission) MISSION="$2"; shift 2 ;;
        --altitude) ALTITUDE="$2"; shift 2 ;;
        --duration) DURATION="$2"; shift 2 ;;
        -h|--help)
            echo "Usage: $0 [options]"
            echo "  --software-render      Force software rendering"
            echo "  --no-software-render   Disable software rendering"
            echo "  --world <file>         World file (default: uav_ugv_search.sdf)"
            echo "  --mission <type>       hover, square, search (default: hover)"
            echo "  --altitude <m>         Takeoff altitude (default: 5.0)"
            echo "  --duration <s>         Hover duration (default: 30.0)"
            exit 0
            ;;
        *) shift ;;
    esac
done

cleanup_all() {
    print_info "Cleaning up ..."
    pkill -f "gz sim" 2>/dev/null || true
    pkill -f "arducopter" 2>/dev/null || true
    pkill -f "sim_vehicle.py" 2>/dev/null || true
    pkill -f "mavproxy" 2>/dev/null || true
    pkill -f "main.py" 2>/dev/null || true
}

cleanup_sitl() {
    pkill -f "arducopter" 2>/dev/null || true
    pkill -f "sim_vehicle.py" 2>/dev/null || true
    pkill -f "mavproxy" 2>/dev/null || true
    pkill -f "main.py" 2>/dev/null || true
}
trap cleanup_all EXIT

export PATH=$PATH:$HOME/ardupilot/Tools/autotest:$HOME/.local/bin
export GZ_SIM_RESOURCE_PATH="$PROJECT_DIR/models:$PROJECT_DIR/worlds:$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds"
export GZ_SIM_SYSTEM_PLUGIN_PATH=$HOME/ardupilot_gazebo/build

if [ "$SOFTWARE_RENDER" = "auto" ]; then
    if grep -qi microsoft /proc/version 2>/dev/null; then
        print_info "WSL detected -> software rendering"
        SOFTWARE_RENDER=true
    else
        SOFTWARE_RENDER=false
    fi
fi

if [ "$SOFTWARE_RENDER" = "true" ]; then
    export LIBGL_ALWAYS_SOFTWARE=1
    export GALLIUM_DRIVER=llvmpipe
    export MESA_GL_VERSION_OVERRIDE=3.3
fi

cleanup_all 2>/dev/null

WORLD_FILE=""
if [ -f "$PROJECT_DIR/worlds/$WORLD" ]; then
    WORLD_FILE="$PROJECT_DIR/worlds/$WORLD"
elif [ -f "$HOME/ardupilot_gazebo/worlds/$WORLD" ]; then
    WORLD_FILE="$HOME/ardupilot_gazebo/worlds/$WORLD"
elif [ -f "$WORLD" ]; then
    WORLD_FILE="$WORLD"
else
    print_error "World file not found: $WORLD"
    exit 1
fi

print_info "==================================="
print_info "  UAV-UGV Simulation"
print_info "==================================="
print_info "World:    $WORLD_FILE"
print_info "Mission:  $MISSION"
print_info "Altitude: ${ALTITUDE}m"
echo ""

print_info "[1/3] Starting Gazebo ..."
gz sim -v4 -r "$WORLD_FILE" &
GZ_PID=$!
sleep 10

if ! kill -0 $GZ_PID 2>/dev/null; then
    print_error "Gazebo failed to start"
    exit 1
fi
print_success "Gazebo running (PID: $GZ_PID)"

print_info "[2/3] Starting ArduPilot SITL ..."
cd ~/ardupilot

SITL_ARGS="-v ArduCopter -f gazebo-iris --model JSON -I0"
SITL_ARGS="$SITL_ARGS --no-mavproxy"

PARAM_FILE="$PROJECT_DIR/config/ardupilot_gps_denied.parm"
if [ -f "$PARAM_FILE" ]; then
    print_info "Loading params: $PARAM_FILE"
    SITL_ARGS="$SITL_ARGS --add-param-file $PARAM_FILE"
fi

sim_vehicle.py $SITL_ARGS &
SITL_PID=$!

print_info "Waiting for SITL (~20s) ..."
sleep 20

if ! pgrep -f "arducopter" > /dev/null 2>&1; then
    print_error "ArduPilot SITL failed to start"
    exit 1
fi
print_success "ArduPilot SITL running (TCP 5760)"

print_info "[3/3] Starting main.py ..."
cd "$PROJECT_DIR"
sleep 3

python3 src/main.py \
    --device sim \
    --connection "tcp:127.0.0.1:5760" \
    --mission "$MISSION" \
    --altitude "$ALTITUDE" \
    --duration "$DURATION" &
MAIN_PID=$!
print_success "main.py running (PID: $MAIN_PID)"

print_info ""
print_info "==================================="
print_info "  Simulation Running"
print_info "==================================="
print_info "  Gazebo -> ArduPilot SITL (TCP:5760)"
print_info "              |"
print_info "          main.py (pymavlink)"
print_info ""
print_info "  Mission: $MISSION  Alt: ${ALTITUDE}m"
print_info "  Press Ctrl+C to stop"
print_info "==================================="

wait $MAIN_PID 2>/dev/null

cleanup_sitl

print_info ""
print_info "=================================="
print_info "  Mission Complete!"
print_info "=================================="
print_info "Gazebo GUI still open."
print_info "Press Ctrl+C to exit."
print_info "=================================="
print_info ""

wait $GZ_PID 2>/dev/null