#!/usr/bin/env bash
# ============================================================
# benchmark_now.sh — Run IMMEDIATELY after GPU driver fix
#
# This is the definitive benchmark script for Qwen 3.6 27B
# on RTX 3090 (24 GB). It tests multiple configs and finds
# the maximum tokens/s without OOM.
#
# Prerequisites:
#   1. GPU CUDA must work (torch.cuda.is_available() == True)
#   2. HF_TOKEN set if model needs authentication
#   3. ~50 GB free on /workspace for models
#
# Usage:  bash scripts/benchmark_now.sh
# ============================================================
set -euo pipefail

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

echo "═══════════════════════════════════════════════════════════"
echo " Qwen SpeedLab — Definitive RTX 3090 Benchmark"
echo "═══════════════════════════════════════════════════════════"
echo ""

# --- GPU check ---
if ! python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available'" 2>/dev/null; then
    echo "❌ CUDA not available. Run fix_gpu_host.sh on the HOST first:"
    echo "   sudo apt-get remove --purge nvidia-dkms-580-open nvidia-driver-580-open"
    echo "   sudo apt-get install nvidia-driver-580"
    echo "   sudo reboot"
    exit 1
fi

GPU=$(python3 -c "import torch; print(torch.cuda.get_device_name(0))")
VRAM=$(python3 -c "import torch; print(f'{torch.cuda.get_device_properties(0).total_memory/1024**3:.1f}')")
echo "✅ GPU: $GPU ($VRAM GB)"
echo ""

# --- Phase 1: Model download ---
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo " Phase 1/4: Downloading best model for RTX 3090..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Download the best models (AWQ and FP8)
python3 -c "
from huggingface_hub import snapshot_download
import os

# AWQ 4-bit — best for Ampere (RTX 3090)
print('Downloading QuantTrio/Qwen3.6-27B-AWQ (15 GB)...')
snapshot_download('QuantTrio/Qwen3.6-27B-AWQ',
                   cache_dir='/workspace/models',
                   token=os.environ.get('HF_TOKEN'))
print('✅ AWQ model ready')

# FP8 — official, higher quality
print('Downloading Qwen/Qwen3.6-27B-FP8 (27 GB)...')
snapshot_download('Qwen/Qwen3.6-27B-FP8',
                   cache_dir='/workspace/models',
                   token=os.environ.get('HF_TOKEN'))
print('✅ FP8 model ready')
"

echo ""

# --- Phase 2: Quick vLLM config sweep ---
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo " Phase 2/4: Quick vLLM sweep (6 best configs)..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
python3 scripts/sweep_vllm_configs.py --quick --port 8000
echo ""

# --- Phase 3: Advanced optimization A/B test ---
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo " Phase 3/4: Advanced optimization sweep..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
python3 scripts/sweep_advanced.py --mode baseline --port 8000
echo ""

# --- Phase 4: Final report ---
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo " Phase 4/4: Final report"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""
echo "📊 Results files:"
ls -la reports/results.csv reports/results_advanced.csv reports/summary.md reports/summary_advanced.md 2>/dev/null || echo "  (some reports not generated)"
echo ""

# Print best result if available
if [[ -f reports/summary.md ]]; then
    echo "Best results from sweep:"
    grep -A15 "Best Configuration" reports/summary.md 2>/dev/null | head -20 || true
fi
if [[ -f reports/summary_advanced.md ]]; then
    echo ""
    echo "Optimization impact analysis:"
    grep -A5 "Optimization Impact" reports/summary_advanced.md 2>/dev/null | head -10 || true
fi

echo ""
echo "═══════════════════════════════════════════════════════════"
echo " Benchmark complete!"
echo " Full report: less reports/summary.md"
echo " Raw data:    cat reports/results.csv"
echo "═══════════════════════════════════════════════════════════"