#!/usr/bin/env bash
#
# run.sh [size ...]
#
# Build the CURRENTLY CHECKED-OUT BRANCH of this repo (its HEAD) as a release,
# then run the buffer-mapping-table probe and print the per-op numbers for it.
# Single arm (just this branch's buf_table.c) — for the flat-vs-dynahash A/B
# use scripts/compare_probe.sh.
#
# The build is cached per commit under buftable_bench/.builds/<commit>, so the
# first run for a commit takes a few minutes and later runs are instant.
# It exports the committed tree (git archive) into a temp dir to build, so your
# working checkout is never touched.
#
# Env: PG_CONFIG (use this prebuilt install instead of building),
#      BUFTABLE_PROBE_ROUNDS (default 10), BUFTABLE_BENCH_WORK.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORK="${BUFTABLE_BENCH_WORK:-$HERE/_work}"
mkdir -p "$WORK"
ROUNDS="${BUFTABLE_PROBE_ROUNDS:-10}"
SIZES="${*:-1GB}"
BUILD_LOG="$WORK/build.log"

# Progress on stderr so result tables on stdout stay parseable.
log() { printf '%s\n' "$*" >&2; }

repo="$(git -C "$HERE" rev-parse --show-toplevel)"
# stg is handy to edit commits (https://stacked-git.github.io/)
commit=""
if commit="$(cd "$repo" && stg series 2>/dev/null | awk '$1 == ">" { print $2; found=1 } END { exit !found }')"; then
	:
else
	commit="$(git -C "$repo" rev-parse --short HEAD)"
fi
branch="$(git -C "$repo" rev-parse --abbrev-ref HEAD)"
label="$branch@$commit"

# Linux installs buftable_bench.so; macOS installs buftable_bench.dylib.
buftable_bench_installed() {
	local pkglibdir="$("$1/bin/pg_config" --pkglibdir)"
	local ext
	for ext in so dylib dll; do
		if [[ -e "$pkglibdir/buftable_bench.$ext" ]]; then
			return 0
		fi
	done
	return 1
}

# ---- locate or build a release install of the current HEAD ------------------
if [[ -n "${PG_CONFIG:-}" ]]; then
	BIN="$("$PG_CONFIG" --bindir)"
	log "==> using PG_CONFIG build: $("$PG_CONFIG" --version) [$BIN]"
else
	prefix="$WORK/.builds/buftable"
	src="$WORK/src"
	log "==> building $label as release (log: $BUILD_LOG)..."
	if [[ -d "$src" ]]; then
		git -C "$src" fetch "$repo"
		git -C "$src" reset --hard "$(git -C "$repo" rev-parse HEAD)"
	else
		git -C "$repo" worktree add -f "$src" HEAD
		(
			cd "$src"
			./configure --prefix="$prefix" \
				--without-icu \
				--without-zlib \
				--without-readline
		) >"$BUILD_LOG" 2>&1 || {
			log "configure failed; see $BUILD_LOG"
			tail -n 40 "$BUILD_LOG" >&2
			exit 1
		}
	fi
	(
		cd "$src"
		# Install the server too.  `make -C …/buftable_bench install` only
		# copies the extension, leaving prefix/bin empty — initdb then fails
		# with "command not found" after a silent compile.
		make -j"$(sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)"
		make install
		make -C src/test/modules/buftable_bench install
	) >>"$BUILD_LOG" 2>&1 || {
		log "build failed; see $BUILD_LOG"
		tail -n 80 "$BUILD_LOG" >&2
		exit 1
	}
	log "==> built $label"
	BIN="$prefix/bin"
fi

if [[ ! -x "$BIN/initdb" || ! -x "$BIN/postgres" ]]; then
	log "error: no postgres install at $BIN (initdb/postgres missing)."
	log "The bench tree must be 'make install'd into the prefix, not only compiled."
	exit 1
fi
if ! buftable_bench_installed "$(dirname "$BIN")"; then
	log "error: buftable_bench extension is not installed under $BIN/.."
	exit 1
fi

# ---- run the probe per size -------------------------------------------------
size_to_bytes() {
	local s
	s="$(printf '%s' "$1" | tr '[:lower:]' '[:upper:]')"
	case "$s" in
		*GB) echo $(( ${s%GB} * 1024 * 1024 * 1024 ));;
		*MB) echo $(( ${s%MB} * 1024 * 1024 ));;
		*KB) echo $(( ${s%KB} * 1024 ));;
		*)   echo "$s";;
	esac
}

echo
echo "===== buftable_bench: $label (random access, rounds=$ROUNDS) ====="
for SIZE in $SIZES; do
	DATADIR="$(mktemp -d "$WORK/pgdata.${SIZE}.XXXX")"
	SOCKDIR="$(mktemp -d /tmp/pgb.XXXXXX)"
	N="$(awk -v b="$(size_to_bytes "$SIZE")" 'BEGIN{printf "%d", 0.8*b/8192}')"
	LOG="$DATADIR/log"

	log "==> initdb + start shared_buffers=$SIZE (n=$N keys, rounds=$ROUNDS)"
	"$BIN/initdb" -D "$DATADIR" --no-sync -A trust >/dev/null 2>&1
	cat >> "$DATADIR/postgresql.conf" <<CONF
shared_buffers = '$SIZE'
jit = off
autovacuum = off
fsync = off
bgwriter_lru_maxpages = 0
listen_addresses = ''
unix_socket_directories = '$SOCKDIR'
CONF
	if ! "$BIN/pg_ctl" -D "$DATADIR" -l "$LOG" -w start >/dev/null; then
		log "pg_ctl start failed; server log:"
		cat "$LOG" >&2 || true
		rm -rf "$DATADIR" "$SOCKDIR"
		exit 1
	fi

	num="numeric(15,2)"
	echo "-- shared_buffers=$SIZE  (n=$N keys) --"
	log "==> running probe (this is silent until psql returns)"
	if ! "$BIN/psql" -h "$SOCKDIR" -d postgres -q -X -P pager=off -v ON_ERROR_STOP=1 \
		-c "CREATE EXTENSION buftable_bench;" \
		-c "SELECT op,
		           avg(avg_ns)::$num AS avg,
		           min(avg_ns)::$num AS min,
				   percentile_disc(0.25) WITHIN GROUP (ORDER BY avg_ns)::$num AS \"[q1\",
				   percentile_disc(0.50) WITHIN GROUP (ORDER BY avg_ns)::$num AS median,
				   percentile_disc(0.75) WITHIN GROUP (ORDER BY avg_ns)::$num AS \"q3]\",
				   percentile_disc(0.99) WITHIN GROUP (ORDER BY avg_ns)::$num AS p99,
				   stddev(avg_ns)::$num AS std
			FROM buftable_bench_probe($N, $ROUNDS, false)
		    GROUP BY op, id
			ORDER BY id;"
	then
		log "psql probe failed; server log:"
		cat "$LOG" >&2 || true
		"$BIN/pg_ctl" -D "$DATADIR" -m immediate stop >/dev/null 2>&1 || true
		rm -rf "$DATADIR" "$SOCKDIR"
		exit 1
	fi

	"$BIN/pg_ctl" -D "$DATADIR" -m immediate stop >/dev/null 2>&1 || true
	rm -rf "$DATADIR" "$SOCKDIR"
done
echo "================================================================"