run.sh

#!/usr/bin/env bash
set -euo pipefail

cd "$(dirname "$0")"

export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"
unset PYENV_VERSION

BENCH_DIR="benchmarks"
BENCHMARKS=("$BENCH_DIR"/*_benchmark.py)
SUMMARY_MODE=false
SUMMARY_FILE="results/summary.md"
declare -a ALL_ERRORS=()
declare -a _bm_rows=()
declare -A _bm_seen=()
declare -a _bm_order=()
declare -a _bm_errors=()

# ==========================================================
# Parse CLI arguments
# ==========================================================
parse_args() {
    for arg in "$@"; do
        case $arg in
            --summary-table)
                SUMMARY_MODE=true
                shift
                ;;
            *)
                echo "Unknown option: $arg"
                echo "Usage: ./run.sh [--summary-table]"
                exit 1
                ;;
        esac
    done
}

# ==========================================================
# Detect installed Python versions via pyenv
# ==========================================================
detect_versions() {
    local raw_versions
    raw_versions=$(pyenv versions --bare | grep -vE 'system|dev|envs' | tr -d '*')

    if [[ -z "$raw_versions" ]]; then
        echo "❌ No valid Python versions found in pyenv!"
        exit 1
    fi

    VERSIONS=($raw_versions)
    echo "📦 Detected Python versions:"
    for version in "${VERSIONS[@]}"; do
        echo " - $version"
    done
    echo
}

# ==========================================================
# Ensure dependencies for each Python version via uv
# ==========================================================
ensure_dependencies() {
    local version="$1"
    echo "🔧 Ensuring dependencies for Python $version"
    local sync_out sync_err
    sync_out=$(PYENV_VERSION="$version" uv sync --all-extras --dev --locked 2>&1) || sync_err=$?
    if [[ -n "${sync_err:-}" ]]; then
        echo "⚠️  uv sync --locked failed for Python $version — retrying without lock"
        sync_out=$(PYENV_VERSION="$version" uv sync --all-extras --dev 2>&1) || sync_err=$?
        if [[ -n "${sync_err:-}" ]]; then
            echo "❌ uv sync failed for Python $version"
            echo "$sync_out"
            return 1
        fi
    fi
}

# ==========================================================
# Execute a single benchmark for one Python version
# ==========================================================
run_benchmarks_for_version() {
    local version="$1"
    local benchmark="$2"

    echo
    echo "🐍 Python $version"
    echo "---------------------------------------"

    if ! pyenv prefix "$version" &>/dev/null; then
        echo "⚠️  Skipping invalid or unavailable version: $version"
        _bm_errors+=("[$version] Skipped: pyenv prefix not found")
        return
    fi

    ensure_dependencies "$version"
    if [[ $? -ne 0 ]]; then
        _bm_errors+=("[$version] Dependency sync failed")
        return
    fi

    echo "▶️  Using Python $version"

    local output exit_code
    if [[ -t 1 ]]; then
        output=$(PYENV_VERSION="$version" uv run --python "$version" python "$benchmark" 2>&1 | tee /dev/tty)
        exit_code=${PIPESTATUS[0]}
    else
        output=$(PYENV_VERSION="$version" uv run --python "$version" python "$benchmark" 2>&1)
        exit_code=$?
        echo "$output"
    fi

    if [[ $exit_code -ne 0 ]]; then
        echo "❌ Benchmark failed for Python $version (exit $exit_code)"
        _bm_errors+=("[$version] Exit code: $exit_code")
        _bm_rows+=("| $version | **ERROR** |")
        return
    fi

    local py_version
    py_version=$(echo "$output" | grep -E '^Python:' | awk '{print $2}')
    local metrics
    metrics=$(echo "$output" | grep -E '^[^:]+: [0-9]+\.[0-9]+' | grep -vE '^Benchmark:|^Python:' | grep -vE 'interpreter at|Creating virtual|Removed virtual|Installed |packages in')

    local row="| $py_version"
    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local key val
        key=$(echo "$line" | cut -d':' -f1 | tr -d '\r')
        val=$(echo "$line" | cut -d':' -f2 | xargs)
        row="$row | $val"
        if [[ -z "${_bm_seen[$key]:-}" ]]; then
            _bm_seen["$key"]=1
            _bm_order+=("$key")
        fi
    done <<< "$metrics"
    row="$row |"
    _bm_rows+=("$row")

    echo "---------------------------------------"
}

# ==========================================================
# Write Markdown summary for a single benchmark
# ==========================================================
write_summary() {
    local bench_name="$1"
    local order_str="$2"
    local errors_str="$3"

    local -a _bm_order_split
    readarray -t _bm_order_split <<< "$order_str"
    local -a _bm_errors_split
    readarray -t _bm_errors_split <<< "$errors_str"

    echo "## 🧩 $bench_name" >> "$SUMMARY_FILE"
    echo "" >> "$SUMMARY_FILE"

    local header="| Python Version"
    for key in "${_bm_order_split[@]}"; do
        header="$header | ${key} (s)"
    done
    echo "$header |" >> "$SUMMARY_FILE"

    local sep_line="|----------------"
    for _ in "${_bm_order_split[@]}"; do
        sep_line="$sep_line|----------------"
    done
    echo "$sep_line|" >> "$SUMMARY_FILE"

    for row in "${_bm_rows[@]}"; do
        echo "$row" >> "$SUMMARY_FILE"
    done
    echo "" >> "$SUMMARY_FILE"

    if [[ ${#_bm_errors_split[@]} -gt 0 && -n "${_bm_errors_split[0]}" ]]; then
        echo "### Errors" >> "$SUMMARY_FILE"
        echo "" >> "$SUMMARY_FILE"
        for err in "${_bm_errors_split[@]}"; do
            echo "- $err" >> "$SUMMARY_FILE"
        done
        echo "" >> "$SUMMARY_FILE"
    fi
}

# ==========================================================
# Main benchmark runner
# ==========================================================
run_all_benchmarks() {
    for benchmark in "${BENCHMARKS[@]}"; do
        local bench_name
        bench_name=$(basename "$benchmark")
        echo
        echo "🚀 Benchmark file: $bench_name"
        echo "======================================="

        _bm_rows=()
        _bm_seen=()
        _bm_order=()
        _bm_errors=()

        for version in "${VERSIONS[@]}"; do
            run_benchmarks_for_version "$version" "$benchmark"
        done

        if [[ "$SUMMARY_MODE" == true ]]; then
            local order_str
            order_str="$(printf "%s\n" "${_bm_order[@]}")"
            local errors_str
            errors_str="$(printf "%s\n" "${_bm_errors[@]}")"
            write_summary "$bench_name" "$order_str" "$errors_str"
        fi

        ALL_ERRORS+=("${_bm_errors[@]}")

        echo "======================================="
        echo "✅ Finished benchmark: $bench_name"
        echo "======================================="
    done
}

# ==========================================================
# Initialization
# ==========================================================
initialize_environment() {
    clear
    echo "======================================="
    echo "========= Running Benchmarks =========="
    echo "======================================="

    if [[ ! -d "$BENCH_DIR" ]]; then
        echo "❌ Benchmarks directory not found: $(pwd)/$BENCH_DIR"
        exit 1
    fi

    if [[ ${#BENCHMARKS[@]} -eq 0 ]]; then
        echo "❌ No benchmark scripts found in $(pwd)/$BENCH_DIR"
        exit 1
    fi

    if [[ "$SUMMARY_MODE" == true ]]; then
        mkdir -p results
        echo "# 🧪 Benchmark Summary" > "$SUMMARY_FILE"
        echo "" >> "$SUMMARY_FILE"
    fi
}

# ==========================================================
# Cleanup
# ==========================================================
cleanup() {
    pyenv shell --unset || true
    unset PYENV_VERSION
    echo
    echo "======================================="
    echo "======= All Benchmarks Complete ======="
    echo "======================================="

    if [[ ${#ALL_ERRORS[@]} -gt 0 ]]; then
        echo
        echo "❌ Errors encountered:"
        for err in "${ALL_ERRORS[@]}"; do
            echo "  - $err"
        done
        echo
    fi

    if [[ "$SUMMARY_MODE" == true ]]; then
        echo
        echo "📊 Summary tables saved to: $SUMMARY_FILE"
        echo
    fi
}

# ==========================================================
# Main entry
# ==========================================================
main() {
    parse_args "$@"
    detect_versions
    initialize_environment
    run_all_benchmarks
    cleanup
}

main "$@"