#!/usr/bin/env bash
# ============================================================
# setup_runpod.sh — Installation environment Qwen SpeedLab
# Pour RTX 3090 / RunPod avec CUDA 12.x
# ============================================================
set -euo pipefail

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

echo "============================================"
echo " Qwen SpeedLab — Setup RunPod RTX 3090"
echo "============================================"

# --- 1. Vérifier GPU --------------------------------------------------
echo ""
echo "[1/8] Vérification GPU..."
if command -v nvidia-smi &>/dev/null; then
    nvidia-smi --query-gpu=name,memory.total,driver_version,cuda.version --format=csv,noheader
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
    VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1 | sed 's/ MiB//')
    echo "  GPU : $GPU_NAME"
    echo "  VRAM: ${VRAM_MB} MiB"

    if [[ "$VRAM_MB" -lt 23000 ]]; then
        echo "  ⚠️  VRAM < 23 Go — les modèles 27B risquent de ne pas passer"
    fi
else
    echo "  ❌ nvidia-smi introuvable — pas de GPU NVIDIA détecté"
    exit 1
fi

# --- 2. Vérifier CUDA / drivers ---------------------------------------
echo ""
echo "[2/8] Vérification CUDA..."
if command -v nvcc &>/dev/null; then
    nvcc --version | grep "release" || true
else
    echo "  nvcc non trouvé (normal sur RunPod, le driver suffit)"
fi

# Vérifier que les libs CUDA sont accessibles
CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
if [[ -d "$CUDA_HOME" ]]; then
    echo "  CUDA_HOME: $CUDA_HOME"
    ls "$CUDA_HOME/lib64/libcudart.so"* 2>/dev/null || echo "  ⚠️  libcudart non trouvée"
fi

# --- 3. Créer venv Python ----------------------------------------------
echo ""
echo "[3/8] Création environnement virtuel Python..."
python3 -m venv "$VENV_DIR" --upgrade-deps
source "${VENV_DIR}/bin/activate"
echo "  Python: $(python3 --version)"
echo "  pip:    $(pip --version)"

# --- 4. Installer PyTorch compatible CUDA ------------------------------
echo ""
echo "[4/8] Installation PyTorch + CUDA..."
# Détecter la version CUDA du driver
CUDA_VER=$(nvidia-smi --query-gpu=cuda.version --format=csv,noheader | head -1 | cut -d. -f1)
echo "  Driver CUDA major: $CUDA_VER"

# PyTorch avec CUDA 12.4 (stable avec vLLM)
pip install --upgrade pip setuptools wheel
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# Vérifier que torch voit CUDA
python3 -c "
import torch
print(f'  PyTorch: {torch.__version__}')
print(f'  CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'  CUDA version: {torch.version.cuda}')
    print(f'  Device count: {torch.cuda.device_count()}')
    print(f'  Device name:  {torch.cuda.get_device_name(0)}')
    print(f'  VRAM total:   {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} Go')
else:
    print('  ⚠️  CUDA non disponible — vérifier driver/compatibilité')
"

# --- 5. Installer vLLM -------------------------------------------------
echo ""
echo "[5/8] Installation vLLM..."
# vLLM avec CUDA 12.4
pip install vllm --no-build-isolation 2>&1 | tail -5 || {
    echo "  ⚠️  Échec vLLM from source, tentative via wheel..."
    pip install vllm 2>&1 | tail -5 || {
        echo "  ❌ Échec installation vLLM"
        echo "  → Essayez: pip install vllm --find-links https://vllm-wheels.s3.us-west-2.amazonaws.com/nightly.html"
    }
}

# Vérifier
python3 -c "import vllm; print(f'  vLLM: {vllm.__version__}')" 2>/dev/null || echo "  ⚠️  vLLM non importable"

# --- 6. Installer SGLang -----------------------------------------------
echo ""
echo "[6/8] Installation SGLang..."
pip install "sglang[all]" 2>&1 | tail -5 || {
    echo "  ⚠️  Échec SGLang — non bloquant"
}

python3 -c "import sglang; print(f'  SGLang: {sglang.__version__}')" 2>/dev/null || echo "  ⚠️  SGLang non importable"

# --- 7. Installer dépendances benchmark --------------------------------
echo ""
echo "[7/8] Installation dépendances benchmark..."
pip install transformers accelerate datasets pandas tqdm psutil openai httpx aiohttp nvidia-ml-py pyyaml rich

# --- 8. Vérification finale --------------------------------------------
echo ""
echo "[8/8] Vérification finale..."
echo ""
echo "  Packages installés :"
pip list 2>/dev/null | grep -iE "torch|vllm|sglang|transformers|openai" || true
echo ""

# Vérifier HF_TOKEN
if [[ -z "${HF_TOKEN:-}" ]]; then
    echo "  ⚠️  HF_TOKEN non défini."
    echo "      → export HF_TOKEN='hf_votre_token'"
    echo "      → ou lancez : huggingface-cli login"
else
    echo "  ✅ HF_TOKEN détecté"
fi

# Vérifier espace disque
DISK_AVAIL=$(df -h /workspace | tail -1 | awk '{print $4}')
echo "  Espace disque disponible: $DISK_AVAIL"

echo ""
echo "============================================"
echo " ✅ Setup terminé !"
echo ""
echo "  Pour activer l'environnement :"
echo "    source ${VENV_DIR}/bin/activate"
echo ""
echo "  Pour lancer le benchmark :"
echo "    bash scripts/serve_vllm.sh"
echo "    python scripts/bench_openai_api.py"
echo "    python scripts/sweep_vllm_configs.py"
echo "============================================"