#!/usr/bin/env bash
# perf-entrypoint.sh — Docker entrypoint for the pg_stat_ch perf benchmark container.
#
# Bootstraps a throwaway PostgreSQL instance with pg_stat_ch loaded and
# use_otel=on, then hands off to bench-otel.sh. Any extra arguments are
# forwarded to bench-otel.sh (e.g. --tps 20000 --duration 60).

set -euo pipefail

OTEL_ENDPOINT="${OTEL_ENDPOINT:-otelcol:4317}"
OTEL_HEALTH="${OTEL_HEALTH:-http://otelcol:13133}"
PGDATA="${PGDATA:-/var/lib/postgresql/data}"
RESULTS_DIR="${RESULTS_DIR:-/results}"

# Ensure results directory exists (may be a volume mount)
mkdir -p "$RESULTS_DIR"

echo "=== pg_stat_ch OTel Benchmark (Docker) ==="
echo "OTel endpoint : $OTEL_ENDPOINT"
echo "OTel health   : $OTEL_HEALTH"
echo "Results dir   : $RESULTS_DIR"
echo

# --- Bootstrap PostgreSQL ---
install -o postgres -g postgres -m 700 -d "$PGDATA"

echo -n "Initializing database cluster..."
gosu postgres initdb \
    -D "$PGDATA" \
    --auth-host=trust \
    --auth-local=trust \
    -U postgres \
    --encoding=UTF8 \
    --locale=C \
    --no-instructions \
    2>/dev/null
echo " done."

# Configure pg_stat_ch with OTel export
cat >> "$PGDATA/postgresql.conf" <<EOF

# pg_stat_ch configuration (injected by perf-entrypoint.sh)
shared_preload_libraries = 'pg_stat_ch'
pg_stat_ch.enabled = on
pg_stat_ch.use_otel = on
pg_stat_ch.otel_endpoint = '${OTEL_ENDPOINT}'
pg_stat_ch.queue_capacity = 131072
EOF

# --- Start PostgreSQL in the background ---
gosu postgres postgres -D "$PGDATA" -k /var/run/postgresql &
PG_PID=$!

echo -n "Waiting for PostgreSQL..."
PG_START_DEADLINE=$(($(date +%s) + 30))
until gosu postgres pg_isready -U postgres -h /var/run/postgresql -q 2>/dev/null; do
    printf '.'
    sleep 1
    if ! kill -0 "$PG_PID" 2>/dev/null; then
        echo
        echo "ERROR: PostgreSQL process exited unexpectedly (magic block mismatch?)"
        exit 1
    fi
    if [[ $(date +%s) -ge $PG_START_DEADLINE ]]; then
        echo
        echo "ERROR: PostgreSQL failed to start within 30s"
        exit 1
    fi
done
echo " ready. (PID $PG_PID)"

# Create extension and initialize pgbench schema
gosu postgres psql -U postgres -h /var/run/postgresql \
    -c "CREATE EXTENSION pg_stat_ch" postgres >/dev/null

echo -n "Initializing pgbench schema..."
gosu postgres pgbench -U postgres -h /var/run/postgresql -i -q postgres 2>/dev/null
echo " done."

# --- Wait for OTel collector ---
echo -n "Waiting for OTel collector at $OTEL_HEALTH..."
OTEL_START_DEADLINE=$(($(date +%s) + 60))
until curl -sf "$OTEL_HEALTH" >/dev/null 2>&1; do
    printf '.'
    sleep 1
    if [[ $(date +%s) -ge $OTEL_START_DEADLINE ]]; then
        echo
        echo "ERROR: OTel collector failed to become healthy within 60s"
        exit 1
    fi
done
echo " ready."

# --- Run the benchmark ---
# Use Unix socket so bench-otel.sh psql calls don't need a password.
export PGHOST=/var/run/postgresql
export PGUSER=postgres
export PGDATABASE=postgres
export OTEL_HEALTH
export OTEL_ENDPOINT

cd "$RESULTS_DIR"
BENCH_EXIT=0
bench-otel.sh "$@" || BENCH_EXIT=$?

# --- Shutdown ---
gosu postgres pg_ctl -D "$PGDATA" stop -m fast 2>/dev/null || kill -TERM "$PG_PID"
wait "$PG_PID" 2>/dev/null || true

exit "$BENCH_EXIT"