#!/usr/bin/env bash
set -euo pipefail

# ----------------------------
# Usage check
# ----------------------------
if [ "$#" -ne 6 ]; then
    echo "Usage: $0 PG_FILE NUM_TABLES NUM_ROWS" >&2
    exit 1
fi

PG_FILE="$1"
NUM_TABLES="$2"
NUM_ROWS="$3"
READY_FILE="$4"
CONNECTIONS="$5"
TIMEOUT="$6"

# ----------------------------
# Create tables
# ----------------------------
psql <<EOF
DO
\$\$
DECLARE
    i INT;
    tbl_name TEXT;
BEGIN
    FOR i IN 1..$NUM_TABLES LOOP
        tbl_name := format('table_%s', i);
        EXECUTE format('
            CREATE UNLOGGED TABLE %I (
                id BIGSERIAL PRIMARY KEY,
                col1 TEXT,
                col2 TEXT,
                col3 TEXT,
                col4 BIGINT,
                col5 BIGINT,
                col6 BIGINT,
                col7 BIGINT,
                col8 BIGINT,
                col9 BIGINT,
                col10 BIGINT,
                col11 TIMESTAMP DEFAULT NOW(),
                col12 TIMESTAMP,
                col13 TIMESTAMP,
                col14 TEXT,
                col15 BIGINT
            );
        ', tbl_name);
    END LOOP;
END
\$\$;
EOF

# ----------------------------
# Populate tables 
# ----------------------------
for i in $(seq 1 "$NUM_TABLES"); do
    echo "Inserting $NUM_ROWS rows into table_$i..."
    psql -c "
    INSERT INTO table_${i} (col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12, col13, col14, col15)
    SELECT
        md5(random()::text),
        md5(random()::text),
        md5(random()::text),
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        (random()*1000000)::BIGINT,
        NOW() - (random() * interval '30 days'),
        NOW() - (random() * interval '30 days'),
        NOW() - (random() * interval '30 days'),
        md5(random()::text),
        (random()*1000000)::BIGINT
    FROM generate_series(1, $NUM_ROWS);
    "
done

# ----------------------------
# VACUUM tables
# ----------------------------
for i in $(seq 1 "$NUM_TABLES"); do
    echo "Vacuuming table_$i..."
    psql -c "VACUUM FREEZE table_$i;"
done

# ----------------------------
# Stop and restart PostgreSQL 
# ----------------------------
echo "Stopping PostgreSQL..."
pg_ctl stop

echo "Waiting for PostgreSQL to fully stop..."
while pgrep -x postgres >/dev/null; do
    echo "  Still shutting down..."
    sleep 1
done
echo "PostgreSQL is fully stopped."

echo "Starting PostgreSQL..."
pg_ctl start

echo "Waiting for PostgreSQL to start up..."
until pg_isready -q; do
    echo "  Still starting..."
    sleep 1
done
echo "PostgreSQL is up and running!"

# ----------------------------
# Signal OLTP readiness
# ----------------------------
touch $READY_FILE

# ----------------------------
# Create pgbench SQL file
# ----------------------------
cat >"$PG_FILE" <<EOF
\set upd1 random(1, $NUM_TABLES)
\set id random(1, $NUM_ROWS)

UPDATE table_:upd1 SET col1 = col1 WHERE id = :id;
EOF

# ----------------------------
# Run pgbench
# ----------------------------
pgbench -f "$PG_FILE" -c$CONNECTIONS -T$TIMEOUT