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

BATCH_SIZE=$1
TMP_DIR=$2
NUM_TABLES=$3
CONNECTIONS="$4"
TIMEOUT="$5"

# ----------------------------
# Usage check
# ----------------------------
if [ "$#" -ne 5 ]; then
    echo "Usage: $0 BATCH_SIZE TMP_DIR NUM_TABLES CONNECTIONS TIMEOUT" >&2
    exit 1
fi

# create tmp directory
mkdir -p "$TMP_DIR"

# create btch tables list
TABLES=()
for i in $(seq 1 "$NUM_TABLES"); do
    TABLES+=("table_batch_$i")
done
echo "Tables: ${TABLES[@]}"

# ----------------------------
# Step 1: Create tables
# ----------------------------
echo "Creating tables..."
for TABLE in "${TABLES[@]}"; do
    psql -c "CREATE UNLOGGED TABLE IF NOT EXISTS $TABLE (
        id BIGINT, id2 BIGINT, id3 BIGINT, id4 BIGINT
    );"
done

# ----------------------------
# Step 2: Create pgbench SQL files
# ----------------------------
echo "Creating pgbench SQL files..."
PG_FILES=()
for TABLE in "${TABLES[@]}"; do
    PG_FILE="$TMP_DIR/${TABLE}.sql"
    cat > "$PG_FILE" <<EOF

INSERT INTO $TABLE (id, id2, id3, id4)
SELECT n, n, n, n
FROM generate_series(1, $BATCH_SIZE) AS n;
EOF
    PG_FILES+=("-f" "$PG_FILE@1")
done

# ----------------------------
# Step 3: Run pgbench serially
# ----------------------------
echo "Starting pgbench inserts (1000 rows per transaction per table)..."
pgbench -c1 -n "${PG_FILES[@]}" -c$CONNECTIONS -T$TIMEOUT

echo "pgbench inserts completed!"