#!/bin/bash
# RUN_RUST, RUN_SQL, RUN_CONCURRENT default to true; set to false to skip a suite.
# PG_VERSION selects the pgrx feature/server (default pg14); SQL uses libpq PG* settings.

set -euo pipefail

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

echo "=============================================="
echo "pgcel Benchmark Suite"
echo "=============================================="
echo ""

RUN_SQL=${RUN_SQL:-true}
RUN_RUST=${RUN_RUST:-true}
RUN_CONCURRENT=${RUN_CONCURRENT:-true}
PG_VERSION=${PG_VERSION:-pg14}

if [ "$RUN_RUST" = "true" ]; then
    echo ">>> Running Rust criterion benchmarks..."
    echo ""
    cd "$PROJECT_DIR"
    cargo bench --no-default-features --features "$PG_VERSION" --bench cel_benchmarks 2>&1 | tee bench/criterion_results.txt
    echo ""
    echo "Criterion results saved to bench/criterion_results.txt"
    echo "HTML report available at: target/criterion/report/index.html"
    echo ""
fi

if [ "$RUN_SQL" = "true" ]; then
    echo ">>> Running SQL stress tests..."
    echo ""
    
    PSQL=$(which psql 2>/dev/null || echo "")
    if [ -z "$PSQL" ]; then
        echo "WARNING: psql not found in PATH. Skipping SQL tests."
        echo "Install PostgreSQL or add psql to PATH to run SQL benchmarks."
    else
        if ! pg_isready -q 2>/dev/null && [ -z "${PGHOST:-}${PGPORT:-}${PGSERVICE:-}" ]; then
            echo "WARNING: PostgreSQL is not running. Attempting to start..."
            if command -v cargo-pgrx &> /dev/null; then
                cargo pgrx start "$PG_VERSION" || true
            fi
        fi
        
        if pg_isready -q 2>/dev/null; then
            cd "$SCRIPT_DIR"
            psql -v ON_ERROR_STOP=1 -f stress_test.sql 2>&1 | tee sql_results.txt
            echo ""
            echo "SQL results saved to bench/sql_results.txt"
        else
            echo "WARNING: Could not connect to PostgreSQL. Skipping SQL tests."
        fi
    fi
    echo ""
fi

if [ "$RUN_CONCURRENT" = "true" ]; then
    echo ">>> Running concurrent stress test..."
    echo ""
    
    PSQL=$(which psql 2>/dev/null || echo "")
    if [ -z "$PSQL" ] || ! pg_isready -q 2>/dev/null; then
        echo "WARNING: PostgreSQL not available. Skipping concurrent tests."
    else
        cd "$SCRIPT_DIR"
        
        if [ -f concurrent_test.sql ]; then
            echo "Starting 4 concurrent connections..."
            
            pids=()
            for i in 1 2 3 4; do
                psql -v ON_ERROR_STOP=1 -f concurrent_test.sql > "concurrent_results_$i.txt" 2>&1 &
                pids+=("$!")
            done

            failed=0
            for pid in "${pids[@]}"; do
                if ! wait "$pid"; then
                    failed=1
                fi
            done
            if [ "$failed" -ne 0 ]; then
                echo "Concurrent test failed. See concurrent_results_*.txt" >&2
                exit 1
            fi

            echo "Concurrent test complete. Results in concurrent_results_*.txt"
        else
            echo "concurrent_test.sql not found. Skipping."
        fi
    fi
    echo ""
fi

echo "=============================================="
echo "Benchmark suite complete!"
echo "=============================================="