#!/bin/bash
# =============================================================================
# Drone Simulation - Arch Linux Installation Script
# =============================================================================
# Installs PyBullet and Python dependencies
# ROS 2 requires AUR (optional - not needed for standalone mode)
# ArduPilot SITL can be built from source (optional)
#
# Usage: ./install_arch.sh
# =============================================================================

set -e

echo "=============================================="
echo "  Drone Simulation - Arch Linux Installation"
echo "=============================================="
echo ""

# Get script directory and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
VENV_DIR="$PROJECT_ROOT/venv"

echo "[INFO] Project root: $PROJECT_ROOT"
echo "[INFO] Virtual environment: $VENV_DIR"

# -----------------------------------------------------------------------------
# Step 1: System Dependencies
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 1/5] Installing system dependencies..."

sudo pacman -Syu --noconfirm
sudo pacman -S --needed --noconfirm \
    python \
    python-pip \
    python-virtualenv \
    base-devel \
    git \
    curl \
    mesa \
    libglvnd

echo "[INFO] System dependencies installed"

# -----------------------------------------------------------------------------
# Step 2: Install AUR Helper (yay) - Optional for ROS 2
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/5] Checking AUR helper..."

if command -v yay &> /dev/null; then
    echo "[INFO] yay already installed"
    HAS_YAY=true
else
    echo "[INFO] Installing yay (AUR helper)..."
    
    cd /tmp
    if [ -d "yay" ]; then
        rm -rf yay
    fi
    git clone https://aur.archlinux.org/yay.git
    cd yay
    makepkg -si --noconfirm
    cd "$PROJECT_ROOT"
    
    HAS_YAY=true
    echo "[INFO] yay installed"
fi

# -----------------------------------------------------------------------------
# Step 3: Create Python Virtual Environment
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 3/5] Creating Python virtual environment..."

if [ -d "$VENV_DIR" ]; then
    rm -rf "$VENV_DIR"
fi

python -m venv "$VENV_DIR"
echo "[INFO] Virtual environment created at: $VENV_DIR"

# -----------------------------------------------------------------------------
# Step 4: Install Python Dependencies
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 4/5] Installing Python dependencies..."

source "$VENV_DIR/bin/activate"
pip install --upgrade pip

if [ -f "$PROJECT_ROOT/requirements.txt" ]; then
    echo "[INFO] Installing from requirements.txt..."
    pip install -r "$PROJECT_ROOT/requirements.txt"
else
    echo "[INFO] Installing packages manually..."
    pip install pybullet numpy pillow pyinstaller
fi

echo "[INFO] Python packages installed"

# -----------------------------------------------------------------------------
# Step 5: Create Activation Script
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 5/5] Creating activation script..."

cat > "$PROJECT_ROOT/activate.sh" << 'EOF'
#!/bin/bash
# RDC Simulation - Environment Activation
# Usage: source activate.sh

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

# =============================================================================
# 1. ROS 2 Setup
# =============================================================================
if [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    echo "[OK] ROS 2 jazzy"
elif [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    echo "[OK] ROS 2 humble"
else
    echo "[WARN] ROS 2 installation not found in /opt/ros/"
fi

# =============================================================================
# 2. ArduPilot Environment
# =============================================================================
export ARDUPILOT_HOME="$HOME/ardupilot"

# Add ArduPilot tools to PATH
# We prepend to ensure we use our versions
export PATH="$ARDUPILOT_HOME/Tools/autotest:$PATH"
export PATH="$ARDUPILOT_HOME/Tools/scripts:$PATH"
export PATH="$HOME/.local/bin:$PATH"

# =============================================================================
# 3. Python Virtual Environment
# =============================================================================
# Deactivate any existing venv to avoid conflicts
if [ -n "$VIRTUAL_ENV" ]; then
    deactivate 2>/dev/null || true
fi

# Activate project venv
if [ -f "$SCRIPT_DIR/venv/bin/activate" ]; then
    source "$SCRIPT_DIR/venv/bin/activate"
    echo "[OK] Python venv active"
else
    echo "[WARN] Python venv not found at $SCRIPT_DIR/venv"
fi

# =============================================================================
# 4. Gazebo & Simulation Paths
# =============================================================================
# Add project models to Gazebo path
export GZ_SIM_RESOURCE_PATH="$SCRIPT_DIR/gazebo/models:$GZ_SIM_RESOURCE_PATH"

# ArduPilot Gazebo Plugin paths
if [ -d "$HOME/ardupilot_gazebo/build" ]; then
    export GZ_SIM_SYSTEM_PLUGIN_PATH="$HOME/ardupilot_gazebo/build:$GZ_SIM_SYSTEM_PLUGIN_PATH"
    export GZ_SIM_RESOURCE_PATH="$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds:$GZ_SIM_RESOURCE_PATH"
else
    echo "[WARN] ArduPilot Gazebo plugin not found at $HOME/ardupilot_gazebo"
fi

echo ""
echo "Environment Configured!"
echo "-----------------------"
echo "Project Root:   $SCRIPT_DIR"
echo "ArduPilot Home: $ARDUPILOT_HOME"
echo "-----------------------"
EOF

chmod +x "$PROJECT_ROOT/activate.sh"
echo "[OK] Created: $PROJECT_ROOT/activate.sh"

# -----------------------------------------------------------------------------
# Verification
# -----------------------------------------------------------------------------
echo ""
echo "Verifying installation..."

source "$PROJECT_ROOT/activate.sh"

echo ""
echo "Checking Python packages:"
python -c "import pybullet; print('  PyBullet: OK')" || echo "  PyBullet: FAILED"
python -c "import numpy; print('  NumPy: OK')" || echo "  NumPy: FAILED"
python -c "from PIL import Image; print('  Pillow: OK')" || echo "  Pillow: FAILED"
python -c "from pymavlink import mavutil; print('  pymavlink: OK')" || echo "  pymavlink: FAILED"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Quick start (3 terminals):"
echo ""
echo "Terminal 1 - Start Gazebo:"
echo "  cd ~/RDC_Simulation"
echo "  ./scripts/run_ardupilot_sim.sh runway"
echo ""
echo "Terminal 2 - Start ArduCopter SITL:"
echo "  source ~/.ardupilot_env"
echo "  sim_vehicle.py -v ArduCopter -f gazebo-iris --model JSON --console"
echo ""
echo "Terminal 3 - Run Controller:"
echo "  cd ~/RDC_Simulation"
echo "  source activate.sh"
echo "  python scripts/run_ardupilot.py --pattern square"
echo ""
echo "=============================================="
echo "  Optional: Install ROS 2 + Gazebo from AUR"
echo "=============================================="
echo ""
echo "If you need ROS 2 for additional features:"
echo ""
echo "  # Install ROS 2 Humble"
echo "  yay -S ros-humble-desktop"
echo ""
echo "  # Install Gazebo bridge"
echo "  yay -S ros-humble-ros-gz"
echo ""
echo "  # Install Gazebo Harmonic"
echo "  yay -S gz-harmonic"
echo ""
echo ""
echo "=============================================="
echo "  Optional: Install ArduPilot SITL"
echo "=============================================="
echo ""
echo "For realistic flight controller simulation:"
echo ""
echo "  # Clone and build ArduPilot"
echo "  git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git ~/ardupilot"
echo "  cd ~/ardupilot"
echo "  pip install -r Tools/environment_install/requirements.txt"
echo ""
echo "  # Clone and build ArduPilot Gazebo plugin"
echo "  git clone https://github.com/ArduPilot/ardupilot_gazebo.git ~/ardupilot_gazebo"
echo "  cd ~/ardupilot_gazebo && mkdir build && cd build"
echo "  cmake .. && make -j\$(nproc)"
echo ""
echo "  # Add to environment"
echo "  export ARDUPILOT_HOME=~/ardupilot"
echo "  export PATH=\$PATH:\$ARDUPILOT_HOME/Tools/autotest"
echo "  export GZ_SIM_SYSTEM_PLUGIN_PATH=~/ardupilot_gazebo/build:\$GZ_SIM_SYSTEM_PLUGIN_PATH"
echo ""
echo "Then run the 3-terminal setup as described above."
echo ""