#!/bin/bash

BACKUP_DIR="/var/tmp/mds-resync-db"
PG_HBA_BACKUP="${BACKUP_DIR}/pg_hba.conf"
MDS_RESYNC_STATE_FILE="/run/mds-resync-db.state"
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

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 "${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"
    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
}

recreate_db() {
    echo "Recreating the database at the new schema..."
    xivo-dcomp upgrade-db
    wait_for_postgres
}

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
    xivo-service start
}

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
}

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

    display_upgrade_notice "Pulling the database image..."
    if ! xivo-dcomp pull db; then
        echo "ERROR: could not pull ${image}." >&2
        return 1
    fi
}

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}"

    recreate_db
    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 ""
    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 ! pull_db_image "${image}"; then
        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

    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
        start_services
    fi
}

main "${@}"
