#!/bin/bash
set -e

# Prompt for container name
read -p "Enter PostgreSQL container name: " CONTAINER_NAME

# Check if username is provided
if [[ -z "$CONTAINER_NAME" ]]; then
    echo "Error: Container cannot be empty!" >&2
    exit 1
fi

# Check if container is running
if [ "$(docker ps -q -f name=$CONTAINER_NAME)" = "" ]; then
    echo "Container $CONTAINER_NAME is not running!"
    exit 1
fi

echo "Using container: $CONTAINER_NAME"

# Prompt for PostgreSQL root credentials
read -p "Enter PostgreSQL root username: " PGROOT_USER

# Prompt for username to drop
read -p "Enter PostgreSQL username to drop: " DB_USER

# Check if username is provided
if [[ -z "$DB_USER" ]]; then
    echo "Error: Username cannot be empty!" >&2
    exit 1
fi
# Check if user exists in PostgreSQL
USER_EXISTS=$(docker exec "$CONTAINER_NAME" psql -U "$PGROOT_USER" -tAc "SELECT 1 FROM pg_roles WHERE rolname='$DB_USER';")
if [[ "$USER_EXISTS" != "1" ]]; then
    echo "Error: User '$DB_USER' does not exist in PostgreSQL!" >&2
    exit 1
fi

# # Check if user exists in PostgreSQL
# if ! docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d $CONTAINER_NAME -tAc "SELECT 1 FROM pg_catalog.pg_roles WHERE rolname='$DB_USER'" | grep -q 1; then
#     echo "Error: User '$DB_USER' does not exist in PostgreSQL!" >&2
#     exit 1
# fi

echo "User '$DB_USER' exists."

# Get all databases (not just those with CONNECT privilege)
echo "Checking databases..."
ALL_DATABASES=$(docker exec $CONTAINER_NAME psql -U "$PGROOT_USER" -tAc "
SELECT datname
FROM pg_database
WHERE datname NOT IN ('template0', 'template1')
  AND datallowconn = true;
")

# Check for database privileges
echo "Checking privileges for user '$DB_USER'..."
DATABASES=$(docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d $CONTAINER_NAME -tAc "
SELECT datname
FROM pg_database d
WHERE has_database_privilege('$DB_USER', d.oid, 'CONNECT')
  AND datname NOT IN ('template0', 'template1');
")

if [[ -n "$DATABASES" && "$DATABASES" != "" ]]; then
    echo "WARNING: User '$DB_USER' has privileges on the following database(s):"
    echo "$DATABASES"
    echo ""
    read -p "Do you want to proceed with removing this user? (yes/no): " CONFIRM
    
    if [[ "$CONFIRM" != "yes" ]]; then
        echo "Operation cancelled."
        exit 0
    fi
    
    # Revoke privileges from each database
    echo "Revoking privileges from databases..."
    while IFS= read -r DB_NAME; do
        if [[ -n "$DB_NAME" ]]; then
            echo "  - Revoking privileges on database: $DB_NAME"
            docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d $CONTAINER_NAME -c "REVOKE ALL PRIVILEGES ON DATABASE \"$DB_NAME\" FROM \"$DB_USER\";"
            
            # Revoke privileges on all tables, sequences, and functions in public schema
            docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d "$DB_NAME" -c "REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM \"$DB_USER\";"
            docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d "$DB_NAME" -c "REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM \"$DB_USER\";"
            docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d "$DB_NAME" -c "REVOKE ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA public FROM \"$DB_USER\";"
            docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d "$DB_NAME" -c "REVOKE ALL PRIVILEGES ON SCHEMA public FROM \"$DB_USER\";"
        fi
    done <<< "$DATABASES"
    
    echo "Privileges revoked successfully."
else
    echo "User '$DB_USER' has no database privileges."
fi

# Reassign and drop owned objects in each database
echo "Reassigning and dropping owned objects in all databases..."
while IFS= read -r DB_NAME; do
    if [[ -n "$DB_NAME" ]]; then
        echo "  - Processing database: $DB_NAME"
        docker exec $CONTAINER_NAME psql -U "$PGROOT_USER" -d "$DB_NAME" -c "REASSIGN OWNED BY \"$DB_USER\" TO \"$PGROOT_USER\";" 2>/dev/null
        docker exec $CONTAINER_NAME psql -U "$PGROOT_USER" -d "$DB_NAME" -c "DROP OWNED BY \"$DB_USER\";" 2>/dev/null
    fi
done <<< "$ALL_DATABASES"

# Drop the user using docker exec and psql
echo "Dropping user '$DB_USER'..."
docker exec $CONTAINER_NAME psql -U $PGROOT_USER -d $CONTAINER_NAME -c "DROP USER \"$DB_USER\";"

echo "User $DB_USER has been dropped successfully."