#!/bin/bash
# Autonomous UAV-UGV Simulation Runner
# Launches Gazebo + ArduPilot SITL + Autonomous Controller

set -e

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

# Colors
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"; }

# Parse arguments
SOFTWARE_RENDER=false
WORLD="iris_runway.sdf"
MISSION="hover"  # hover, square, circle

while [[ $# -gt 0 ]]; do
    case $1 in
        --software-render) SOFTWARE_RENDER=true; shift ;;
        --world) WORLD="$2"; shift 2 ;;
        --mission) MISSION="$2"; shift 2 ;;
        -h|--help)
            echo "Usage: $0 [options]"
            echo "  --software-render   Use software rendering (WSL/no GPU)"
            echo "  --world <file>      World file to load (default: iris_runway.sdf)"
            echo "  --mission <type>    Mission type: hover, square, circle (default: hover)"
            exit 0
            ;;
        *) shift ;;
    esac
done

# Cleanup function
cleanup() {
    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 "autonomous_controller.py" 2>/dev/null || true
}
trap cleanup EXIT

# Setup environment
export PATH=$PATH:$HOME/ardupilot/Tools/autotest:$HOME/.local/bin
export GZ_SIM_RESOURCE_PATH=$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds
export GZ_SIM_SYSTEM_PLUGIN_PATH=$HOME/ardupilot_gazebo/build

if [ "$SOFTWARE_RENDER" = true ]; then
    print_info "Using software rendering (Mesa)"
    export LIBGL_ALWAYS_SOFTWARE=1
    export GALLIUM_DRIVER=llvmpipe
fi

# Kill any existing processes
cleanup 2>/dev/null

print_info "==================================="
print_info "  Autonomous UAV-UGV Simulation"
print_info "==================================="
print_info "World: $WORLD"
print_info "Mission: $MISSION"
echo ""

# Step 1: Start Gazebo
print_info "Starting Gazebo Harmonic..."
gz sim -v4 -r $HOME/ardupilot_gazebo/worlds/$WORLD &
GZ_PID=$!
sleep 8

# Check if Gazebo started
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)"

# Step 2: Start ArduPilot SITL
print_info "Starting ArduPilot SITL..."
cd ~/ardupilot

# Build SITL command
# We enable MAVProxy (--console --map) for visualization
# MAVProxy will sit between SITL and our controller
# SITL -> TCP:5760 -> MAVProxy -> UDP:14550 (GCS)
#                              -> UDP:14551 (Autonomous Controller)
SITL_ARGS="-v ArduCopter -f gazebo-iris --model JSON -I0"
SITL_ARGS="$SITL_ARGS --console --map"

# Add output ports
# 14550: External GCS (QGroundControl)
# 14551: Autonomous Controller
SITL_ARGS="$SITL_ARGS --out 127.0.0.1:14550 --out 127.0.0.1:14551"

# Check if custom param file exists
PARAM_FILE="$PROJECT_DIR/config/ardupilot_gps_denied.parm"
if [ -f "$PARAM_FILE" ]; then
    print_info "Loading custom parameters..."
    SITL_ARGS="$SITL_ARGS --add-param-file $PARAM_FILE"
fi

print_info "SITL will output status to this terminal"
print_info "Connect GCS to udp:127.0.0.1:14550 for monitoring"

# Start SITL
sim_vehicle.py $SITL_ARGS &
SITL_PID=$!

# Wait for SITL to initialize
print_info "Waiting for SITL to initialize (~20 seconds)..."
sleep 20

# Check if SITL started (check for arducopter process, not sim_vehicle.py)
if ! pgrep -f "arducopter" > /dev/null 2>&1; then
    print_error "ArduPilot SITL failed to start"
    exit 1
fi
print_success "ArduPilot SITL running"

# Step 3: Start the autonomous controller
print_info "Starting Autonomous Controller..."
print_info "Mission: $MISSION"

# Wait a bit more, then start controller with retry
sleep 5

python3 "$PROJECT_DIR/src/autonomous_controller.py" --mission "$MISSION" --connection "udp:127.0.0.1:14551" &
CTRL_PID=$!

print_success "Autonomous Controller started (PID: $CTRL_PID)"
print_info ""
print_info "==================================="
print_info "  Simulation Running"
print_info "==================================="
print_info "The UAV will automatically:"
print_info "  1. Wait for EKF initialization (~30-60s)"
print_info "  2. Arm and enter GUIDED mode"
print_info "  3. Take off to 5 meters"
print_info "  4. Execute mission: $MISSION"
print_info ""
print_info "Press Ctrl+C to stop"
print_info "==================================="

# Wait for controller to finish or user interrupt
wait $CTRL_PID