#!/bin/bash
# Full simulation launch script with ArduPilot SITL
# This script launches Gazebo, ArduPilot SITL, and MAVROS

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 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"
    exit 1
fi

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

# Check ArduPilot installation
ARDUPILOT_HOME="${ARDUPILOT_HOME:-$HOME/ardupilot}"
if [ ! -d "$ARDUPILOT_HOME" ]; then
    echo -e "${RED}ERROR: ArduPilot not found at $ARDUPILOT_HOME${NC}"
    echo ""
    echo "Please install ArduPilot first:"
    echo "  bash scripts/install_ardupilot.sh"
    echo ""
    exit 1
fi

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

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="$PROJECT_DIR/worlds/iris_runway.world"
VEHICLE="copter"
USE_SOFTWARE_RENDER=false
INSTANCE=0

while [[ $# -gt 0 ]]; do
    case $1 in
        --world)
            WORLD="$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 FILE        World file (default: iris_runway.world)"
            echo "  --vehicle TYPE      Vehicle type: copter, rover (default: copter)"
            echo "  --software-render   Force software rendering (for WSL)"
            echo "  --instance N        Vehicle instance number (default: 0)"
            echo "  --help              Show this help"
            exit 0
            ;;
        *)
            if [ -f "$1" ]; then
                WORLD="$1"
            elif [ -f "$PROJECT_DIR/worlds/$1" ]; then
                WORLD="$PROJECT_DIR/worlds/$1"
            fi
            shift
            ;;
    esac
done

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

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/Display setup
if $IS_WSL; then
    if [ -f "$PROJECT_DIR/wsl_env.sh" ]; then
        source "$PROJECT_DIR/wsl_env.sh"
    fi
    if [ -z "$DISPLAY" ]; then
        export DISPLAY=:0
    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}"
    
    # Kill processes
    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 "gzserver" 2>/dev/null || true
    pkill -f "gzclient" 2>/dev/null || true
    pkill -f "gazebo" 2>/dev/null || true
    pkill -f "mavros" 2>/dev/null || true
    
    sleep 2
    echo -e "${GREEN}Cleanup complete.${NC}"
}

trap cleanup EXIT INT TERM

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

# Use ardupilot_gazebo iris world if no custom world specified
if [ ! -f "$WORLD" ]; then
    if [ -f "$HOME/ardupilot_gazebo/worlds/iris_arducopter_runway.world" ]; then
        WORLD="$HOME/ardupilot_gazebo/worlds/iris_arducopter_runway.world"
    else
        echo -e "${RED}ERROR: World file not found: $WORLD${NC}"
        exit 1
    fi
fi

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

# Calculate ports
MAVLINK_PORT=$((14550 + INSTANCE * 10))
SITL_PORT=$((5760 + INSTANCE * 10))

# Start Gazebo in background
echo -e "${GREEN}Starting Gazebo...${NC}"
gazebo --verbose "$WORLD" &
GAZEBO_PID=$!

sleep 5

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

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

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

if [ "$VEHICLE" = "copter" ]; then
    VEHICLE_DIR="ArduCopter"
    FRAME="gazebo-iris"
elif [ "$VEHICLE" = "rover" ]; then
    VEHICLE_DIR="Rover"
    FRAME="gazebo-rover"
else
    VEHICLE_DIR="ArduCopter"
    FRAME="gazebo-iris"
fi

cd "$ARDUPILOT_HOME/$VEHICLE_DIR"

# Start sim_vehicle.py in background
python3 "$ARDUPILOT_HOME/Tools/autotest/sim_vehicle.py" \
    -v "$VEHICLE_DIR" \
    -f "$FRAME" \
    --no-mavproxy \
    -I "$INSTANCE" \
    --no-rebuild \
    &
SITL_PID=$!

sleep 10

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

echo -e "${GREEN}ArduPilot SITL running (PID: $SITL_PID)${NC}"

# Start MAVROS
echo -e "${GREEN}Starting MAVROS...${NC}"
ros2 run mavros mavros_node --ros-args \
    -p fcu_url:="udp://127.0.0.1:${MAVLINK_PORT}@" \
    -p gcs_url:="" \
    -p target_system_id:=1 \
    -p target_component_id:=1 \
    &>/dev/null &
MAVROS_PID=$!

sleep 5

echo ""
echo -e "${GREEN}==========================================${NC}"
echo -e "${GREEN}  Simulation Running${NC}"
echo -e "${GREEN}==========================================${NC}"
echo ""
echo -e "${BLUE}Components:${NC}"
echo "  - Gazebo: PID $GAZEBO_PID"
echo "  - ArduPilot SITL: PID $SITL_PID"
echo "  - MAVROS: PID $MAVROS_PID"
echo ""
echo -e "${BLUE}MAVROS Topics:${NC}"
echo "  ros2 topic list | grep mavros"
echo ""
echo -e "${BLUE}Control UAV:${NC}"
echo "  # Arm"
echo "  ros2 service call /mavros/cmd/arming mavros_msgs/srv/CommandBool \"{value: true}\""
echo ""
echo "  # Set mode to GUIDED"
echo "  ros2 service call /mavros/set_mode mavros_msgs/srv/SetMode \"{base_mode: 0, custom_mode: 'GUIDED'}\""
echo ""
echo "  # Takeoff to 5 meters"
echo "  ros2 service call /mavros/cmd/takeoff mavros_msgs/srv/CommandTOL \"{altitude: 5}\""
echo ""
echo "  # Send position setpoint (x, y, z in meters)"
echo "  ros2 topic pub /mavros/setpoint_position/local geometry_msgs/PoseStamped \\"
echo "    \"{header: {frame_id: 'map'}, pose: {position: {x: 10, y: 5, z: 5}}}\""
echo ""
echo -e "${YELLOW}Press Ctrl+C to stop simulation${NC}"
echo ""

# Wait for all processes
wait