#!/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 the upgrade order, its replication and its statistics."
    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 "******************************************************************************"
}
