#!/bin/bash

declare -r SCRIPT_NAME="$(basename "$0")"
declare -r SCRIPT_PATH="$(readlink -f "$0")"
declare -r SCRIPT_DIR="$(dirname "$SCRIPT_PATH")"

declare -r VERBOSE_ERROR="1"
declare -r VERBOSE_WARNING="2"
declare -r VERBOSE_NOTICE="3"
declare -r VERBOSE_INFO="4"

declare -r VERBOSE_MAX="${VERBOSE_INFO}"
declare -r VERBOSE_MIN="${VERBOSE_ERROR}"

declare -r YES_PARAM="yes"
declare -r NO_PARAM="no"

declare -r EXIT_SUCCESS="0"
declare -r EXIT_FAILURE="1"

declare -r QUERY_CMD="mlxfwmanager --query"
declare -r BURN_CMD="mlxfwmanager -u -f -y"

declare -r FW_FILE="/etc/mlnx/fw-SPC.mfa"
declare -r QUERY_FILE="/tmp/mlxfwmanager-query.txt"

declare -r FW_REQUIRED="{{ MLNX_FW_VERSION }}"

IMAGE_UPGRADE="${NO_PARAM}"
CPLD_UPGRADE="${NO_PARAM}"
VERBOSE_LEVEL="${VERBOSE_MIN}"

function PrintHelp() {
    echo
    echo "Usage: ./${SCRIPT_NAME} [OPTIONS]"
    echo
    echo "OPTIONS:"
    echo "  -u, --upgrade  Upgrade firmware using next boot image (useful after SONiC-To-SONiC update)"
    echo "  -c, --cpld     Upgrade CPLD firmware (requires -u|--upgrade)"
    echo "  -v, --verbose  Verbose mode"
    echo "  -h, --help     Print help"
    echo
    echo "Examples:"
    echo "  ./${SCRIPT_NAME}"
    echo "  ./${SCRIPT_NAME} --upgrade --cpld --verbose"
    echo "  ./${SCRIPT_NAME} --upgrade --verbose"
    echo "  ./${SCRIPT_NAME} --help"
    echo
}

function ParseArguments() {
    while [[ $# -ge 1 ]]; do
        case "$1" in
            -u|--upgrade)
                IMAGE_UPGRADE="${YES_PARAM}"
                VERBOSE_LEVEL="${VERBOSE_MAX}"
            ;;
            -c|--cpld)
                CPLD_UPGRADE="${YES_PARAM}"
            ;;
            -v|--verbose)
                VERBOSE_LEVEL="${VERBOSE_MAX}"
            ;;
            -h|--help)
                PrintHelp
                exit "${EXIT_SUCCESS}"
            ;;
        esac
        shift
    done
}

function LogError() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_ERROR}" ]]; then
        echo "ERROR: $*"
    fi
}

function LogWarning() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_WARNING}" ]]; then
        echo "WARNING: $*"
    fi
}

function LogNotice() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_NOTICE}" ]]; then
        echo "NOTICE: $*"
    fi
}

function LogInfo() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_INFO}" ]]; then
        echo "INFO: $*"
    fi
}

function ExitFailure() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_ERROR}" ]]; then
        echo
        LogError "$@"
        echo
    fi

    exit "${EXIT_FAILURE}"
}

function ExitSuccess() {
    if [[ "${VERBOSE_LEVEL}" -ge "${VERBOSE_INFO}" ]]; then
        echo
        LogInfo "$@"
        echo
    fi

    exit "${EXIT_SUCCESS}"
}

function ParseMachineConf() {
    ONIE_MACHINE="$(cat /host/machine.conf | grep 'onie_machine=' | cut -f2 -d'=')"
    ONIE_PLATFORM="$(cat /host/machine.conf | grep 'onie_platform=' | cut -f2 -d'=')"
}

function ShowProgressBar() {
    local -rA SPIN=(
        [0]="-"
        [1]="\\"
        [2]="|"
        [3]="/"
    )

    if [[ "${VERBOSE_LEVEL}" -lt "${VERBOSE_INFO}" ]]; then
        sleep 2s
        return "${EXIT_SUCCESS}"
    fi

    # Print progress bar: use carriage return to overwrite command line content
    for i in "${SPIN[@]}"; do
        echo -ne "\r[${i}] ${1}"
        sleep 0.5s
    done

    # Clear command line content + carriage return
    echo -ne "\033[1K\r"
}

function WaitForDevice() {
    local -i QUERY_RETRY_COUNT_MAX="10"
    local -i QUERY_RETRY_COUNT="0"

    ${QUERY_CMD} &> /dev/null

    while [[ ("${QUERY_RETRY_COUNT}" -lt "${QUERY_RETRY_COUNT_MAX}") && ("$?" -ne "0") ]]; do
        sleep 1s
        ((QUERY_RETRY_COUNT++))
        ${QUERY_CMD} &> /dev/null
    done
}

function RunCmd() {
    local _EXIT_CODE="${EXIT_SUCCESS}"

    if [[ "${VERBOSE_LEVEL}" -eq "${VERBOSE_MAX}" ]]; then
        eval "$@"
    else
        eval "$@" &>/dev/null
    fi

    _EXIT_CODE="$?"
    if [[ "${_EXIT_CODE}" != "${EXIT_SUCCESS}" ]]; then
        ExitFailure "command failed: $@"
    fi
}

function UpgradeASICFW() {
    local _FW_FILE="$1"

    if [[ ! -z "${_FW_FILE}" ]]; then
        if [[ ! -f "${_FW_FILE}" ]]; then
            ExitFailure "no such file: ${_FW_FILE}"
        fi

        RunCmd "${QUERY_CMD} -i ${_FW_FILE} > ${QUERY_FILE}"

        local -r _FW_INFO="$(grep FW ${QUERY_FILE})"
        local -r _FW_CURRENT="$(echo ${_FW_INFO} | cut -f2 -d' ')"
        local -r _FW_AVAILABLE="$(echo ${_FW_INFO} | cut -f3 -d' ')"
    else
        RunCmd "${QUERY_CMD} > ${QUERY_FILE}"

        local -r _FW_INFO="$(grep FW ${QUERY_FILE})"
        local -r _FW_CURRENT="$(echo ${_FW_INFO} | cut -f2 -d' ')"
        local -r _FW_AVAILABLE="${FW_REQUIRED}"

        _FW_FILE="${FW_FILE}"
    fi

    if [[ -z "${_FW_CURRENT}" ]]; then
        ExitFailure "could not retrieve current ASIC firmware version"
    fi

    if [[ -z "${_FW_AVAILABLE}" ]]; then
        ExitFailure "could not retrieve available ASIC firmware version"
    fi

    if [[ "${_FW_CURRENT}" = "${_FW_AVAILABLE}" ]]; then
        LogInfo "ASIC firmware is up to date"
        return "${EXIT_SUCCESS}"
    fi

    LogNotice "ASIC firmware upgrade is required. Installing compatible version..."

    LogInfo "current ASIC firmware version: ${_FW_CURRENT}"
    LogInfo "target ASIC firmware version: ${_FW_AVAILABLE}"

    LogInfo "ASIC firmware file: ${_FW_FILE}"

    RunCmd "${BURN_CMD} -i ${_FW_FILE}"
}

function UpgradeCPLDFW_Worker() {
    local -r _CPLD_BURN_FILE="${1}"
    local -r _CPLD_REFRESH_FILE="${2}"
    local -r _ASIC_DEV="$(find /dev/mst -iname '*_pciconf0')"

    if [[ -f /tmp/cpld_fw_updated ]]; then
        RunCmd "cpldupdate --dev ${_ASIC_DEV} ${_CPLD_REFRESH_FILE}"
        return "${EXIT_SUCCESS}"
    fi

    RunCmd "cpldupdate --dev ${_ASIC_DEV} ${_CPLD_BURN_FILE}"
    RunCmd "cpldupdate --dev ${_ASIC_DEV} ${_CPLD_REFRESH_FILE}"
}

function UpgradeCPLDFW() {
    local -r _CPLD_ARCHIVE="$1"

    if [[ -f /tmp/cpld_fw_updated ]]; then
        LogWarning "forced CPLD refresh was requested for ${ONIE_PLATFORM}"
        CPLD_UPGRADE="${YES_PARAM}"
    fi

    if [[ "${CPLD_UPGRADE}" != "${YES_PARAM}" ]]; then
        LogNotice "CPLD upgrade was not requested for ${ONIE_PLATFORM}"
        return "${EXIT_SUCCESS}"
    fi

    if [[ ! -f "${_CPLD_ARCHIVE}" ]]; then
        LogNotice "CPLD update $(basename ${_CPLD_ARCHIVE}) was not provided for ${ONIE_PLATFORM}"
        return "${EXIT_SUCCESS}"
    fi

    CPLD_DIR="$(mktemp -d)"

    RunCmd "tar xzf ${_CPLD_ARCHIVE} -C ${CPLD_DIR}"

    local -r _CPLD_BURN_FILE="${CPLD_DIR}/$(cat ${CPLD_DIR}/bundle.txt | grep 'burn=' | cut -f2 -d'=')"
    local -r _CPLD_REFRESH_FILE="${CPLD_DIR}/$(cat ${CPLD_DIR}/bundle.txt | grep 'refresh=' | cut -f2 -d'=')"
    local -r _CPLD_VERSION="$(cat ${CPLD_DIR}/bundle.txt | grep 'version=' | cut -f2 -d'=')"

    local _CURRENT_CPLD_VERSION="${_CPLD_VERSION}"
    local _TARGET_CPLD_VERSION="${_CPLD_VERSION}"

    if [[ -f /bsp/cpld/cpld_mgmt_version ]]; then
        _CURRENT_CPLD_VERSION="$(cat /bsp/cpld/cpld_mgmt_version)"
    elif [[ -f /var/run/hw-management/system/cpld1_version ]]; then
        _CURRENT_CPLD_VERSION="$(cat /var/run/hw-management/system/cpld1_version)"
    else
        ExitFailure "could not retrieve current CPLD firmware version"
    fi

    if [[ "${_CURRENT_CPLD_VERSION}" = "${_TARGET_CPLD_VERSION}" ]]; then
        LogInfo "CPLD firmware is up to date"
        return "${EXIT_SUCCESS}"
    fi

    LogNotice "CPLD firmware upgrade is required. Installing compatible version..."

    LogInfo "current CPLD firmware version: ${_CURRENT_CPLD_VERSION}"
    LogInfo "target CPLD firmware version: ${_TARGET_CPLD_VERSION}"

    LogInfo "CPLD burn firmware file: ${_CPLD_BURN_FILE}"
    LogInfo "CPLD refresh firmware file: ${_CPLD_REFRESH_FILE}"

    UpgradeCPLDFW_Worker "${_CPLD_BURN_FILE}" "${_CPLD_REFRESH_FILE}" &
    local -r _PID="$!"

    while $(ps -e -o pid | grep -E "^[[:blank:]]*${_PID}$" &> /dev/null); do
        ShowProgressBar "CPLD update..."
    done

    RunCmd "wait ${_PID}"
}

function UpgradeFWFromImage() {
    local -r _NEXT_SONIC_IMAGE="$(sonic_installer list | grep 'Next: ' | cut -f2 -d' ')"
    local -r _CURRENT_SONIC_IMAGE="$(sonic_installer list | grep 'Current: ' | cut -f2 -d' ')"

    if [[ "${_CURRENT_SONIC_IMAGE}" = "${_NEXT_SONIC_IMAGE}" ]]; then
        ExitSuccess "firmware is up to date"
    fi

    FS_PATH="/host/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}/fs.squashfs"
    FS_MOUNTPOINT="/tmp/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}-fs"

    RunCmd "mkdir -p ${FS_MOUNTPOINT}"
    RunCmd "mount -t squashfs ${FS_PATH} ${FS_MOUNTPOINT}"

    UpgradeASICFW "${FS_MOUNTPOINT}${FW_FILE}"
    UpgradeCPLDFW "${FS_MOUNTPOINT}/etc/mlnx/cpld/${ONIE_MACHINE#mlnx_}_cpld.tar.gz"
}

function Cleanup() {
    if [[ -d "${FS_MOUNTPOINT}" ]]; then
        umount -rf "${FS_MOUNTPOINT}"
        rm -rf "${FS_MOUNTPOINT}"
    fi

    if [[ -d "${CPLD_DIR}" ]]; then
        rm -rf "${CPLD_DIR}"
    fi
}

trap Cleanup EXIT

ParseMachineConf
ParseArguments "$@"

WaitForDevice

if [[ "${CPLD_UPGRADE}" = "${YES_PARAM}" ]]; then
    if [[ "${IMAGE_UPGRADE}" = "${NO_PARAM}" ]]; then
        ExitFailure "mandatory parameter was not provided: -u|--upgrade"
    fi
fi

if [[ "${IMAGE_UPGRADE}" = "${YES_PARAM}" ]]; then
    UpgradeFWFromImage
else
    UpgradeASICFW
fi

ExitSuccess "firmware upgrade is completed"