# bench_all.sh — Run 3 pgss benchmarks on the current build
#
# Usage:
#   ./bench_all.sh <patch|upstream>
#
# Prerequisites:
#   - PostgreSQL running with pg_stat_statements loaded
#   - max_connections >= 300
#   - Source the activate script before running
#   - pg_stat_statements.max > 5000 and < 5500.
#
# Runs: 5k, 100k, spinlock (5 min each, caller should restart+reset stats in between)
# Output: /tmp/pgss_results_<build>.txt
#

BUILD="${1:?Usage: bench_all.sh <patch|upstream>}"
DURATION=300
CLIENTS=256
JOBS=16
RESULTS="/tmp/pgss_results_${BUILD}.txt"

cat /dev/null > "$RESULTS"

header() {
    echo "" >> "$RESULTS"
    echo "============================================================" >> "$RESULTS"
    echo "  $1" >> "$RESULTS"
    echo "============================================================" >> "$RESULTS"
}

run_test() {
    local test="$1"
    local workload="/tmp/bench_${test}.sql"

    header "TEST: ${test} / BUILD: ${BUILD}"

    # Restart server
    pg_ctl restart -w >> "$RESULTS" 2>&1
    sleep 2

    # Reset stats
    psql -Xc "SELECT pg_stat_statements_reset();" >/dev/null

    # Record config
    echo "" >> "$RESULTS"
    echo "-- Config --" >> "$RESULTS"
    psql -XAtc "SELECT 'pg_stat_statements.max = ' || current_setting('pg_stat_statements.max');" >> "$RESULTS"
    psql -XAtc "SELECT 'max_connections = ' || current_setting('max_connections');" >> "$RESULTS"
    echo "clients = ${CLIENTS}, duration = ${DURATION}s" >> "$RESULTS"

    # Start wait event sampler
    rm -f "/tmp/pgss_waits_${test}.log"
    (for i in $(seq 1 $DURATION); do
        psql -XAtF"|" -c "SELECT wait_event_type, wait_event FROM pg_stat_activity WHERE state = 'active' AND pid != pg_backend_pid() AND wait_event IS NOT NULL;" >> "/tmp/pgss_waits_${test}.log" 2>/dev/null
        sleep 1
    done) &
    local sample_pid=$!

    # Start retention poller (hot vs churn category counts every 1s)
    local retention_pid=0
    if [[ "$test" != "spinlock" ]]; then
        rm -f "/tmp/pgss_retention_${test}.log"
        (for i in $(seq 1 $DURATION); do
            psql -XAtF"|" >> "/tmp/pgss_retention_${test}.log" 2>/dev/null <<'SQL'
SELECT
    CASE WHEN query LIKE 'WITH hot%' THEN 'hot' ELSE 'churn' END AS category,
    count(*),
    sum(calls),
    min(calls),
    max(calls),
    round(avg(calls))
FROM pg_stat_statements
WHERE query LIKE 'WITH hot%' OR query LIKE 'WITH t%'
GROUP BY 1 ORDER BY 1;
SQL
            sleep 1
        done) &
        retention_pid=$!
    fi

    # Run pgbench
    pgbench -f "$workload" -c $CLIENTS -j $JOBS -T $DURATION -P 30 > "/tmp/pgss_pgbench_${test}.log" 2>&1

    # Stop samplers
    kill $sample_pid 2>/dev/null; wait $sample_pid 2>/dev/null || true
    if [[ $retention_pid -ne 0 ]]; then
        kill $retention_pid 2>/dev/null; wait $retention_pid 2>/dev/null || true
    fi

    # --- Collect results ---

    echo "" >> "$RESULTS"
    echo "-- TPS --" >> "$RESULTS"
    grep -E "^(tps|number of transactions|latency average|number of failed)" "/tmp/pgss_pgbench_${test}.log" >> "$RESULTS"

    echo "" >> "$RESULTS"
    echo "-- Progress (TPS every 30s) --" >> "$RESULTS"
    grep "^progress:" "/tmp/pgss_pgbench_${test}.log" >> "$RESULTS"

    echo "" >> "$RESULTS"
    echo "-- Wait Events (top 15) --" >> "$RESULTS"
    printf "%-20s %-30s %s\n" "type" "event" "samples" >> "$RESULTS"
    sort "/tmp/pgss_waits_${test}.log" 2>/dev/null | grep -v "^$" | uniq -c | sort -rn | head -15 | \
    while read -r count line; do
        type="${line%%|*}"
        event="${line##*|}"
        printf "%-20s %-30s %d\n" "$type" "$event" "$count" >> "$RESULTS"
    done

    if [[ "$test" != "spinlock" ]]; then
        echo "" >> "$RESULTS"
        echo "-- Entry Retention --" >> "$RESULTS"
        psql -X >> "$RESULTS" <<'SQL'
SELECT
    CASE WHEN query LIKE 'WITH hot%' THEN 'hot' ELSE 'churn' END AS category,
    count(*) AS entries,
    sum(calls) AS total_calls,
    min(calls) AS min_calls,
    max(calls) AS max_calls,
    round(avg(calls)) AS avg_calls
FROM pg_stat_statements
WHERE query LIKE 'WITH hot%' OR query LIKE 'WITH t%'
GROUP BY 1 ORDER BY 1;
SQL
    else
        echo "" >> "$RESULTS"
        echo "-- Spinlock Entry --" >> "$RESULTS"
        psql -X >> "$RESULTS" <<'SQL'
SELECT calls, total_exec_time::bigint as total_exec_time_ms,
       mean_exec_time, stddev_exec_time
FROM pg_stat_statements
WHERE query = 'SELECT'
LIMIT 1;
SQL
    fi

    if [[ "$test" != "spinlock" ]]; then
        local rlog="/tmp/pgss_retention_${test}.log"
        echo "" >> "$RESULTS"
        echo "-- Retention poll (category|entries|calls|min|max|avg every 1s) --" >> "$RESULTS"
        echo "first 10:" >> "$RESULTS"
        head -10 "$rlog" >> "$RESULTS" 2>/dev/null
        echo "last 10:" >> "$RESULTS"
        tail -10 "$rlog" >> "$RESULTS" 2>/dev/null

        echo "" >> "$RESULTS"
        echo "-- Retention stability (min/max entries across run) --" >> "$RESULTS"
        awk -F'|' '
            /^hot/  { if (hot_min == "" || $2 < hot_min) hot_min = $2;
                      if ($2 > hot_max) hot_max = $2; hot_n++ }
            /^churn/ { if (churn_min == "" || $2 < churn_min) churn_min = $2;
                       if ($2 > churn_max) churn_max = $2; churn_n++ }
            END {
                printf "hot:   samples=%d  min_entries=%s  max_entries=%s\n", hot_n, hot_min, hot_max
                printf "churn: samples=%d  min_entries=%s  max_entries=%s\n", churn_n, churn_min, churn_max
                if (hot_min != "" && hot_max != "" && hot_min + 0 < hot_max * 0.9)
                    printf "WARNING: hot entries dropped >10%% during run (min=%s max=%s)\n", hot_min, hot_max
            }
        ' "$rlog" >> "$RESULTS" 2>/dev/null
    fi

    echo "" >> "$RESULTS"
    echo "-- pg_stat_statements_info --" >> "$RESULTS"
    psql -X >> "$RESULTS" -c "SELECT * FROM pg_stat_statements_info;"
}

# --- Generate workload files ---

cat > /tmp/bench_5k.sql <<'EOF'
\set roll random(1, 100)
\set hot random(1, 1000)
\set churn random(1, 4000)
\if :roll <= 80
WITH hot:hot AS (SELECT 1) SELECT FROM hot:hot
\else
WITH t:churn AS (SELECT 1) SELECT FROM t:churn
\endif
EOF

cat > /tmp/bench_100k.sql <<'EOF'
\set roll random(1, 100)
\set hot random(1, 1000)
\set churn random(1, 100000)
\if :roll <= 80
WITH hot:hot AS (SELECT 1) SELECT FROM hot:hot
\else
WITH t:churn AS (SELECT 1) SELECT FROM t:churn
\endif
EOF

cat > /tmp/bench_spinlock.sql <<'EOF'
SELECT;
EOF

# --- Header ---
{
    echo "pg_stat_statements benchmark — build: ${BUILD}"
    echo "Date: $(date '+%Y-%m-%d %H:%M:%S %Z')"
    echo "Host: $(uname -srm), $(nproc) CPUs, $(free -h | awk '/Mem:/{print $2}') RAM"
    echo "PostgreSQL: $(psql -XAtc 'SELECT version();')"
} >> "$RESULTS"

# --- Run tests ---
run_test 5k
run_test 100k
run_test spinlock

echo "" >> "$RESULTS"
echo "=== DONE ===" >> "$RESULTS"

echo "All tests complete. Results in: ${RESULTS}"
cat "$RESULTS"