#!/usr/bin/env bash
set -Eeuo pipefail

die() { echo "postgres driver: $*" >&2; exit 1; }

load_config() {
    : "${DRIVER_CONFIG:?DRIVER_CONFIG is required}"
    [[ -r "$DRIVER_CONFIG" ]] || die "cannot read $DRIVER_CONFIG"
    set -a
    # shellcheck disable=SC1090
    source "$DRIVER_CONFIG"
    if [[ -n "${DRIVER_SECRET_FILE:-}" ]]; then
        [[ -r "$DRIVER_SECRET_FILE" ]] || die "cannot read $DRIVER_SECRET_FILE"
        # shellcheck disable=SC1090
        source "$DRIVER_SECRET_FILE"
    fi
    set +a
    : "${PGHOST:?PGHOST is required}"; : "${PGPORT:?PGPORT is required}"; : "${PGUSER:?PGUSER is required}"
    PG_BACKUP_MODE="${PG_BACKUP_MODE:-physical}"
    PGDATABASE="${PGDATABASE:-postgres}"
    PG_COMPRESSION="${PG_COMPRESSION:-client-gzip:level=6}"
    PG_CHECKPOINT="${PG_CHECKPOINT:-fast}"
    PG_LOGICAL_FORMAT="${PG_LOGICAL_FORMAT:-custom}"
    PG_LOGICAL_FILE_NAMING="${PG_LOGICAL_FILE_NAMING:-sha256}"
    PG_LOGICAL_COMPRESSION="${PG_LOGICAL_COMPRESSION:-6}"
    PG_LOGICAL_GLOBALS_GZIP_LEVEL="${PG_LOGICAL_GLOBALS_GZIP_LEVEL:-6}"
    PG_DATABASE_INCLUDE="${PG_DATABASE_INCLUDE:-}"
    PG_DATABASE_EXCLUDE="${PG_DATABASE_EXCLUDE:-}"
    [[ "$PG_BACKUP_MODE" == physical || "$PG_BACKUP_MODE" == logical ]] || \
        die "PG_BACKUP_MODE must be physical or logical"
    [[ "$PG_LOGICAL_FORMAT" == custom || "$PG_LOGICAL_FORMAT" == plain ]] || \
        die "PG_LOGICAL_FORMAT must be custom or plain"
    [[ "$PG_LOGICAL_FILE_NAMING" == sha256 || "$PG_LOGICAL_FILE_NAMING" == plain ]] || \
        die "PG_LOGICAL_FILE_NAMING must be sha256 or plain"
}

postgres_ready() {
    pg_isready -q -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$PGDATABASE" || \
        die "local PostgreSQL is unavailable"
}

prepare_physical() {
    local destination="$1"
    pg_basebackup --host="$PGHOST" --port="$PGPORT" --username="$PGUSER" \
        --pgdata="$destination" --format=tar --wal-method=stream \
        --compress="$PG_COMPRESSION" --checkpoint="$PG_CHECKPOINT" \
        --manifest-checksums=SHA256 --label="$BACKUP_NAME"
}

database_is_listed() {
    local database="$1" list="$2" item
    while IFS= read -r item || [[ -n "$item" ]]; do
        [[ -n "$item" && "$item" == "$database" ]] && return 0
    done <<<"$list"
    return 1
}

validate_includes() {
    local available="$1" requested
    [[ -n "$PG_DATABASE_INCLUDE" ]] || return 0
    while IFS= read -r requested || [[ -n "$requested" ]]; do
        [[ -n "$requested" ]] || continue
        database_is_listed "$requested" "$available" || \
            die "included database does not exist or cannot be dumped: $requested"
    done <<<"$PG_DATABASE_INCLUDE"
}

prepare_logical() {
    local destination="$1" database_output database digest dump_file extension index temporary count=0
    local -a dump_format_arguments
    case "$PG_LOGICAL_FORMAT" in
        custom)
            extension=dump
            dump_format_arguments=(--format=custom --compress="$PG_LOGICAL_COMPRESSION")
            ;;
        plain)
            extension=sql
            dump_format_arguments=(--format=plain)
            ;;
    esac
    database_output="$(psql --no-psqlrc --tuples-only --no-align \
        --host="$PGHOST" --port="$PGPORT" --username="$PGUSER" \
        --dbname="$PGDATABASE" --set=ON_ERROR_STOP=1 \
        --command="SELECT datname FROM pg_database WHERE datallowconn AND NOT datistemplate ORDER BY datname")" || \
        die "could not list PostgreSQL databases"
    validate_includes "$database_output"

    mkdir -p "$destination/databases"
    index="$destination/databases.json"
    printf '{"format":"backmaster-postgres-logical-v1","dump_format":"%s","file_naming":"%s","databases":[]}\n' \
        "$PG_LOGICAL_FORMAT" "$PG_LOGICAL_FILE_NAMING" >"$index"
    while IFS= read -r database || [[ -n "$database" ]]; do
        [[ -n "$database" ]] || continue
        if [[ -n "$PG_DATABASE_INCLUDE" ]] && \
            ! database_is_listed "$database" "$PG_DATABASE_INCLUDE"; then
            continue
        fi
        if database_is_listed "$database" "$PG_DATABASE_EXCLUDE"; then
            continue
        fi

        case "$PG_LOGICAL_FILE_NAMING" in
            sha256)
                digest="$(printf '%s' "$database" | sha256sum)"
                dump_file="${digest%% *}.${extension}"
                ;;
            plain)
                [[ "$database" != *'/'* && "$database" != *$'\n'* && \
                    "$database" != *$'\r'* ]] || \
                    die "database name cannot be used as a plain filename: $database"
                dump_file="${database}.${extension}"
                ;;
        esac
        pg_dump --host="$PGHOST" --port="$PGPORT" --username="$PGUSER" \
            --dbname="$database" "${dump_format_arguments[@]}" \
            --file="$destination/databases/$dump_file"
        temporary="${index}.tmp"
        jq --arg database "$database" --arg file "databases/$dump_file" \
            '.databases += [{database: $database, file: $file}]' \
            "$index" >"$temporary"
        mv -- "$temporary" "$index"
        count=$((count + 1))
    done <<<"$database_output"
    ((count > 0)) || die "database filters selected no databases"

    pg_dumpall --host="$PGHOST" --port="$PGPORT" --username="$PGUSER" \
        --database="$PGDATABASE" --globals-only | \
        gzip "-$PG_LOGICAL_GLOBALS_GZIP_LEVEL" -c >"$destination/globals.sql.gz"
}

prepare() {
    local destination="$1"
    [[ -d "$destination" ]] || die "staging payload does not exist: $destination"
    [[ -z "$(find "$destination" -mindepth 1 -print -quit)" ]] || die "staging payload is not empty"
    postgres_ready
    case "$PG_BACKUP_MODE" in
        physical) prepare_physical "$destination" ;;
        logical) prepare_logical "$destination" ;;
    esac
}

connectivitycheck() {
    command -v pg_isready >/dev/null || die "pg_isready is unavailable"
    case "$PG_BACKUP_MODE" in
        physical) command -v pg_basebackup >/dev/null || die "pg_basebackup is unavailable" ;;
        logical)
            command -v psql >/dev/null || die "psql is unavailable"
            command -v pg_dump >/dev/null || die "pg_dump is unavailable"
            command -v pg_dumpall >/dev/null || die "pg_dumpall is unavailable"
            command -v jq >/dev/null || die "jq is unavailable"
            ;;
    esac
    postgres_ready
}

healthcheck() {
    connectivitycheck
    echo "OK: local PostgreSQL can produce a $PG_BACKUP_MODE backup"
}

wal_archive() {
    local wal_path="$1" temporary
    [[ "$PG_BACKUP_MODE" == physical ]] || die "WAL archival requires PG_BACKUP_MODE=physical"
    : "${EXPORTER_EXEC:?EXPORTER_EXEC is required}"
    mkdir -p "${STAGING_ROOT}/${INSTANCE_NAME}"
    temporary="$(mktemp "${STAGING_ROOT}/${INSTANCE_NAME}/.wal.XXXXXX.gz")"
    if ! gzip -c "$wal_path" >"$temporary" || \
        ! "$EXPORTER_EXEC" put-file "$temporary" "wal/$(basename "$wal_path").gz"; then
        rm -f -- "$temporary"; return 1
    fi
    rm -f -- "$temporary"
}

wal_restore() {
    local wal_name="$1" destination="$2" temporary
    [[ "$PG_BACKUP_MODE" == physical ]] || die "WAL restore requires PG_BACKUP_MODE=physical"
    : "${EXPORTER_EXEC:?EXPORTER_EXEC is required}"
    mkdir -p "${STAGING_ROOT}/${INSTANCE_NAME}"
    temporary="$(mktemp "${STAGING_ROOT}/${INSTANCE_NAME}/.wal-restore.XXXXXX.gz")"
    if ! "$EXPORTER_EXEC" get-file "wal/${wal_name}.gz" "$temporary" || \
        ! gzip -dc "$temporary" >"$destination"; then
        rm -f -- "$temporary"; return 1
    fi
    rm -f -- "$temporary"
}

main() {
    load_config
    case "${1:-}" in
        prepare) [[ $# -eq 2 ]] || exit 64; prepare "$2" ;;
        connectivitycheck) connectivitycheck ;;
        healthcheck) healthcheck ;;
        wal-archive) [[ $# -eq 2 ]] || exit 64; wal_archive "$2" ;;
        wal-restore) [[ $# -eq 3 ]] || exit 64; wal_restore "$2" "$3" ;;
        *) die "unsupported command: ${1:-none}" ;;
    esac
}

main "$@"
