#!/bin/sh
set -eu

umask 077

confirmation_phrase="replace-bowli-production-database"

if [ "${RESTORE_CONFIRM:-}" != "$confirmation_phrase" ]; then
    echo "Database restoration was not confirmed." >&2
    echo "Run it with RESTORE_CONFIRM=$confirmation_phrase" >&2
    exit 1
fi

require_variable() {
    variable_name="$1"
    eval "variable_value=\${$variable_name:-}"

    if [ -z "$variable_value" ]; then
        echo "$variable_name is required." >&2
        exit 1
    fi
}

aws_cli() {
    endpoint="${BACKUP_S3_ENDPOINT:-${AWS_ENDPOINT:-}}"

    if [ -n "$endpoint" ]; then
        aws --endpoint-url "$endpoint" "$@"
    else
        aws "$@"
    fi
}

for variable in DB_DATABASE DB_HOST DB_PASSWORD DB_USERNAME
do
    require_variable "$variable"
done

case "$DB_DATABASE" in
    *[!A-Za-z0-9_-]*)
        echo "DB_DATABASE may only contain letters, numbers, underscores and hyphens for restoration." >&2
        exit 1
        ;;
esac

source_name="${1:-latest-s3}"
backup_directory="${BACKUP_DIRECTORY:-/backups}"
backup_prefix="${BACKUP_S3_PREFIX:-database-backups/postgres}"
backup_prefix="${backup_prefix#/}"
backup_prefix="${backup_prefix%/}"
temporary_directory="$(mktemp -d)"
restore_path=""

cleanup() {
    rm -rf "$temporary_directory"
}

trap cleanup EXIT INT TERM

download_s3_backup() {
    s3_uri="$1"
    file_name="${s3_uri##*/}"
    restore_path="$temporary_directory/$file_name"

    aws_cli s3 cp "$s3_uri" "$restore_path" --only-show-errors
    aws_cli s3 cp "$s3_uri.sha256" "$restore_path.sha256" --only-show-errors
}

case "$source_name" in
    latest-local)
        restore_path="$(find "$backup_directory" -maxdepth 1 -type f -name "*.dump" | sort | tail -n 1)"
        ;;
    latest-s3)
        require_variable AWS_ACCESS_KEY_ID
        require_variable AWS_DEFAULT_REGION
        require_variable AWS_SECRET_ACCESS_KEY
        require_variable BACKUP_S3_BUCKET

        s3_directory="s3://$BACKUP_S3_BUCKET"
        if [ -n "$backup_prefix" ]; then
            s3_directory="$s3_directory/$backup_prefix"
        fi

        latest_key="$(aws_cli s3 ls "$s3_directory/" --recursive \
            | awk '$4 ~ /\.dump$/ { print $4 }' \
            | sort \
            | tail -n 1)"

        if [ -n "$latest_key" ]; then
            download_s3_backup "s3://$BACKUP_S3_BUCKET/$latest_key"
        fi
        ;;
    s3://*)
        require_variable AWS_ACCESS_KEY_ID
        require_variable AWS_DEFAULT_REGION
        require_variable AWS_SECRET_ACCESS_KEY
        download_s3_backup "$source_name"
        ;;
    /*)
        restore_path="$source_name"
        ;;
    *)
        restore_path="$backup_directory/$source_name"
        ;;
esac

if [ -z "$restore_path" ] || [ ! -f "$restore_path" ]; then
    echo "Backup not found: $source_name" >&2
    exit 1
fi

if [ ! -f "$restore_path.sha256" ]; then
    echo "Checksum not found for: $restore_path" >&2
    exit 1
fi

(
    cd "$(dirname "$restore_path")"
    sha256sum -c "$(basename "$restore_path").sha256"
)
pg_restore --list "$restore_path" >/dev/null

if [ "$(dirname "$restore_path")" != "$temporary_directory" ]; then
    cp "$restore_path" "$temporary_directory/"
    cp "$restore_path.sha256" "$temporary_directory/"
    restore_path="$temporary_directory/$(basename "$restore_path")"
fi

export PGPASSWORD="$DB_PASSWORD"
export PGCONNECT_TIMEOUT="${PGCONNECT_TIMEOUT:-10}"

if [ "${RESTORE_SKIP_SAFETY_BACKUP:-false}" != "true" ]; then
    echo "Creating a safety backup before restoration."
    /usr/local/bin/backup-database pre-restore
fi

timestamp="$(date -u +"%Y%m%dT%H%M%SZ")"
database_prefix="$(printf "%s" "$DB_DATABASE" | cut -c1-28)"
restore_database="${database_prefix}_restore_${timestamp}_$$"
previous_database="${database_prefix}_before_${timestamp}_$$"

createdb \
    --host="$DB_HOST" \
    --port="${DB_PORT:-5432}" \
    --username="$DB_USERNAME" \
    --owner="$DB_USERNAME" \
    "$restore_database"

restore_failed=true
current_database_renamed=false

cleanup_restore_database() {
    if [ "$restore_failed" = "true" ]; then
        if [ "$current_database_renamed" = "true" ]; then
            psql \
                --host="$DB_HOST" \
                --port="${DB_PORT:-5432}" \
                --username="$DB_USERNAME" \
                --dbname=postgres \
                --command="ALTER DATABASE \"$previous_database\" RENAME TO \"$DB_DATABASE\";" \
                >/dev/null 2>&1 || true
        fi

        dropdb \
            --host="$DB_HOST" \
            --port="${DB_PORT:-5432}" \
            --username="$DB_USERNAME" \
            --force \
            --if-exists \
            "$restore_database" 2>/dev/null || true
    fi

    cleanup
}

trap cleanup_restore_database EXIT INT TERM

echo "Restoring and validating temporary database $restore_database."
pg_restore \
    --host="$DB_HOST" \
    --port="${DB_PORT:-5432}" \
    --username="$DB_USERNAME" \
    --dbname="$restore_database" \
    --exit-on-error \
    --no-owner \
    --no-acl \
    "$restore_path"
vacuumdb \
    --host="$DB_HOST" \
    --port="${DB_PORT:-5432}" \
    --username="$DB_USERNAME" \
    --dbname="$restore_database" \
    --analyze-in-stages

database_exists="$(psql \
    --host="$DB_HOST" \
    --port="${DB_PORT:-5432}" \
    --username="$DB_USERNAME" \
    --dbname=postgres \
    --tuples-only \
    --no-align \
    --command="SELECT 1 FROM pg_database WHERE datname = '$DB_DATABASE';")"

if [ "$database_exists" = "1" ]; then
    psql \
        --host="$DB_HOST" \
        --port="${DB_PORT:-5432}" \
        --username="$DB_USERNAME" \
        --dbname=postgres \
        --set=ON_ERROR_STOP=1 \
        --command="SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$DB_DATABASE' AND pid <> pg_backend_pid();" \
        --command="ALTER DATABASE \"$DB_DATABASE\" RENAME TO \"$previous_database\";"
    current_database_renamed=true
fi

psql \
    --host="$DB_HOST" \
    --port="${DB_PORT:-5432}" \
    --username="$DB_USERNAME" \
    --dbname=postgres \
    --set=ON_ERROR_STOP=1 \
    --command="ALTER DATABASE \"$restore_database\" RENAME TO \"$DB_DATABASE\";"

restore_failed=false
trap cleanup EXIT INT TERM

echo "Database restoration completed successfully."
if [ "$database_exists" = "1" ]; then
    echo "Previous database kept for rollback as: $previous_database"
fi