#!/bin/bash
set -e

echo "Starting database refresh at $(date)"

# Wait for PostgreSQL to be ready
until pg_isready; do
    echo "Waiting for PostgreSQL to be ready..."
    sleep 2
done

# Create user if it doesn't exist
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
    DO
    \$do\$
    BEGIN
        IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '$POSTGRES_USER') THEN
            CREATE USER $POSTGRES_USER WITH PASSWORD '$POSTGRES_PASSWORD' SUPERUSER;
        END IF;
    END
    \$do\$;
EOSQL

# Drop and recreate the target database
psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" <<-EOSQL
    DROP DATABASE IF EXISTS $TARGET_DB_NAME;
    CREATE DATABASE $TARGET_DB_NAME OWNER $POSTGRES_USER;
EOSQL

# Copy the database if SOURCE_DB_URL is provided
if [ ! -z "$SOURCE_DB_URL" ]; then
    echo "Copying database from $SOURCE_DB_URL to $TARGET_DB_NAME..."

    # Dump the source database and filter out transaction_timeout
    pg_dump "$SOURCE_DB_URL" | grep -v "transaction_timeout" | psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$TARGET_DB_NAME"

    # Update ownership of all objects
    psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$TARGET_DB_NAME" <<-EOSQL
        DO \$\$
        DECLARE
            r record;
        BEGIN
            FOR r IN SELECT tablename FROM pg_tables WHERE schemaname = 'public'
            LOOP
                EXECUTE 'ALTER TABLE public.' || quote_ident(r.tablename) || ' OWNER TO $POSTGRES_USER';
            END LOOP;

            FOR r IN SELECT sequence_name FROM information_schema.sequences WHERE sequence_schema = 'public'
            LOOP
                EXECUTE 'ALTER SEQUENCE public.' || quote_ident(r.sequence_name) || ' OWNER TO $POSTGRES_USER';
            END LOOP;

            FOR r IN SELECT typname FROM pg_type WHERE typnamespace = 'public'::regnamespace
            LOOP
                EXECUTE 'ALTER TYPE public.' || quote_ident(r.typname) || ' OWNER TO $POSTGRES_USER';
            END LOOP;
        END
        \$\$;
EOSQL

    echo "Database refresh completed successfully at $(date)"
else
    echo "No SOURCE_DB_URL provided, skipping database refresh"
fi