#!/bin/bash
set -euo pipefail

# Repeatable pgbench benchmark harness.
# Usage example:
#   PGHOST=localhost PGPORT=5432 PGUSER=postgres PGPASSWORD=secret \
#   ./pgbench_benchmark.sh
#
# Optional environment variables:
#   BENCH_DB=pgbench
#   SCALE=50
#   DURATION=60
#   CLIENTS="1 4 8 16 32"
#   JOBS=4
#   WARMUP_SECONDS=15
#   RUN_READ_ONLY=1
#   RUN_CUSTOM=0
#   CUSTOM_SQL_FILE=pgbench_custom_workload.sql
#   CUSTOM_MODE_NAME=custom_sql
#   CUSTOM_READ_PCT=80
#   CUSTOM_HOT_PCT=90
#   CUSTOM_TXN_SIZE=2
#   CUSTOM_HOT_ACCOUNTS=10000
#   REINIT=0

PGHOST="${PGHOST:-localhost}"
PGPORT="${PGPORT:-5432}"
PGUSER="${PGUSER:-postgres}"
BENCH_DB="${BENCH_DB:-pgbench}"
SCALE="${SCALE:-50}"
DURATION="${DURATION:-60}"
CLIENTS="${CLIENTS:-1 4 8 16 32}"
JOBS="${JOBS:-4}"
WARMUP_SECONDS="${WARMUP_SECONDS:-15}"
RUN_READ_ONLY="${RUN_READ_ONLY:-1}"
RUN_CUSTOM="${RUN_CUSTOM:-0}"
CUSTOM_SQL_FILE="${CUSTOM_SQL_FILE:-pgbench_custom_workload.sql}"
CUSTOM_MODE_NAME="${CUSTOM_MODE_NAME:-custom_sql}"
CUSTOM_READ_PCT="${CUSTOM_READ_PCT:-80}"
CUSTOM_HOT_PCT="${CUSTOM_HOT_PCT:-90}"
CUSTOM_TXN_SIZE="${CUSTOM_TXN_SIZE:-2}"
CUSTOM_HOT_ACCOUNTS="${CUSTOM_HOT_ACCOUNTS:-10000}"
REINIT="${REINIT:-0}"

TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
RESULT_DIR="pgbench_results_${TIMESTAMP}"
SUMMARY_CSV="${RESULT_DIR}/summary.csv"

PSQL=(psql -X -v ON_ERROR_STOP=1 -h "$PGHOST" -p "$PGPORT" -U "$PGUSER")
PGBENCH=(pgbench -h "$PGHOST" -p "$PGPORT" -U "$PGUSER")

require_command() {
    local cmd="$1"
    if ! command -v "$cmd" >/dev/null 2>&1; then
        echo "Missing required command: $cmd" >&2
        exit 1
    fi
}

create_database_if_missing() {
    local exists
    exists="$("${PSQL[@]}" -d postgres -Atc "SELECT 1 FROM pg_database WHERE datname = '${BENCH_DB}';" || true)"
    if [[ "$exists" != "1" ]]; then
        echo "Creating benchmark database: ${BENCH_DB}"
        "${PSQL[@]}" -d postgres -c "CREATE DATABASE ${BENCH_DB};"
    fi
}

needs_init() {
    local rel
    rel="$("${PSQL[@]}" -d "$BENCH_DB" -Atc "SELECT to_regclass('public.pgbench_accounts') IS NOT NULL;")"
    [[ "$rel" != "t" ]]
}

init_or_reinit_data() {
    if [[ "$REINIT" == "1" ]] || needs_init; then
        echo "Initializing pgbench schema/data (scale=${SCALE})"
        "${PGBENCH[@]}" -i -s "$SCALE" "$BENCH_DB"
    else
        echo "Using existing pgbench data in ${BENCH_DB}"
    fi
}

extract_tps() {
    local output_file="$1"
    awk -F'= ' '/tps =/{print $2}' "$output_file" | awk '{print $1}' | tail -n 1
}

extract_avg_latency() {
    local output_file="$1"
    awk -F'= ' '/latency average =/{print $2}' "$output_file" | awk '{print $1}' | tail -n 1
}

select_percentile_ms_from_prefix() {
    local log_prefix="$1"
    local percentile="$2"
    local -a log_files
    local -a values_us
    local count rank value_us

    shopt -s nullglob
    log_files=("${log_prefix}"*)
    shopt -u nullglob

    if (( ${#log_files[@]} == 0 )); then
        echo ""
        return
    fi

    mapfile -t values_us < <(awk 'NF >= 3 && $3 ~ /^[0-9]+(\.[0-9]+)?$/ {print $3}' "${log_files[@]}" | sort -n)
    count="${#values_us[@]}"
    if (( count == 0 )); then
        echo ""
        return
    fi

    rank=$(( (percentile * count + 99) / 100 ))
    if (( rank < 1 )); then rank=1; fi
    if (( rank > count )); then rank=count; fi

    value_us="${values_us[$((rank - 1))]}"
    awk -v us="$value_us" 'BEGIN { printf "%.3f", us/1000.0 }'
}

run_case() {
    local mode="$1"
    local clients="$2"
    shift 2
    local -a extra_args=("$@")
    local log_file="${RESULT_DIR}/${mode}_c${clients}.txt"
    local tx_log_prefix="${RESULT_DIR}/${mode}_c${clients}_txlog"

    echo "Running mode=${mode}, clients=${clients}, duration=${DURATION}s"
    # Warm cache and avoid measuring first-run effects.
    "${PGBENCH[@]}" -n -M prepared -j "$JOBS" -c "$clients" -T "$WARMUP_SECONDS" "${extra_args[@]}" "$BENCH_DB" > /dev/null

    "${PGBENCH[@]}" -n -M prepared -j "$JOBS" -c "$clients" -T "$DURATION" -r -l --log-prefix="$tx_log_prefix" "${extra_args[@]}" "$BENCH_DB" > "$log_file"

    local tps latency p95 p99
    tps="$(extract_tps "$log_file")"
    latency="$(extract_avg_latency "$log_file")"
    p95="$(select_percentile_ms_from_prefix "$tx_log_prefix" 95)"
    p99="$(select_percentile_ms_from_prefix "$tx_log_prefix" 99)"

    echo "${mode},${clients},${DURATION},${JOBS},${tps},${latency},${p95},${p99},${log_file}" >> "$SUMMARY_CSV"
}

main() {
    require_command psql
    require_command pgbench

    mkdir -p "$RESULT_DIR"
    echo "mode,clients,duration_seconds,jobs,tps,latency_ms,p95_ms,p99_ms,details_file" > "$SUMMARY_CSV"

    create_database_if_missing
    init_or_reinit_data

    if [[ "$RUN_CUSTOM" == "1" ]]; then
        if [[ ! -f "$CUSTOM_SQL_FILE" ]]; then
            echo "Custom SQL file not found: $CUSTOM_SQL_FILE" >&2
            exit 1
        fi
    fi

    for c in $CLIENTS; do
        run_case "read_write" "$c"

        if [[ "$RUN_READ_ONLY" == "1" ]]; then
            run_case "read_only" "$c" "-S"
        fi

        if [[ "$RUN_CUSTOM" == "1" ]]; then
            run_case "$CUSTOM_MODE_NAME" "$c" \
                "-D" "read_pct=${CUSTOM_READ_PCT}" \
                "-D" "hot_pct=${CUSTOM_HOT_PCT}" \
                "-D" "txn_size=${CUSTOM_TXN_SIZE}" \
                "-D" "hot_accounts=${CUSTOM_HOT_ACCOUNTS}" \
                "-f" "$CUSTOM_SQL_FILE"
        fi
    done

    echo
    echo "Benchmark completed."
    echo "Summary: $SUMMARY_CSV"
    echo "Raw outputs: $RESULT_DIR"
}

main "$@"