#!/bin/bash

BACKUP_DIR="/var/tmp/mds-resync-db"
PG_HBA_BACKUP="${BACKUP_DIR}/pg_hba.conf"
REBUILT_DB_WAIT_TIMEOUT=600
UNSENT_STATS_TABLES="cel queue_log"
CEL_SENT_COLUMNS="eventtype,eventtime,userdeftype,cid_name,cid_num,cid_ani,cid_rdnis,cid_dnid,exten,context,channame,appname,appdata,amaflags,accountcode,peeraccount,uniqueid,linkedid,userfield,peer,extra"
QUEUE_LOG_SENT_COLUMNS="time,callid,queuename,agent,event,data1,data2,data3,data4,data5"

#shellcheck disable=SC1091
source /usr/bin/xivo-upgrade-functions
#shellcheck disable=SC1091
source /usr/bin/mds-upgrade-functions

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
}

get_image_env_value() {
    local image="${1}"; shift
    local key="${1}"; shift

    docker image inspect --format '{{range .Config.Env}}{{println .}}{{end}}' "${image}" 2> /dev/null \
        | grep -oP -m 1 "^${key}=\K.*"
}

get_image_alembic_head() {
    local image="${1}"; shift

    docker run --rm --entrypoint /bin/bash "${image}" \
        -c 'cd /usr/share/xivo-db && alembic heads' 2> /dev/null \
        | head -n 1 | cut -d ' ' -f 1
}

get_main_alembic_version_through_replication() {
	xivo-dcomp exec -T db bash -s 2> /dev/null <<-'EOF' | head -n 1
	conninfo=$(psql -U postgres -d asterisk -Aqt -c "SELECT subconninfo FROM pg_subscription LIMIT 1;")
	[ -n "${conninfo}" ] && psql "${conninfo}" -Aqt -c "SELECT version_num FROM alembic_version;"
	EOF
}

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_main_unreachable_error() {
	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 through the replication connection: the db           *
	*  container must be running and subscribed to the XiVO Main. When it is     *
	*  not, check the XiVO Main by hand and re-run mds-resync-db with -f.        *
	******************************************************************************
	EOF
}

check_upgrade_order() {
    local image="${1}"; shift
    local image_head 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

    main_version=$(get_main_alembic_version_through_replication)
    if [ -z "${main_version}" ]; then
        display_main_unreachable_error
        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."
}

normalize_locale() {
    tr '[:upper:]' '[:lower:]' | sed -e 's/utf-8/utf8/g'
}

warn_about_locale_change() {
    local image="${1}"; shift
    local cluster_locale="${1}"; shift
    local image_locale

    if [ -z "${cluster_locale}" ]; then
        echo "Note: current cluster locale unknown, skipping the locale check."
        return 0
    fi

    image_locale=$(get_image_env_value "${image}" "LANG")
    if [ -z "${image_locale}" ]; then
        echo "Note: locale of ${image} unknown, skipping the locale check."
        return 0
    fi

    if [ "$(echo "${cluster_locale}" | normalize_locale)" != "$(echo "${image_locale}" | normalize_locale)" ]; then
        echo "Note: the database cluster locale (${cluster_locale}) differs from the one of the image (${image_locale})."
        echo "      The recreated cluster will use ${image_locale}."
    fi
}

get_pgdata_dir_from_image() {
    local image="${1}"; shift

    get_image_env_value "${image}" "PGDATA"
}

backup_pg_hba() {
    local pgdata="${1}"; shift

    if [ -f "${PG_HBA_BACKUP}" ]; then
        echo "Keeping ${PG_HBA_BACKUP}, saved by a resynchronization that did not restore it"
        return 0
    fi
    if [ ! -f "${pgdata}/pg_hba.conf" ]; then
        return 0
    fi
    mkdir -p "${BACKUP_DIR}"
    cp -a "${pgdata}/pg_hba.conf" "${PG_HBA_BACKUP}"
    echo "Saved ${pgdata}/pg_hba.conf to ${PG_HBA_BACKUP}"
}

restore_pg_hba() {
    local pgdata="${1}"; shift

    if [ ! -f "${PG_HBA_BACKUP}" ]; then
        return 0
    fi

    cp -a "${PG_HBA_BACKUP}" "${pgdata}/pg_hba.conf"
    echo "Restored ${PG_HBA_BACKUP} to ${pgdata}/pg_hba.conf"
    rm -f "${PG_HBA_BACKUP}"
    xivo-dcomp reload db
}

get_sent_columns() {
    local table="${1}"; shift

    case "${table}" in
        cel) echo "${CEL_SENT_COLUMNS}";;
        queue_log) echo "${QUEUE_LOG_SENT_COLUMNS}";;
    esac
}

get_unsent_stats_file() {
    local table="${1}"; shift

    echo "${BACKUP_DIR}/${table}.unsent"
}

save_unsent_stats_of_table() {
    local table="${1}"; shift
    local columns file new_rows_file

    columns=$(get_sent_columns "${table}")
    file=$(get_unsent_stats_file "${table}")
    new_rows_file="${file}.new"

    if ! psql -U postgres -d asterisk -q -v ON_ERROR_STOP=1 -c "COPY (SELECT ${columns} FROM ${table} WHERE id > COALESCE((SELECT val FROM replication_state WHERE name = '${table}'), 0) ORDER BY id) TO STDOUT;" > "${new_rows_file}"; then
        rm -f "${new_rows_file}"
        echo "ERROR: could not save the ${table} rows not sent to the XiVOCC yet." >&2
        return 1
    fi

    if [ ! -s "${new_rows_file}" ]; then
        rm -f "${new_rows_file}"
        echo "No ${table} rows waiting to be sent to the XiVOCC."
        return 0
    fi

    echo "Saved $(wc -l < "${new_rows_file}") ${table} rows not sent to the XiVOCC yet to ${file}"
    cat "${new_rows_file}" >> "${file}"
    rm -f "${new_rows_file}"
}

save_unsent_stats() {
    local table

    mkdir -p "${BACKUP_DIR}"
    for table in ${UNSENT_STATS_TABLES}; do
        save_unsent_stats_of_table "${table}" || return 1
    done
}

display_unsent_stats_restore_failure() {
    local table="${1}"; shift
    local file="${1}"; shift
    local columns

    columns=$(get_sent_columns "${table}")
    echo "WARNING: the ${table} rows not sent to the XiVOCC yet could not be restored."
    echo "         They are kept in ${file}. Restore them before starting the services with:"
    echo "           psql -U postgres -d asterisk -c \"COPY ${table}(${columns}) FROM STDIN;\" < ${file}"
}

restore_unsent_stats() {
    local table columns file

    for table in ${UNSENT_STATS_TABLES}; do
        file=$(get_unsent_stats_file "${table}")
        [ -f "${file}" ] || continue

        columns=$(get_sent_columns "${table}")
        if psql -U postgres -d asterisk -q -v ON_ERROR_STOP=1 -c "COPY ${table}(${columns}) FROM STDIN;" < "${file}"; then
            echo "Restored $(wc -l < "${file}") ${table} rows: they will be sent to the XiVOCC once the services are started."
            rm -f "${file}"
        else
            display_unsent_stats_restore_failure "${table}" "${file}"
        fi
    done
}

wipe_pgdata() {
    local pgdata="${1}"; shift

    echo "Stopping the database container..."
    xivo-dcomp stop db
    xivo-dcomp rm -f db > /dev/null

    echo "Removing the PostgreSQL data directory ${pgdata}..."
    rm -rf "${pgdata:?}"

    if [ -e "${pgdata}" ]; then
        echo "ERROR: ${pgdata} could not be removed." >&2
        return 1
    fi
}

wait_for_rebuilt_db() {
    local waited

    echo -n "Waiting up to ${REBUILT_DB_WAIT_TIMEOUT}s for the rebuilt database..."
    for ((waited = 0; waited < REBUILT_DB_WAIT_TIMEOUT; waited++)); do
        if is_postgres_running_on_port; then
            echo
            return 0
        fi
        sleep 1
        echo -n "."
    done
    echo
    echo "ERROR: the rebuilt database did not answer within ${REBUILT_DB_WAIT_TIMEOUT}s." >&2
    return 1
}

recreate_db() {
    echo "Recreating the database at the new schema..."
    if ! xivo-dcomp start-db; then
        echo "ERROR: the database container could not be started." >&2
        return 1
    fi
    wait_for_rebuilt_db
}

write_mds_resync_state() {
    local wipe_epoch="${1}"; shift

	cat > "${MDS_RESYNC_STATE_FILE}" <<-EOF
	MDS_DB_RECREATED=1
	MDS_DB_WIPE_EPOCH=${wipe_epoch}
	EOF
}

is_database_erased() {
    [ -f "${MDS_RESYNC_STATE_FILE}" ]
}

ask_to_continue_resync() {
    local answer

    echo ""
    echo "The PostgreSQL data directory of this media server is about to be ERASED."
    echo "Its database will be rebuilt from the XiVO Main; the local call_log table"
    echo "will be lost. The cel and queue_log rows not sent to the XiVOCC yet are kept."
    read -r -p 'Would you like to continue [N/y]? ' answer
    answer="${answer:-N}"
    if [ "$answer" != 'y' ] && [ "$answer" != 'Y' ]; then
        exit 0
    fi
}

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

    echo "Stopping the media server services..."
    xivo-service stop
    if ! save_unsent_stats; then
        [ "${force}" -eq 1 ] || return 1
        echo "WARNING: going on anyway, the cel and queue_log rows not sent to the XiVOCC yet are lost."
    fi
    xivo-service stop all
}

start_services() {
    echo "Starting the media server services..."
    xivo-dcomp up -d --remove-orphans
    wait_for_initial_copy
    xivo-service start
    catch_up_initial_copy
}

db_is_subscribed_to_main() {
    local subscriptions

    subscriptions=$(psql -U postgres -d asterisk -Aqt -c "SELECT COUNT(*) FROM pg_subscription;" 2> /dev/null)

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

display_subscription_enable_failure() {
	cat <<-EOF
	******************************************************************************
	*  ERROR: the subscription to the XiVO Main could not be enabled.            *
	*                                                                            *
	*  The local replica is EMPTY and will stay empty: this media server must    *
	*  not be put back in service as is.                                         *
	*                                                                            *
	*  You MUST:                                                                 *
	*    1. enable it once the database answers:                                 *
	*         xivo-enable-mds-subscription                                       *
	*    2. or re-run mds-resync-db                                              *
	******************************************************************************
	EOF
}

display_subscription_creation_failure() {
	cat <<-EOF
	******************************************************************************
	*  ERROR: the database was rebuilt but no subscription to the XiVO Main      *
	*  was created.                                                              *
	*                                                                            *
	*  The local replica is EMPTY and will stay empty: this media server must    *
	*  not be put back in service as is. The usual cause is a replication slot   *
	*  left behind on the Main, since CREATE SUBSCRIPTION also creates a slot    *
	*  of the same name and fails when it already exists.                        *
	*                                                                            *
	*  You MUST:                                                                 *
	*    1. on the XiVO Main, list the surviving slots:                          *
	*         SELECT slot_name FROM pg_replication_slots;                        *
	*       and drop the one named after this media server:                      *
	*         SELECT pg_drop_replication_slot('main_<MDS_NAME>');                *
	*    2. and then, re-run mds-resync-db                                       *
	******************************************************************************
	EOF
}

is_db_image_present() {
    local image="${1}"; shift

    docker image inspect "${image}" > /dev/null 2>&1
}

resync_db() {
    local image="${1}"; shift
    local pgdata cluster_locale wipe_epoch

    pgdata=$(get_pgdata_dir_from_image "${image}")
    if [ -z "${pgdata}" ]; then
        echo "ERROR: could not read PGDATA from ${image}." >&2
        return 1
    fi
    if [ ! -d "${pgdata}" ]; then
        echo "ERROR: ${pgdata} is not a directory of this host." >&2
        echo "       The data directory of the db container is expected to be bind mounted on the same path." >&2
        return 1
    fi

    cluster_locale="${MDS_CLUSTER_LOCALE:-$(get_db_locale)}"
    warn_about_locale_change "${image}" "${cluster_locale}"

    backup_pg_hba "${pgdata}"

    display_upgrade_notice "Rebuilding the media server database..."
    wipe_epoch=$(date +%s)
    if ! wipe_pgdata "${pgdata}"; then
        return 1
    fi
    write_mds_resync_state "${wipe_epoch}"

    if ! recreate_db; then
        echo "Check the db container with: xivo-dcomp logs db" >&2
        return 1
    fi
    restore_pg_hba "${pgdata}"

    if ! db_is_subscribed_to_main; then
        display_subscription_creation_failure
        return 1
    fi

    if ! xivo-enable-mds-subscription; then
        display_subscription_enable_failure
        return 1
    fi

    restore_unsent_stats

    echo "The database was recreated and re-subscribed to the XiVO Main."
    echo "The initial copy runs in the background."
}

usage() {
    local progname="${1}"; shift

    echo "usage: ${progname} [--no-service-control] [--check-order-only] [--save-unsent-stats] [-f] [-h]"
    echo ""
    echo "Erase the PostgreSQL data directory of this media server and let the db"
    echo "container rebuild it from the XiVO Main, subscription included."
    echo "The db image already present on this media server is used, none is pulled."
    echo ""
    echo -e "\t--no-service-control: do not stop/start the services (they already are stopped)"
    echo -e "\t--check-order-only: only check that the XiVO Main was upgraded first, change nothing"
    echo -e "\t--save-unsent-stats: only save the cel and queue_log rows not sent to the XiVOCC yet"
    echo -e "\t-f: skip the upgrade order check and the confirmation, go on if the unsent rows cannot be saved"
    echo -e "\t-h: print usage"
}

main() {
    local progname="mds-resync-db"
    local service_control=1
    local check_order_only=0
    local force=0
    local image

    while [ $# -gt 0 ]; do
        case "${1}" in
            --no-service-control) service_control=0;;
            --check-order-only) check_order_only=1;;
            --save-unsent-stats)
                save_unsent_stats
                exit
            ;;
            -f) force=1;;
            -h)
                usage "${progname}"
                exit 0
            ;;
            *)
                echo "${progname} : option ${1} is not valid" >&2
                usage "${progname}"
                exit 1
            ;;
        esac
        shift
    done

    if ! image=$(get_db_image); then
        echo "ERROR: could not resolve the database image from the docker compose configuration." >&2
        exit 1
    fi

    if ! is_db_image_present "${image}"; then
        echo "ERROR: ${image} is not present on this media server, pull it with: xivo-dcomp pull db" >&2
        exit 1
    fi

    if [ "${force}" -eq 0 ]; then
        display_upgrade_notice "Checking that the XiVO Main was upgraded first..."
        if ! check_upgrade_order "${image}"; then
            exit 1
        fi
    fi

    if [ "${check_order_only}" -eq 1 ]; then
        exit 0
    fi

    if ! command -v xivo-enable-mds-subscription > /dev/null; then
        echo "ERROR: xivo-enable-mds-subscription is missing, install it with: apt install xivo-manage-db-mds" >&2
        exit 1
    fi

    rm -f "${MDS_RESYNC_STATE_FILE}"

    if [ "${service_control}" -eq 1 ]; then
        [ "${force}" -eq 0 ] && ask_to_continue_resync
        if ! stop_services "${force}"; then
            echo "Resynchronization aborted before erasing anything, starting the services back." >&2
            start_services
            exit 1
        fi
    fi

    if ! resync_db "${image}"; then
        if is_database_erased; then
            echo "The database was erased: the services were left stopped." >&2
        elif [ "${service_control}" -eq 1 ]; then
            start_services
        fi
        exit 1
    fi

    if [ "${service_control}" -eq 1 ]; then
        import_mds_resync_state
        start_services
    fi
}

main "${@}"
