#!/bin/bash

# shellcheck disable=SC2034
RESYNC_FAILED_UNTOUCHED=1
# shellcheck disable=SC2034
RESYNC_FAILED_ERASED=2
INITIAL_COPY_TIMEOUT=300
INITIAL_COPY_POLL_INTERVAL=5
MDS_STATS_TABLES="cel queue_log"
STOPPED_DB_START_TIMEOUT=120
SUBSCRIPTION_RESTART_TIMEOUT=60
MDS_CUSTOM_ENV_FILE=/etc/docker/mds/custom.env
PGPASS_FILE=/root/.pgpass
MAIN_DB_CONNECT_TIMEOUT=10

get_db_image() {
    local resolved_images

    resolved_images=$(xivo-dcomp config --images db 2> /dev/null) || return 1
    if [ -z "${resolved_images}" ]; then
        return 1
    fi

    echo "${resolved_images}" | head -n 1
}

count_subscriptions() {
    psql -U postgres -d asterisk -Aqt -c "SELECT COUNT(*) FROM pg_subscription;" 2> /dev/null
}

count_relations_still_copying() {
    psql -U postgres -d asterisk -Aqt -c "SELECT COUNT(*) FROM pg_subscription_rel WHERE srsubstate <> 'r';" 2> /dev/null
}

check_mds_stats_replication() {
    local force="${1}"; shift
    local ignore_stats="${1}"; shift
    local db_rebuilt="${1}"; shift

    report_stats_replication "${MDS_STATS_TABLES}" && return 0
    display_stalled_stats_replication_warning "${db_rebuilt}"
    ! stats_replication_blocks_upgrade "${force}" "${ignore_stats}" "${db_rebuilt}"
}

get_xivo_host() {
    if [ ! -f "${MDS_CUSTOM_ENV_FILE}" ]; then
        return 0
    fi

    grep -oP '^\s*XIVO_HOST=\K.*' "${MDS_CUSTOM_ENV_FILE}" | tail -n 1
}

ensure_main_pgpass_line() {
    local xivo_host="${1}"; shift
    local main_line="${xivo_host}:5432:asterisk:asterisk:proformatique"

    touch "${PGPASS_FILE}"
    chmod 600 "${PGPASS_FILE}"
    grep -Fxq "${main_line}" "${PGPASS_FILE}" || echo "${main_line}" >> "${PGPASS_FILE}"
}

get_main_alembic_version() {
    local xivo_host="${1}"; shift

    PGCONNECT_TIMEOUT="${MAIN_DB_CONNECT_TIMEOUT}" psql -h "${xivo_host}" -p 5432 -U asterisk -d asterisk -w -Aqt -c "SELECT version_num FROM alembic_version;"
}

display_upgrade_order_error() {
    local main_version="${1}"; shift
    local image_head="${1}"; shift

	cat <<-EOF
	******************************************************************************
	*  ERROR: the XiVO Main was not upgraded yet.                                *
	*                                                                            *
	*  Recreating the database now would subscribe a brand new schema to an      *
	*  old Main: the tables added by the new migrations would stay empty         *
	*  forever. The database of this media server was left untouched.            *
	*                                                                            *
	*  You MUST:                                                                 *
	*    1. upgrade the XiVO Main (xivo-upgrade), then the XiVOCC,               *
	*    2. and then, re-run this media server upgrade.                          *
	******************************************************************************
	EOF
    echo "  XiVO Main schema revision   : ${main_version}"
    echo "  New database image revision : ${image_head}"
}

display_xivo_host_missing_error() {
	cat <<-EOF
	******************************************************************************
	*  ERROR: XIVO_HOST is not set in /etc/docker/mds/custom.env.                *
	*                                                                            *
	*  It is needed to read the schema revision of the XiVO Main. The database   *
	*  of this media server was left untouched.                                  *
	******************************************************************************
	EOF
}

display_main_unreachable_error() {
    local xivo_host="${1}"; shift

	cat <<-EOF
	******************************************************************************
	*  ERROR: could not read the schema revision of the XiVO Main.               *
	*                                                                            *
	*  The XiVO Main must be up and already upgraded before a media server is    *
	*  resynchronized. The database of this media server was left untouched.     *
	*                                                                            *
	*  The revision is read from this host, with the XiVO Main credentials of    *
	*  /root/.pgpass: the XiVO Main database must be reachable and accept them.  *
	*  Check the error above, then test the connection by hand with:             *
	******************************************************************************
	EOF
    echo "  psql -h ${xivo_host} -U asterisk asterisk"
}

check_upgrade_order() {
    local image="${1}"; shift
    local image_head xivo_host main_version

    image_head=$(get_image_alembic_head "${image}")
    if [ -z "${image_head}" ]; then
        echo "ERROR: could not read the alembic head embedded in ${image}." >&2
        return 1
    fi

    xivo_host=$(get_xivo_host)
    if [ -z "${xivo_host}" ]; then
        display_xivo_host_missing_error
        return 1
    fi

    ensure_main_pgpass_line "${xivo_host}"
    main_version=$(get_main_alembic_version "${xivo_host}")
    if [ -z "${main_version}" ]; then
        display_main_unreachable_error "${xivo_host}"
        return 1
    fi

    if [ "${main_version}" != "${image_head}" ]; then
        display_upgrade_order_error "${main_version}" "${image_head}"
        return 1
    fi

    echo "XiVO Main schema revision (${main_version}) matches the new database image."
}

release_main_replication_slot() {
	xivo-dcomp exec -T db bash -s <<-'EOF'
	set -e
	run_local() { psql -U postgres -d asterisk -Aqt -v ON_ERROR_STOP=1 "$@"; }
	subscription=$(run_local -c "SELECT subname FROM pg_subscription LIMIT 1;")
	[ -n "${subscription}" ] || exit 0
	slot=$(run_local -c "SELECT subslotname FROM pg_subscription WHERE subname = '${subscription}';")
	conninfo=$(run_local -c "SELECT subconninfo FROM pg_subscription WHERE subname = '${subscription}';")
	slots_on_main=0
	if [ -n "${slot}" ]; then
	    slots_on_main=$(psql "${conninfo}" -Aqt -v ON_ERROR_STOP=1 -c "SELECT count(*) FROM pg_replication_slots WHERE slot_name = '${slot}';")
	fi
	if [ "${slots_on_main}" -eq 0 ]; then
	    run_local -c "ALTER SUBSCRIPTION \"${subscription}\" DISABLE;"
	    run_local -c "ALTER SUBSCRIPTION \"${subscription}\" SET (slot_name = NONE);"
	fi
	run_local -c "DROP SUBSCRIPTION \"${subscription}\";"
	EOF
}

release_main_replication_slot_if_db_answers() {
    is_postgres_running_on_port || return 0

    display_upgrade_notice "Releasing the replication slot of this media server on the XiVO Main..."
    release_main_replication_slot
}

is_main_slot_present() {
	local slots_on_main

	slots_on_main=$(xivo-dcomp exec -T db bash -s 2> /dev/null <<-'EOF'
	slot=$(psql -U postgres -d asterisk -Aqt -c "SELECT subslotname FROM pg_subscription LIMIT 1;")
	conninfo=$(psql -U postgres -d asterisk -Aqt -c "SELECT subconninfo FROM pg_subscription LIMIT 1;")
	[ -n "${slot}" ] && psql "${conninfo}" -Aqt -c "SELECT count(*) FROM pg_replication_slots WHERE slot_name = '${slot}';"
	EOF
	)

	[ -n "${slots_on_main}" ] && [ "${slots_on_main}" -gt 0 ]
}

is_subscription_streaming() {
    local streaming

    streaming=$(psql -U postgres -d asterisk -Aqt -c "SELECT COUNT(*) FROM pg_subscription s JOIN pg_stat_subscription w ON w.subid = s.oid WHERE s.subenabled AND w.pid IS NOT NULL AND w.relid IS NULL;" 2> /dev/null)

    [ -n "${streaming}" ] && [ "${streaming}" -gt 0 ]
}

is_initial_copy_done() {
    local subscriptions still_copying

    subscriptions=$(count_subscriptions)
    if [ -z "${subscriptions}" ] || [ "${subscriptions}" -eq 0 ]; then
        return 1
    fi

    still_copying=$(count_relations_still_copying)
    [ -n "${still_copying}" ] && [ "${still_copying}" -eq 0 ]
}

wait_for_started_db() {
    local waited

    echo -n "Waiting up to ${STOPPED_DB_START_TIMEOUT}s for the database..."
    for ((waited = 0; waited < STOPPED_DB_START_TIMEOUT; waited++)); do
        if is_postgres_running_on_port; then
            echo
            return 0
        fi
        sleep 1
        echo -n "."
    done
    echo
    return 1
}

is_db_container_present() {
    [ -n "$(xivo-dcomp ps -aq db 2> /dev/null)" ]
}

ask_to_start_stopped_db() {
    local answer

    read -r -p 'Would you like to start it [N/y]? ' answer
    answer="${answer:-N}"
    [ "${answer}" = 'y' ] || [ "${answer}" = 'Y' ]
}

start_db_if_stopped() {
    local force="${1}"; shift
    local ignore_stats="${1}"; shift

    if is_postgres_running_on_port; then
        return 0
    fi

    display_upgrade_notice "The database of this media server is stopped."
    echo "It is started to check its replication and its statistics, and to release its slot."
    echo "The other services stay stopped."
    if [ "${force}" -eq 0 ] && [ "${ignore_stats}" -eq 0 ] && ! ask_to_start_stopped_db; then
        return 1
    fi

    if ! is_db_container_present || ! xivo-dcomp start db || ! wait_for_started_db; then
        echo "ERROR: the database of this media server could not be started, check it with: xivo-dcomp logs db" >&2
        if [ "${ignore_stats}" -eq 0 ]; then
            echo "       Run again with -s to go on without it: the database will be rebuilt." >&2
            return 1
        fi
        echo "WARNING: going on with -s: the database will be rebuilt." >&2
        return 0
    fi
    wait_for_restarted_subscription
}

wait_for_restarted_subscription() {
    local waited

    if ! is_main_slot_present; then
        return 0
    fi

    echo -n "Waiting up to ${SUBSCRIPTION_RESTART_TIMEOUT}s for the replication from the XiVO Main to restart..."
    for ((waited = 0; waited < SUBSCRIPTION_RESTART_TIMEOUT; waited++)); do
        if is_subscription_streaming; then
            echo
            return 0
        fi
        sleep 1
        echo -n "."
    done
    echo
}

log_initial_copy_duration() {
    local rebuild_started="${1}"; shift

    echo "Database rebuild: $(( $(date +%s) - rebuild_started ))s between the start of the rebuild and the end of the initial copy."
}

display_initial_copy_timeout_warning() {
    cmn_echo_important "WARNING: the initial replication copy is still running after ${INITIAL_COPY_TIMEOUT}s."
    cmn_echo_important "         Starting the services anyway; the copy is checked once more after they"
    cmn_echo_important "         start, otherwise the asterisk configuration must be reloaded by hand."
}

display_initial_copy_unfinished_warning() {
    cmn_echo_important "WARNING: the initial replication copy is STILL not finished."
    cmn_echo_important "         Once it is (no row left in pg_subscription_rel with srsubstate <> 'r'),"
    cmn_echo_important "         reload the asterisk configuration with:"
    cmn_echo_important "           /usr/share/mds-upgrade/post-start.d/80-reload-asterisk-configuration"
}

wait_for_initial_copy() {
    local rebuild_started="${1}"; shift
    local waited=0

    [ -n "${rebuild_started}" ] || return 0

    echo -n "Waiting up to ${INITIAL_COPY_TIMEOUT}s for the initial replication copy to complete..."
    while [ "${waited}" -lt "${INITIAL_COPY_TIMEOUT}" ]; do
        if is_initial_copy_done; then
            echo
            log_initial_copy_duration "${rebuild_started}"
            return 0
        fi
        sleep "${INITIAL_COPY_POLL_INTERVAL}"
        waited=$(( waited + INITIAL_COPY_POLL_INTERVAL ))
        echo -n "."
    done
    echo
    display_initial_copy_timeout_warning
    MDS_INITIAL_COPY_TIMED_OUT=1
}

catch_up_initial_copy() {
    local rebuild_started="${1}"; shift

    [ "${MDS_INITIAL_COPY_TIMED_OUT}" = "1" ] || return 0

    echo "Checking again whether the initial replication copy completed..."
    if is_initial_copy_done; then
        log_initial_copy_duration "${rebuild_started}"
        display_upgrade_notice "Initial copy completed, regenerating the asterisk configuration..."
        /usr/share/mds-upgrade/post-start.d/80-reload-asterisk-configuration
    else
        display_initial_copy_unfinished_warning
    fi
}

display_mds_database_recovery() {
    echo "******************************************************************************"
    format_box_line "The services of this media server were left stopped."
    format_box_line ""
    format_box_line "Rebuild its database once the cause above is fixed:"
    format_box_line "  mds-resync-db -f -s"
    format_box_line "then run mds-upgrade again if the upgrade did not finish."
    echo "******************************************************************************"
}
