#!/bin/bash
# =============================================================================
# Drone Simulation - Arch Linux Installation Script
# =============================================================================
# Installs PyBullet and Python dependencies
# ROS 2 requires AUR (optional - not needed for standalone mode)
#
# 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
# Drone Simulation - Environment Activation (Arch Linux)

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

# Try to source ROS 2 if available (from AUR)
if [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    echo "[OK] ROS 2 humble sourced"
elif [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    echo "[OK] ROS 2 jazzy sourced"
else
    echo "[INFO] ROS 2 not found - standalone mode available"
fi

# Activate Python venv
if [ -f "$SCRIPT_DIR/venv/bin/activate" ]; then
    source "$SCRIPT_DIR/venv/bin/activate"
    echo "[OK] Python venv activated"
fi

# Set Gazebo model path if available
export GZ_SIM_RESOURCE_PATH="$SCRIPT_DIR/gazebo/models:$GZ_SIM_RESOURCE_PATH"

echo ""
echo "Environment ready! Run:"
echo "  python standalone_simulation.py"
echo ""
EOF

chmod +x "$PROJECT_ROOT/activate.sh"
echo "[INFO] 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"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Quick start:"
echo "  source activate.sh"
echo "  python standalone_simulation.py"
echo ""
echo "With moving rover:"
echo "  python standalone_simulation.py --pattern circular --speed 0.3"
echo ""
echo "=============================================="
echo "  Optional: Install ROS 2 + Gazebo from AUR"
echo "=============================================="
echo ""
echo "If you need ROS 2 for the full simulation mode:"
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 Fortress (provides 'ign' command)"
echo "  yay -S ignition-fortress"
echo ""
echo "After installing, use:"
echo "  ign gazebo gazebo/worlds/drone_landing.sdf  # Terminal 1"
echo "  python run_gazebo.py --pattern circular     # Terminal 2"
echo ""