#!/bin/bash

set -euo pipefail

if (( $# != 2 )); then
	echo "usage: $0 prepared-transaction-count port" >&2
	exit 1
fi

PGHOST=${PGHOST:-localhost}
PGPORT=$2
PGDATABASE=${PGDATABASE:-postgres}
PREPARED_XACT_COUNT=$1
GID_PREFIX=${GID_PREFIX:-promotion_test_$(date +%Y%m%d%H%M%S)_$$}

if [[ ! $PREPARED_XACT_COUNT =~ ^[1-9][0-9]*$ ]]; then
	echo "prepared-transaction-count must be a positive integer" >&2
	exit 1
fi

if [[ ! $PGPORT =~ ^[0-9]+$ ]] || (( PGPORT < 1 || PGPORT > 65535 )); then
	echo "port must be an integer between 1 and 65535" >&2
	exit 1
fi

if [[ -n ${PSQL:-} ]]; then
	psql=$PSQL
elif [[ -x ./TMP/postgres/bin/psql ]]; then
	psql=./TMP/postgres/bin/psql
else
	psql=$(command -v psql || true)
fi

if [[ -z $psql ]]; then
	echo "psql was not found; set PSQL to the client executable" >&2
	exit 1
fi

psql_args=(
	-X
	-v ON_ERROR_STOP=1
	-h "$PGHOST"
	-p "$PGPORT"
	-d "$PGDATABASE"
)

read -r max_prepared existing_prepared in_recovery <<EOF
$($psql "${psql_args[@]}" -At -F ' ' -c \
	"SELECT current_setting('max_prepared_transactions'), count(*), pg_is_in_recovery() FROM pg_prepared_xacts")
EOF

if [[ $in_recovery != f ]]; then
	echo "$PGHOST:$PGPORT is in recovery; connect this test to the primary" >&2
	exit 1
fi

available=$((max_prepared - existing_prepared))
if (( available < PREPARED_XACT_COUNT )); then
	echo "insufficient prepared-transaction capacity: requested $PREPARED_XACT_COUNT, available $available" >&2
	echo "max_prepared_transactions=$max_prepared, existing prepared transactions=$existing_prepared" >&2
	exit 1
fi

echo "Creating $PREPARED_XACT_COUNT unresolved prepared transactions on $PGHOST:$PGPORT/$PGDATABASE"
echo "GID prefix: $GID_PREFIX"

$psql "${psql_args[@]}" -q \
	-v prepared_count="$PREPARED_XACT_COUNT" \
	-v gid_prefix="$GID_PREFIX" <<'SQL'
SELECT format('BEGIN; PREPARE TRANSACTION %L;',
              :'gid_prefix' || '_' || generate_series)
FROM generate_series(1, :prepared_count)
\gexec
SQL

created=$($psql "${psql_args[@]}" -At \
	-v gid_prefix="$GID_PREFIX" <<'SQL'
SELECT count(*)
FROM pg_prepared_xacts
WHERE starts_with(gid, :'gid_prefix' || '_');
SQL
)

if (( created != PREPARED_XACT_COUNT )); then
	echo "expected $PREPARED_XACT_COUNT prepared transactions, found $created with this prefix" >&2
	exit 1
fi

flush_lsn=$($psql "${psql_args[@]}" -At -c "SELECT pg_current_wal_flush_lsn()")

echo "Created and left unresolved: $created"
echo "Primary flush LSN after creation: $flush_lsn"
echo "Wait for the standby replay LSN to reach $flush_lsn before promotion."