#!/bin/bash

COMPOSE_PATH="/etc/docker/mds"
FACTORY_ENV_FILE="${COMPOSE_PATH}/factory.env"
CUSTOM_ENV_FILE="${COMPOSE_PATH}/custom.env"
MAIN_DB_USER="asterisk"
PG_HBA_BACKUP_DIR="/var/tmp/mds-resync-db"
PG_HBA_BACKUP="${PG_HBA_BACKUP_DIR}/pg_hba.conf"
MDS_RESYNC_STATE_FILE="/run/mds-resync-db.state"

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

get_env_value() {
    local key="${1}"; shift
    local value=""

    if [ -f "${CUSTOM_ENV_FILE}" ]; then
        value=$(grep -oP -m 1 "^${key}=\K.*" "${CUSTOM_ENV_FILE}")
    fi
    if [ -z "${value}" ] && [ -f "${FACTORY_ENV_FILE}" ]; then
        value=$(grep -oP -m 1 "^${key}=\K.*" "${FACTORY_ENV_FILE}")
    fi

    echo "${value}"
}

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

    psql -h "${main_host}" -p 5432 -U "${MAIN_DB_USER}" -d asterisk -Aqt -w \
        -c "SELECT version_num FROM alembic_version;" 2> /dev/null | head -n 1
}

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() {
    local main_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 connection uses /root/.pgpass: check that it holds an entry for this  *
	*  Main as the asterisk user.                                                *
	******************************************************************************
	EOF
    echo "  XiVO Main host: ${main_host}"
}

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

    main_host=$(get_env_value "XIVO_HOST")
    if [ -z "${main_host}" ]; then
        echo "ERROR: XIVO_HOST is not set in ${CUSTOM_ENV_FILE} nor ${FACTORY_ENV_FILE}." >&2
        return 1
    fi

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

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

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
}

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 cel, queue_log"
    echo "and call_log tables will be lost."
    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() {
    echo "Stopping the media server services..."
    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_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
}

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

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

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

    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

    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] [-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-f: skip the upgrade order check and the confirmation"
    echo -e "\t-h: print usage"
}

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

    while [ $# -gt 0 ]; do
        case "${1}" in
            --no-service-control) service_control=0;;
            -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

    rm -f "${MDS_RESYNC_STATE_FILE}"

    if [ "${service_control}" -eq 1 ]; then
        [ "${force}" -eq 0 ] && ask_to_continue_resync
        stop_services
    fi

    if ! resync_db "${image}" "${force}"; then
        [ "${service_control}" -eq 1 ] && start_services
        exit 1
    fi

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

main "${@}"
