#!/bin/bash
set -e

# Directories
BASE_DIR="$HOME/pg_rewind_test"
PRIMARY="$BASE_DIR/primary"
STANDBY="$BASE_DIR/standby"

# Random ports
PORT_PRIMARY=$((50000 + RANDOM % 10000))
PORT_STANDBY=$((50000 + RANDOM % 10000))
echo "Selected ports: PRIMARY=$PORT_PRIMARY, STANDBY=$PORT_STANDBY"

echo "Cleaning workspace..."
rm -rf "$BASE_DIR"
mkdir -p "$PRIMARY" "$STANDBY"

# Ensure correct permissions on data directories
chmod 700 "$PRIMARY"
chmod 700 "$STANDBY"

# 1. Init primary
echo "Initializing primary cluster..."
initdb -D "$PRIMARY" --no-locale --encoding=UTF8 --data-checksums > /dev/null

# 2. Start primary
pg_ctl -D "$PRIMARY" -o "-p $PORT_PRIMARY" -l "$PRIMARY/logfile" start
sleep 2

# 3. Create test DB and table
createdb -p $PORT_PRIMARY testdb
psql -p $PORT_PRIMARY -d testdb -c "CREATE TABLE t_test(id int, data text);"
psql -p $PORT_PRIMARY -d testdb -c "INSERT INTO t_test VALUES (1,'primary');"

# Generate extra WAL (simulate shutdown-only checkpoint case)
psql -p $PORT_PRIMARY -d testdb -c "INSERT INTO t_test VALUES (2,'wal1');"
psql -p $PORT_PRIMARY -d testdb -c "INSERT INTO t_test VALUES (3,'wal2');"
psql -p $PORT_PRIMARY -d testdb -c "CHECKPOINT;"

# Capture old primary WAL position BEFORE stopping
OLD_PRIMARY_LAST_WAL=$(psql -p $PORT_PRIMARY -d testdb -t -c "SELECT pg_current_wal_lsn();" | tr -d '[:space:]')

# 4. Basebackup to standby
echo "Creating standby cluster via basebackup..."
pg_basebackup -D "$STANDBY" -R -X stream -c fast -p $PORT_PRIMARY --progress

# 5. Start standby
pg_ctl -D "$STANDBY" -o "-p $PORT_STANDBY" -l "$STANDBY/logfile" start
sleep 2
psql -p $PORT_STANDBY -d testdb -c "SELECT pg_is_in_recovery();"

# 6. Promote standby
pg_ctl -D "$STANDBY" promote
sleep 2
psql -p $PORT_STANDBY -d testdb -c "SELECT pg_is_in_recovery();"

# Capture new primary WAL position after promotion
NEW_PRIMARY_LAST_WAL=$(psql -p $PORT_STANDBY -d testdb -t -c "SELECT pg_current_wal_lsn();" | tr -d '[:space:]')

# 7. Stop old primary
echo "Stopping old primary..."
pg_ctl -D "$PRIMARY" stop -m fast
sleep 2

echo "=== WAL state before rewind ==="
echo "Old primary WAL: $OLD_PRIMARY_LAST_WAL"
echo "New primary WAL: $NEW_PRIMARY_LAST_WAL"

# 8. Decide whether rewind is needed
if [[ "$OLD_PRIMARY_LAST_WAL" < "$NEW_PRIMARY_LAST_WAL" ]]; then
    echo "No rewind required (only shutdown checkpoint difference)."
else
    echo "Running pg_rewind to sync old primary with new primary..."
    pg_rewind -D "$PRIMARY" --source-server="port=$PORT_STANDBY dbname=testdb"
fi

# 9. Start old primary
pg_ctl -D "$PRIMARY" -o "-p $PORT_PRIMARY" -l "$PRIMARY/logfile" start
sleep 2

# 10. Verify data
echo "=== Data check on old primary ==="
psql -p $PORT_PRIMARY -d testdb -c "SELECT * FROM t_test;"

echo "=== Data check on new primary ==="
psql -p $PORT_STANDBY -d testdb -c "SELECT * FROM t_test;"

echo "Test completed successfully."