#!/usr/bin/env bash
set -e

echo "=== Elden Ring AI — Dependency Installer ==="

# ── 1. System packages ────────────────────────────────────────────────────────
echo "[1/4] Installing system packages..."
sudo pacman -S --needed --noconfirm \
    base-devel \
    linux-headers \
    v4l2loopback-dkms \
    gpu-screen-recorder \
    python \
    python-pip

# ── 2. Persistent v4l2loopback config ─────────────────────────────────────────
echo "[2/4] Configuring v4l2loopback..."
echo 'v4l2loopback' | sudo tee /etc/modules-load.d/v4l2loopback.conf > /dev/null
echo 'options v4l2loopback devices=1 video_nr=2 card_label="EldenVision" exclusive_caps=1' \
    | sudo tee /etc/modprobe.d/v4l2loopback.conf > /dev/null
sudo modprobe v4l2loopback devices=1 video_nr=2 card_label="EldenVision" exclusive_caps=1 2>/dev/null || true

# Add user to input group (needed for evdev/uinput)
sudo usermod -aG input "$USER"
echo "    NOTE: Log out and back in for the input group to take effect."

# ── 3. Python venv ────────────────────────────────────────────────────────────
echo "[3/4] Setting up Python virtual environment..."
if [ ! -f "bin/activate" ]; then
    python -m venv .
fi
# shellcheck disable=SC1091
source bin/activate

# ── 4. Python packages ────────────────────────────────────────────────────────
echo "[4/4] Installing Python packages..."

# Detect CUDA version from driver and map to a PyTorch wheel tag
CUDA_RAW=$(nvidia-smi 2>/dev/null | grep -oP "CUDA Version: \K[\d.]+" || echo "0")
CUDA_MAJOR=$(echo "$CUDA_RAW" | cut -d. -f1)
CUDA_MINOR=$(echo "$CUDA_RAW" | cut -d. -f2)
CUDA_NUM=$((CUDA_MAJOR * 10 + CUDA_MINOR))  # e.g. 13.2 → 132

if   [ "$CUDA_NUM" -ge 130 ]; then TORCH_CU="cu130"
elif [ "$CUDA_NUM" -ge 124 ]; then TORCH_CU="cu124"
elif [ "$CUDA_NUM" -ge 121 ]; then TORCH_CU="cu121"
elif [ "$CUDA_NUM" -ge 118 ]; then TORCH_CU="cu118"
else
    echo "    WARNING: No supported CUDA found, installing CPU-only PyTorch."
    TORCH_CU="cpu"
fi

TORCH_INDEX="https://download.pytorch.org/whl/${TORCH_CU}"
echo "    CUDA ${CUDA_RAW} detected → using index: ${TORCH_INDEX}"

pip install --upgrade pip
pip install torch torchvision --index-url "$TORCH_INDEX"
pip install -r requirements.txt

echo ""
echo "=== Done! Activate the venv with: source bin/activate ==="
echo "=== Then start training with:     python train.py      ==="