# Sourced by managed-container; uses the shared OCF and parameter environment.

mc_pg_rejoin_basebackup_enabled()
{
    case "${REJOIN_ALLOW_BASEBACKUP}" in
        true|1|yes)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

mc_pg_rejoin_basebackup_fallback()
{
    local pgdata="$1"
    local reason="${2:-rejoin fallback requested}"
    local source_bytes free_bytes required_bytes headroom_bytes min_headroom_bytes
    local timestamp preserved_pgdata failed_pgdata source_role_retry

    if ! mc_pg_rejoin_basebackup_enabled; then
        ocf_exit_reason \
            "${reason}; rejoin_allow_basebackup is not enabled"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log warn \
        "${reason}; evaluating pg_basebackup fallback"

    # Revalidate the source immediately before replacing local PGDATA.
    source_role_retry="$(pg_remote_query \
        'SELECT pg_is_in_recovery();' 2>/dev/null)" || {
        ocf_exit_reason \
            "Unable to revalidate rejoin source before pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    }

    if [ "${source_role_retry}" != "f" ]; then
        ocf_exit_reason \
            "Rejoin source is no longer a writable primary; refusing pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    if pg_running; then
        ocf_exit_reason \
            "PostgreSQL is still active; refusing pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    if ! service_running || ! container_running; then
        ocf_exit_reason \
            "Maintenance service/container is unavailable; refusing pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    if ! ensure_remote_slot \
        "${REJOIN_SOURCE_HOST}" \
        "${REPLICATION_SLOT}"
    then
        ocf_exit_reason \
            "Unable to revalidate replication slot ${REPLICATION_SLOT} before pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    # Estimate the replacement payload. Existing PGDATA is renamed in place on
    # the same filesystem, so it consumes no additional free space. Require one
    # replacement copy plus configurable safety headroom.
    source_bytes="$(pg_remote_query \
        "SELECT COALESCE(sum(pg_database_size(datname)),0)
         FROM pg_database
         WHERE datallowconn;" 2>/dev/null)" || {
        ocf_exit_reason \
            "Unable to estimate source database size before pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    }

    case "${source_bytes}" in
        ""|*[!0-9]*)
            ocf_exit_reason \
                "Invalid source database size returned before pg_basebackup fallback"
            return "${OCF_ERR_GENERIC}"
            ;;
    esac

    free_bytes="$(
        pg_command '
            blocks="$(df -k --output=avail "$PGDATA" | tail -n 1 | tr -d "[:space:]")"
            case "$blocks" in
                ""|*[!0-9]*) exit 1 ;;
            esac
            printf "%s\\n" "$((blocks * 1024))"
        ' 2>/dev/null
    )" || {
        ocf_exit_reason \
            "Unable to determine free space for pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    }

    case "${free_bytes}" in
        ""|*[!0-9]*)
            ocf_exit_reason \
                "Invalid free-space value before pg_basebackup fallback"
            return "${OCF_ERR_GENERIC}"
            ;;
    esac

    min_headroom_bytes=$((1024 * 1024 * 1024))
    headroom_bytes=$((source_bytes * REJOIN_BASEBACKUP_HEADROOM_PERCENT / 100))
    if [ "${headroom_bytes}" -lt "${min_headroom_bytes}" ]; then
        headroom_bytes=${min_headroom_bytes}
    fi
    required_bytes=$((source_bytes + headroom_bytes))

    ocf_log info \
        "pg_basebackup space check: source-estimate=${source_bytes} headroom=${headroom_bytes} required=${required_bytes} free=${free_bytes}"

    if [ "${free_bytes}" -lt "${required_bytes}" ]; then
        ocf_exit_reason \
            "Insufficient free space for pg_basebackup fallback: free=${free_bytes} estimated-required=${required_bytes} source-estimate=${source_bytes} headroom=${headroom_bytes}"
        return "${OCF_ERR_GENERIC}"
    fi

    timestamp="$(date -u +%Y%m%d-%H%M%S)"
    preserved_pgdata="${pgdata}.pre-basebackup-${timestamp}"
    failed_pgdata="${pgdata}.failed-basebackup-${timestamp}"

    ocf_log warn \
        "Preserving existing PGDATA as ${preserved_pgdata}"

    if ! pg_command "
        set -e
        [ -d '${pgdata}' ]
        [ ! -e '${preserved_pgdata}' ]
        [ ! -e '${failed_pgdata}' ]
        mv '${pgdata}' '${preserved_pgdata}'
        install -d -m 0700 '${pgdata}'
    "
    then
        ocf_exit_reason \
            "Unable to preserve existing PGDATA for pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log warn \
        "Running pg_basebackup from ${REJOIN_SOURCE_HOST}"

    if ! pg_replication_exec \
        pg_basebackup \
        --dbname="host=${REJOIN_SOURCE_HOST} port=${REJOIN_SOURCE_PORT} user=${REPLICATION_USER} application_name=${APPLICATION_NAME}" \
        --pgdata="${pgdata}" \
        --format=plain \
        --wal-method=stream \
        --slot="${REPLICATION_SLOT}" \
        --write-recovery-conf \
        --progress
    then
        ocf_log err \
            "pg_basebackup failed; preserving failed copy and restoring previous PGDATA"

        if ! pg_command "
            set -e
            [ ! -e '${failed_pgdata}' ]
            if [ -e '${pgdata}' ]; then
                mv '${pgdata}' '${failed_pgdata}'
            fi
            mv '${preserved_pgdata}' '${pgdata}'
        "
        then
            ocf_exit_reason \
                "pg_basebackup failed and automatic PGDATA rollback also failed; manual recovery required"
            return "${OCF_ERR_GENERIC}"
        fi

        ocf_exit_reason \
            "pg_basebackup failed; original PGDATA restored and failed copy preserved as ${failed_pgdata}"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log info \
        "pg_basebackup fallback completed; previous PGDATA preserved as ${preserved_pgdata}"
    return "${OCF_SUCCESS}"
}

mc_pg_rejoin_configure_standby()
{
    if [ -n "${OCF_RESKEY_replication_peers:-}" ]; then
        mc_pg_write_peer_recovery || {
            ocf_exit_reason "Failed to configure native peer recovery"
            return "${OCF_ERR_GENERIC}"
        }
        return "${OCF_SUCCESS}"
    fi

    if ! mc_runtime_exec \
        --user "${PG_OS_USER}" \
        --env "RA_PRIMARY_HOST=${REJOIN_SOURCE_HOST}" \
        --env "RA_PRIMARY_PORT=${REJOIN_SOURCE_PORT}" \
        --env "RA_REPLICATION_USER=${REPLICATION_USER}" \
        --env "RA_REPLICATION_SLOT=${REPLICATION_SLOT}" \
        --env "RA_APPLICATION_NAME=${APPLICATION_NAME}" \
        "${CONTAINER_NAME}" \
        sh -c '
            set -e
            sed -i \\
                -e "/^[[:space:]]*primary_conninfo[[:space:]]*=/d" \\
                -e "/^[[:space:]]*primary_slot_name[[:space:]]*=/d" \\
                "$PGDATA/postgresql.auto.conf"
            printf "%s\\n" \\
                "primary_conninfo = '\''host=${RA_PRIMARY_HOST} port=${RA_PRIMARY_PORT} user=${RA_REPLICATION_USER} application_name=${RA_APPLICATION_NAME}'\''" \\
                "primary_slot_name = '\''${RA_REPLICATION_SLOT}'\''" \\
                >> "$PGDATA/postgresql.auto.conf"
            touch "$PGDATA/standby.signal"
            chmod 600 "$PGDATA/postgresql.auto.conf" "$PGDATA/standby.signal"
        '
    then
        ocf_exit_reason "Failed to configure standby settings"
        return "${OCF_ERR_GENERIC}"
    fi

    return "${OCF_SUCCESS}"
}

mc_pg_rejoin_start_and_wait_streaming()
{
    local role receiver_status i

    ocf_log info "Starting PostgreSQL as standby"

    if ! pg_command '
        BINDIR="$(pg_config --bindir)" || exit 1
        exec "$BINDIR/pg_ctl" \
            -D "$PGDATA" \
            -l "$PGDATA/postgresql.log" \
            start -w
    '
    then
        ocf_exit_reason "PostgreSQL failed to start after rejoin"
        return "${OCF_ERR_GENERIC}"
    fi

    for i in $(seq 1 30); do
        role="$(pg_role 2>/dev/null)" || {
            sleep 1
            continue
        }
        receiver_status="$(pg_wal_receiver_status 2>/dev/null || true)"

        if [ "${role}" = "t" ] && [ "${receiver_status}" = "streaming" ]; then
            return "${OCF_SUCCESS}"
        fi
        sleep 1
    done

    return "${OCF_ERR_GENERIC}"
}

managed_rejoin()
{
    local source_role pgdata
    local fallback_used=false

    case "${REJOIN_ALLOW_BASEBACKUP}" in
        true|false|1|0|yes|no) ;;
        *)
            ocf_exit_reason "rejoin_allow_basebackup must be true or false"
            return "${OCF_ERR_CONFIGURED}"
            ;;
    esac

    if [ -z "${REJOIN_SOURCE_HOST}" ]; then
        ocf_exit_reason "rejoin_source_host is required for rejoin"
        return "${OCF_ERR_CONFIGURED}"
    fi

    if ! service_running; then
        ocf_exit_reason "Maintenance service ${SERVICE_NAME} must be running before rejoin"
        return "${OCF_ERR_GENERIC}"
    fi

    if ! container_running; then
        ocf_exit_reason "Maintenance container ${CONTAINER_NAME} must be running before rejoin"
        return "${OCF_ERR_GENERIC}"
    fi

    mc_pg_require_replication_slot || return "$?"

    if pg_running; then
        ocf_exit_reason "PostgreSQL must be stopped before rejoin"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log info "Verifying rejoin source ${REJOIN_SOURCE_HOST}:${REJOIN_SOURCE_PORT}"
    source_role="$(pg_remote_query 'SELECT pg_is_in_recovery();' 2>/dev/null)" || {
        ocf_exit_reason "Unable to connect to rejoin source ${REJOIN_SOURCE_HOST}"
        return "${OCF_ERR_GENERIC}"
    }

    if [ "${source_role}" != "f" ]; then
        ocf_exit_reason "Rejoin source ${REJOIN_SOURCE_HOST} is not a writable primary"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log info "Ensuring replication slot ${REPLICATION_SLOT} on rejoin source ${REJOIN_SOURCE_HOST}"
    if ! ensure_remote_slot "${REJOIN_SOURCE_HOST}" "${REPLICATION_SLOT}"; then
        ocf_exit_reason "Unable to ensure physical replication slot ${REPLICATION_SLOT} on rejoin source ${REJOIN_SOURCE_HOST}"
        return "${OCF_ERR_GENERIC}"
    fi

    pgdata="$(pgdata_path 2>/dev/null)" || {
        ocf_exit_reason "Unable to determine PostgreSQL PGDATA inside ${CONTAINER_NAME}"
        return "${OCF_ERR_GENERIC}"
    }
    if [ -z "${pgdata}" ]; then
        ocf_exit_reason "PostgreSQL PGDATA is empty inside ${CONTAINER_NAME}"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log info "Running pg_rewind against ${REJOIN_SOURCE_HOST}"
    if ! pg_replication_exec \
        pg_rewind \
        --target-pgdata="${pgdata}" \
        --source-server="host=${REJOIN_SOURCE_HOST} port=${REJOIN_SOURCE_PORT} dbname=postgres user=${REPLICATION_USER}" \
        --progress
    then
        mc_pg_rejoin_basebackup_fallback "${pgdata}" "pg_rewind failed" || return "$?"
        fallback_used=true
    fi

    ocf_log info "Configuring PostgreSQL as standby of ${REJOIN_SOURCE_HOST}"
    mc_pg_rejoin_configure_standby || return "$?"

    if mc_pg_rejoin_start_and_wait_streaming; then
        pg_clear_rejoin_pending >/dev/null 2>&1 || true
        ocf_log info "PostgreSQL successfully rejoined as streaming standby"
        return "${OCF_SUCCESS}"
    fi

    # A successful pg_rewind does not guarantee that the WAL required to join
    # the new timeline is still available. If the standby cannot reach
    # streaming, stop it and perform the same guarded full-resync fallback.
    if [ "${fallback_used}" = "true" ]; then
        ocf_exit_reason \
            "PostgreSQL did not become a streaming standby even after pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    if ! mc_pg_rejoin_basebackup_enabled; then
        ocf_exit_reason \
            "PostgreSQL started in recovery but did not become streaming; rejoin_allow_basebackup is not enabled"
        return "${OCF_ERR_GENERIC}"
    fi

    ocf_log warn \
        "PostgreSQL entered recovery but did not reach streaming after pg_rewind; stopping it for pg_basebackup fallback"

    # The failed post-rewind start may already have exited. Stopping an
    # already-stopped server is success for fallback purposes.
    if pg_running; then
        if ! pg_command '
            BINDIR="$(pg_config --bindir)" || exit 1
            exec "$BINDIR/pg_ctl" -D "$PGDATA" stop -m fast -w
        '
        then
            ocf_exit_reason \
                "Unable to stop non-streaming standby before pg_basebackup fallback"
            return "${OCF_ERR_GENERIC}"
        fi
    else
        ocf_log info \
            "PostgreSQL is already stopped; proceeding directly to pg_basebackup fallback"
    fi

    mc_pg_rejoin_basebackup_fallback \
        "${pgdata}" \
        "pg_rewind completed but standby could not establish WAL streaming" || return "$?"

    ocf_log info "Reconfiguring PostgreSQL standby after pg_basebackup fallback"
    mc_pg_rejoin_configure_standby || return "$?"

    if ! mc_pg_rejoin_start_and_wait_streaming; then
        ocf_exit_reason \
            "PostgreSQL did not become a streaming standby after pg_basebackup fallback"
        return "${OCF_ERR_GENERIC}"
    fi

    pg_clear_rejoin_pending >/dev/null 2>&1 || true
    ocf_log info \
        "PostgreSQL successfully rejoined as streaming standby after pg_basebackup fallback"
    return "${OCF_SUCCESS}"
}