#!/bin/bash

#
# Codex Vault AVR update script
#

if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root" 1>&2
   exit 1
fi

AVRFirmwareLocation="/usr/share/codex/firmware"

#-------------------------------------------------
# Command line options
#-------------------------------------------------

stopProcesses=1
updateModules=1
updateMain=1
updateBay=0
updatePeripherals=1
startProcesses=1

while [[ ${#} -gt 0 ]]; do

    case "${1}" in

        "--skipProcessStop" )
            stopProcesses=0
            shift 1
            ;;

        "--mainAVROnly" )
            updateModules=0
            updatePeripherals=0
            shift 1
            ;;

        "--modulesOnly" )
            updateMain=0
            updatePeripherals=0
            shift 1
            ;;

        "--peripheralsOnly" )
            updateMain=0
            updateModules=0
            shift 1
            ;;

        "--skipProcessStart" )
            startProcesses=0
            shift 1
            ;;

        * )
            echo "Unrecognised argument \"${1}\""
            exit 1
            ;;
    esac
done

#-------------------------------------------------
# Hardware identification
#-------------------------------------------------

if "/usr/bin/isVaultXL" ; then

    # This is a Vault XL
    mainAVRFirmwarePrefix="vaultxl_v"
    updateModules=0 #Vault XLs don't have any modules
    updateBay=1 # Might be a workstation with Bay CD2 fitted
    bayAVRFirmwarePrefix="baycd2_v"

elif "/usr/bin/isVaultS" ; then

    # This is a Vault S
    mainAVRFirmwarePrefix="vault_v"

else

    # Is this even a Vault at all?
    echo "Unable to identify Vault type - This does not appear to be a Codex Vault device" >&2
    exit 1

fi

#-------------------------------------------------
# Functions
#-------------------------------------------------

wait_for_exit() {
    PROCESS=$1
    WAIT=30

    PID=$(pidof $PROCESS)
    if [[ -n "$PID" ]]; then
        echo -n "Waiting for $PROCESS to exit"
        i=0
        while true; do
            echo -n "."
            sleep 1
            PID=$(pidof $PROCESS)
            if [[ -z "$PID" ]]; then
                #exited
                break
            fi
            if ((i > $WAIT)); then
                echo "Failed, force killing it" 1>&2
                kill -KILL $PID
                break
            fi
            i=$((i+1))
        done
        echo ""
    fi
}

#---------------------------------------

unload_kernel_module() {
    local moduleName=${1}
    local moduleDetails="$(lsmod | grep "^${moduleName} ")"

    if [ -n "${moduleDetails}" ] ; then

        local dependencyList=( $(echo "${moduleDetails}" | sed "s/^[^ \t]*[ \t]*[^ \t]*[ \t]*[^ \t]*[ \t]*//" | sed "s/,/ /g") )

        #-------------------
        #
        # Recursively unload the dependencies
        #

        if [ ${#dependencyList[@]} -gt 0 ] ; then

            local dependentModule
            for dependentModule in ${dependencyList[@]} ; do

                unload_kernel_module ${dependentModule}
                local unloadResult=${?}

                #
                # Bail out if we hit an error - this dependency can't be unloaded
                #
                if [ ${unloadResult} -ne 0 ] ; then
                    return ${unloadResult}
                fi
            done
        fi

        #-------------------
        #
        # Unload the module itself
        #

        echo "Unloading ${moduleName} kernel module"
        rmmod ${moduleName} > /dev/null 2>&1
        return ${?}

        #-------------------
    fi

    return 0
}

#-------------------------------------------------
# Shutting down anything that may interfere
#-------------------------------------------------

if [ ${stopProcesses} -ne 0 ] ; then

    echo
    echo "--------- Shutting down running processes for firmware update ----------"
    echo

    systemctl stop drserver
    wait_for_exit drserver

fi

if [ ${updateModules} -ne 0 ] ; then

    # make sure the Nvidia stuff is unloaded from
    # the kernel so we don't brick review modules

    unload_kernel_module nvidia

    if [ ${?} -ne 0 ] ; then

        echo "Failed to unload kernel modules - update is unsafe - aborting"
        updateModules=0
        updateMain=0

    fi
fi

#-------------------------------------------------
# Actual AVR updates
#-------------------------------------------------

if [ ${updateModules} -ne 0 ] ; then

    echo
    echo "--------------------- Updating module firmware -------------------------"
    echo

    # Upgrade any Vault modules
    sleep 2 # Protective wait against the Vault board being unready when the update happens
    AVRModUpgrade -a ${AVRFirmwareLocation}/storage_v*.fw
    sleep 2
    AVRModUpgrade -a ${AVRFirmwareLocation}/lto_v*.fw
    sleep 2
    AVRModUpgrade -a ${AVRFirmwareLocation}/pcie_v*.fw
fi

if [ ${updatePeripherals} -ne 0 ] ; then
    echo
    echo "------------------- Updating peripheral firmware -----------------------"
    echo

    # Upgrade any Codex peripherals
    sleep 2
    AVRUpgrade -a ${AVRFirmwareLocation}/xferdual_v*.fw
fi

if [ ${updateMain} -ne 0 ] ; then
    echo
    echo "------------------------ Updating main AVR -----------------------------"
    echo
    echo "NOTE: Vault S will power down automatically if firmware upgrade is required"
    echo

    sleep 2
    AVRUpgrade -a ${AVRFirmwareLocation}/${mainAVRFirmwarePrefix}*.fw
    if [ ${updateBay} -ne 0 ] ; then
        AVRUpgrade -a ${AVRFirmwareLocation}/${bayAVRFirmwarePrefix}*.fw
    fi
fi

#-------------------------------------------------
# Resume normal operations
#-------------------------------------------------

if [ ${startProcesses} -ne 0 ] ; then

    echo
    echo "----------------------- Restoring processes ----------------------------"
    echo

    systemctl start default.target

    echo
fi

#-------------------------------------------------
#
#-------------------------------------------------
