#!/bin/bash

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

declare -r EXIT_SUCCESS="0"
declare -r EXIT_ERROR="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/mlnxfwmanager-query.txt"

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

IMAGE_UPGRADE="no"

function PrintHelp() {
    echo
    echo "Usage: ./${SCRIPT_NAME} [OPTIONS]"
    echo
    echo "OPTIONS:"
    echo "  -u, --upgrade  Upgrade MLNX ASIC firmware using next boot image (useful after SONiC-To-SONiC update)"
    echo "  -h, --help     Print help"
    echo
    echo "Examples:"
    echo "  ./${SCRIPT_NAME}"
    echo "  ./${SCRIPT_NAME} --upgrade"
    echo "  ./${SCRIPT_NAME} --help"
    echo
}

function ParseArguments() {
    while [ $# -ge 1 ]; do
        case "$1" in
            -u|--upgrade)
                IMAGE_UPGRADE="yes"
            ;;
            -h|--help)
                PrintHelp
                exit "${EXIT_SUCCESS}"
            ;;
        esac
        shift
    done
}

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() {
    $1
    if [[ $? != 0 ]]; then
        echo "Command failed: cmd=$1, errno=$?"
        exit "${EXIT_ERROR}"
    fi
}

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

    if [ ! -z "${_FW_FILE}" ]; then
        if [ ! -f "${_FW_FILE}" ]; then
            echo "No such file: ${_FW_FILE}"
            exit "${EXIT_ERROR}"
        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
        echo "Could not retreive current FW version"
        exit "${EXIT_ERROR}"
    fi

    if [[ -z "${_FW_AVAILABLE}" ]]; then
        echo "Could not retreive available FW version"
        exit "${EXIT_ERROR}"
    fi

    if [[ "${_FW_CURRENT}" == "${_FW_AVAILABLE}" ]]; then
        echo "Mellanox firmware is up to date"
    else
        echo "Mellanox firmware upgrade is required. Installing compatible version..."
        RunCmd "${BURN_CMD} -i ${_FW_FILE}"
    fi
}

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' ')"

    local -r _FS_PATH="/host/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}/fs.squashfs"
    local -r _FS_MOUNTPOINT="/tmp/image-${_NEXT_SONIC_IMAGE#SONiC-OS-}-fs"

    if [[ "${_CURRENT_SONIC_IMAGE}" == "${_NEXT_SONIC_IMAGE}" ]]; then
        echo "Mellanox firmware is up to date"
    else
        mkdir -p "${_FS_MOUNTPOINT}"
        mount -t squashfs "${_FS_PATH}" "${_FS_MOUNTPOINT}"

        UpgradeFW "${_FS_MOUNTPOINT}/etc/mlnx/fw-SPC.mfa"

        umount -rf "${_FS_MOUNTPOINT}"
        rm -rf "${_FS_MOUNTPOINT}"
    fi
}

ParseArguments "$@"

WaitForDevice

if [ "${IMAGE_UPGRADE}" != "yes" ]; then
    UpgradeFW
else
    UpgradeFWFromImage
fi

exit "${EXIT_SUCCESS}"