#!/usr/bin/env bash
# ============================================================
# run_all.sh — Master orchestrator for Qwen SpeedLab
#
# Runs the full pipeline:
#   1. Setup environment
#   2. Inspect available models
#   3. Download best model
#   4. Quick sweep (7 configs, ~30 min)
#   5. Advanced A/B sweep (15 configs, ~1h)
#   6. Generate final report
#
# Usage:
#   bash scripts/run_all.sh              # Full pipeline
#   bash scripts/run_all.sh --quick      # Quick only (skip advanced)
#   bash scripts/run_all.sh --dry-run    # Print what would be done
# ============================================================
set -euo pipefail

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

MODE="${1:---full}"
DRY_RUN=false
if [[ "${1:-}" == "--dry-run" ]]; then
    DRY_RUN=true
fi

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
NC='\033[0m'

banner() {
    echo ""
    echo -e "${BLUE}╔══════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║${NC} ${BOLD}$1${NC}"
    echo -e "${BLUE}╚══════════════════════════════════════════════════════════╝${NC}"
    echo ""
}

step() {
    echo -e "${GREEN}[$1/$TOTAL_STEPS]${NC} ${BOLD}$2${NC}"
    echo ""
}

# Count steps
TOTAL_STEPS=6
if [[ "$MODE" == "--quick" ]]; then
    TOTAL_STEPS=4
fi

banner "Qwen SpeedLab — Full Pipeline"
echo "  Mode: $MODE"
echo "  GPU:  $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'unknown')"
echo ""

# ---- Step 1: Setup ----
step 1 "Environment check"
if $DRY_RUN; then
    echo "  [DRY RUN] Would run: bash scripts/setup_runpod.sh"
else
    if [[ -d ".venv" ]]; then
        echo "  ✅ .venv already exists"
        source .venv/bin/activate
    else
        echo "  Running setup_runpod.sh..."
        bash scripts/setup_runpod.sh
        source .venv/bin/activate
    fi
fi

# ---- Step 2: Inspect models ----
step 2 "Model availability scan"
if $DRY_RUN; then
    echo "  [DRY RUN] Would run: python scripts/inspect_model.py --all-variants"
else
    python scripts/inspect_model.py --all-variants 2>/dev/null || {
        echo -e "  ${YELLOW}⚠️  inspect_model.py failed (httpx not installed?) — continuing${NC}"
    }
fi

# ---- Step 3: Download best model ----
step 3 "Download best model"
if $DRY_RUN; then
    echo "  [DRY RUN] Would run: python scripts/download_best_model.py --list"
else
    python scripts/download_best_model.py --list 2>/dev/null || {
        echo -e "  ${YELLOW}⚠️  download_best_model.py failed — continuing${NC}"
    }
    echo ""
    echo -e "  ${YELLOW}⚠️  Skipping actual download (vLLM will pull on first use).${NC}"
    echo "  To download now: python scripts/download_best_model.py"
fi

# ---- Step 4: Quick sweep ----
step 4 "Quick configuration sweep (7 configs, ~30 min)"
if $DRY_RUN; then
    echo "  [DRY RUN] Would run: python scripts/sweep_vllm_configs.py --quick"
else
    echo -e "  ${YELLOW}⚠️  This takes ~30 min and will use the GPU exclusively.${NC}"
    echo "  Press Ctrl+C to skip, or Enter to continue..."
    read -r -t 5 || true
    echo ""
    python scripts/sweep_vllm_configs.py --quick 2>&1 | tail -40 || {
        echo -e "  ${RED}❌ Quick sweep failed — check logs${NC}"
    }
fi

# ---- Step 5: Advanced sweep ----
if [[ "$MODE" != "--quick" ]]; then
    step 5 "Advanced optimization A/B sweep (15 configs, ~1h)"
    if $DRY_RUN; then
        echo "  [DRY RUN] Would run: python scripts/sweep_advanced.py --mode baseline"
    else
        echo -e "  ${YELLOW}⚠️  This takes ~1h. Press Ctrl+C to skip, or Enter to continue...${NC}"
        read -r -t 5 || true
        echo ""
        python scripts/sweep_advanced.py --mode shootout 2>&1 | tail -40 || {
            echo -e "  ${YELLOW}⚠️  Advanced sweep had issues — check reports/summary_advanced.md${NC}"
        }
    fi
fi

# ---- Step 6: Final report ----
step "$TOTAL_STEPS" "Final report"
echo "  Reports generated:"
for f in reports/summary.md reports/summary_advanced.md reports/results.csv reports/results_advanced.csv; do
    if [[ -f "$f" ]]; then
        size=$(wc -l < "$f" 2>/dev/null || echo "?")
        echo -e "    ${GREEN}✅${NC} $f ($size lines)"
    else
        echo -e "    ${YELLOW}—${NC}  $f (not generated — run sweeps to populate)"
    fi
done

# Print best result if available
echo ""
if [[ -f reports/summary.md ]]; then
    echo "────────────────────────────────────────────────────"
    if grep -A5 "Best Configuration" reports/summary.md 2>/dev/null | head -12; then
        true
    fi
    echo "────────────────────────────────────────────────────"
fi

banner "✅ Pipeline complete!"
echo ""
echo "  Next steps:"
echo "    - View report:    cat reports/summary.md"
echo "    - Advanced report: cat reports/summary_advanced.md"
echo "    - Raw data:       cat reports/results.csv"
echo "    - Re-run sweep:   python scripts/sweep_advanced.py --mode grid"
echo "    - Start server:   bash scripts/serve_vllm.sh awq-8k"
echo ""