#!/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
    SELECT pg_terminate_backend(pid)
    FROM pg_stat_activity
    WHERE datname = '$TARGET_DB_NAME'
    AND pid != pg_backend_pid();
EOSQL

# Drop and create database as separate commands
dropdb -U "$POSTGRES_USER" --if-exists "$TARGET_DB_NAME"
createdb -U "$POSTGRES_USER" -O "$POSTGRES_USER" "$TARGET_DB_NAME"

# 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;

            -- Handle regular types, excluding arrays and special types
            FOR r IN SELECT t.typname
                    FROM pg_type t
                    LEFT JOIN pg_type arr ON arr.typelem = t.oid
                    WHERE t.typnamespace = 'public'::regnamespace
                    AND t.typtype NOT IN ('a', 'c')  -- Exclude arrays and composite types
                    AND NOT EXISTS (
                        SELECT 1 FROM pg_type arr
                        WHERE arr.typarray = t.oid
                    )
            LOOP
                BEGIN
                    EXECUTE 'ALTER TYPE public.' || quote_ident(r.typname) || ' OWNER TO $POSTGRES_USER';
                EXCEPTION WHEN OTHERS THEN
                    RAISE NOTICE 'Could not alter ownership of type %: %', r.typname, SQLERRM;
                END;
            END LOOP;
        END
        \$\$;
EOSQL

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