NIC_FW_UPDATE_DONE=0
FORCE_NIC_FW_UPDATE=${FORCE_NIC_FW_UPDATE:-"no"}
NIC_FW_RESET=${NIC_FW_RESET:-"yes"}
FORCE_NIC_FW_RESET=${FORCE_NIC_FW_RESET:-"no"}
NIC_FW_RESET_REQUIRED=0
NIC_FW_FOUND=0
FW_UPDATER=${FW_UPDATER:-"/opt/mellanox/mlnx-fw-updater/mlnx_fw_updater.pl"}
FW_DIR=${FW_DIR:-"/opt/mellanox/mlnx-fw-updater/firmware/"}
is_nic_mode=${is_nic_mode:-0}
cx_pcidev=$(lspci -nD 2> /dev/null | grep 15b3:a2d[26c] | awk '{print $1}' | head -1)
flint_dev=$cx_pcidev
FLINT=mstflint
if [ ! -x /usr/bin/mstflint ]; then
	FLINT=flint
fi

PROVIDED_NIC_FW_VERSION=""
CHROOT=
CHROOT_DIR=${CHROOT_DIR:-"/mnt"}

is_SecureBoot=0
if (mokutil --sb-state 2>&1 | grep -q "SecureBoot enabled"); then
        is_SecureBoot=1
fi

if [ $is_SecureBoot -eq 1 ]; then
        mst_dev=$(/bin/ls -1 /dev/mst/mt*_pciconf0 2> /dev/null)
        if [ ! -n "${mst_dev}" ]; then
                mst start > /dev/null 2>&1
        fi
		flint_dev=$(/bin/ls -1 /dev/mst/mt*_pciconf0 2> /dev/null)
		FLINT=flint
fi

if [ "$FORCE_NIC_FW_RESET" == "yes" ]; then
	NIC_FW_RESET="yes"
fi

running_nic_fw()
{
	$FLINT -d $flint_dev q 2>&1 | grep -w 'FW Version:' | awk '{print $NF}'
}

provided_nic_fw()
{
	if [ ! -z "${PSID}" ]; then
		if [ -n "$NIC_FW_BIN" ]; then
			PROVIDED_NIC_FW_VERSION=$($FLINT -i $NIC_FW_BIN q 2>&1 | grep -w 'FW Version:' | awk '{print $NF}')
		else
			PROVIDED_NIC_FW_VERSION=$(${CHROOT_DIR}${FW_DIR}/mlxfwmanager_sriov_dis_aarch64_${cx_dev_id} --list 2> /dev/null | grep -w "${PSID}" | awk '{print $4}')
		fi
		echo $PROVIDED_NIC_FW_VERSION
	fi
}

is_nic_fw_update_required()
{
	if [ "${FORCE_NIC_FW_UPDATE}" == "yes" ]; then
		log "FORCE_NIC_FW_UPDATE is set."
		return 0
	fi

	if [ "$(running_nic_fw)" == "$(provided_nic_fw)" ]; then
		log "Installed NIC Firmware is the same as provided. Skipping NIC Firmware update."
		return 1
	fi
	return 0
}

fw_update()
{
	if [[ -x ${CHROOT_DIR}/${FW_UPDATER} && -d ${CHROOT_DIR}/${FW_DIR} ]]; then
		NIC_FW_FOUND=1
	elif [[ -x ${FW_UPDATER} && -d ${FW_DIR} ]]; then
		NIC_FW_FOUND=1
		CHROOT_DIR=
	else
		NIC_FW_FOUND=0
	fi

	if [ $NIC_FW_FOUND -eq 1 ]; then
		# Check if firmware image was already updated on flash
		if ($FLINT -d $cx_pcidev q 2>&1 | grep -q 'FW Version(Running)'); then
			log "Reactivating previous firmware image on the NIC"
			ilog "$($FLINT -d $cx_pcidev ir)"
		fi

		if ! is_nic_fw_update_required; then
			return 0
		fi

		log "Updating NIC firmware..."
		${CHROOT_DIR}${FW_UPDATER} --log /tmp/mlnx_fw_update.log -v \
			--force-fw-update \
			--fw-dir ${CHROOT_DIR}${FW_DIR} > /tmp/mlnx_fw_update.out 2>&1
		rc=$?
		sync
		if [ -e /tmp/mlnx_fw_update.out ]; then
			cat /tmp/mlnx_fw_update.out > /dev/hvc0
			cat /tmp/mlnx_fw_update.out > /dev/ttyAMA0
			cat /tmp/mlnx_fw_update.out >> $LOG
		fi
		if [ -e /tmp/mlnx_fw_update.log ]; then
			cat /tmp/mlnx_fw_update.log > /dev/hvc0
			cat /tmp/mlnx_fw_update.log > /dev/ttyAMA0
			cat /tmp/mlnx_fw_update.log >> $LOG
		fi
		NIC_FW_UPDATE_DONE=1
		if [ $rc -ne 0 ] || (grep -q '\-E- Failed' /tmp/mlnx_fw_update.log); then
			NIC_FW_UPDATE_PASSED=0
			log "ERR NIC firmware update failed"
			return 1
		else
			NIC_FW_UPDATE_PASSED=1
			log "NIC firmware update done: $(provided_nic_fw)"
			return 0
		fi
	else
		log "ERR NIC Firmware files were not found"
		return 1
	fi
}

fw_bin_update()
{
	local rc=0
	if ! is_nic_fw_update_required; then
		return 0
	fi
	output=$($FLINT -d $flint_dev -i $NIC_FW_BIN -y b 2>&1)
	rc=$?
	if [ $rc -eq 0 ]; then
		NIC_FW_UPDATE_PASSED=1
		log "NIC firmware update done: $(provided_nic_fw). NIC Firmware reset or Host power cycle is required to activate the new NIC Firmware."
	else
		NIC_FW_UPDATE_PASSED=0
		log "ERR NIC firmware update failed: $output"
	fi
	return $rc
}

fw_reset()
{
	local rc=0
	if [ $is_nic_mode -eq 1 ]; then
		log "Run mlxfwreset or system-level reset to load new NIC firmware"
		return 1
	fi

	if ($CHROOT mlxfwreset -d /dev/mst/mt*_pciconf0 q | grep live-Patch | grep -qw "\-Supported"); then
		log "Live Patch NIC Firmware reset is supported."
		msg=$($CHROOT mlxfwreset -d /dev/mst/mt*_pciconf0 -y -l 0 r 2>&1)
		rc=$?
		if [ $rc -ne 0 ]; then
			log "ERR Live Patch NIC Firmware reset failed."
			log "$msg"
			if [ "$FORCE_NIC_FW_RESET" != "yes" ]; then
				return $rc
			fi
		else
			log "Live Patch NIC Firmware reset done"
			return 0
		fi
	else
		if [ "$LFWP" == "yes" ]; then
		        log "ERR Live Patch NIC Firmware reset is not supported."
		        if [ "$FORCE_NIC_FW_RESET" != "yes" ]; then
		                return 1
		        fi
		fi
	fi

	if [ "$RUNTIME_UPGRADE" == "yes" ]; then
		return 0
	fi

	MLXFWRESET_TIMEOUT=${MLXFWRESET_TIMEOUT:-60}
	SECONDS=0
	while ! ($CHROOT mlxfwreset -d /dev/mst/mt*_pciconf0 q 2>&1 | grep -w "Driver is the owner" | grep -qw "\-Supported")
	do
		if [ $SECONDS -gt $MLXFWRESET_TIMEOUT ]; then
			log "ERR NIC Firmware reset is not supported. Host power cycle is required"
			return 1
		fi
		sleep 1
	done

	if ! ($CHROOT mlxfwreset -d /dev/mst/mt*_pciconf0 q 2>&1 | grep -w "Driver restart and PCI reset" | grep -qw "\-Supported"); then
		log "ERR NIC Firmware reset is not supported. Host power cycle is required"
		return 1
	fi

	log "Running NIC Firmware reset"
	# Wait for these messages to be pulled by the rshim service
	# as mlxfwreset will restart the DPU
	sleep 3

	msg=$($CHROOT mlxfwreset -d /dev/mst/mt*_pciconf0 -y -l 3 --sync 1 r 2>&1)
	rc=$?
	if [ $rc -ne 0 ]; then
		log "ERR NIC Firmware reset failed. Host power cycle is required"
		log "$msg"
	else
		log "NIC Firmware reset done"
		if [ "$(running_nic_fw)" == "${PROVIDED_NIC_FW_VERSION}" ]; then
			log "NIC Firmware reset passed. Running NIC FW: ${PROVIDED_NIC_FW_VERSION}"
		else
			log "ERR NIC Firmware reset failed. Host power cycle is required"
			rc=1
		fi
	fi
	return $rc
}

update_nic_firmware()
{
	local rc=0
	if [ $NIC_FW_UPDATE_DONE -eq 0 ]; then
		if [ -n "$NIC_FW_BIN" ]; then
			fw_bin_update
		else
			fw_update
		fi
		rc=$?
	fi
	return $rc
}

nic_firmware_update()
{
	local update_nic_firmware_status
	update_nic_firmware
	update_nic_firmware_status=$?
	update_progress nic_firmware $update_nic_firmware_status
	if [ $update_nic_firmware_status -eq 0 ]; then
		touch ${BFB_DIR}/nic_firmware_updated
	else
		touch ${BFB_DIR}/nic_firmware_update_failed
	fi
	return $update_nic_firmware_status
}

reset_nic_firmware()
{
	local rc=0
	if [ "$NIC_FW_RESET" != "yes" ]; then
		log "Skip NIC Firmware reset"
		return 0
	fi

	if [ $COMPATIBILITY_FLOW -eq 1 ]; then
		if [ $NIC_FW_UPDATE_DONE -eq 1 ]; then
			if [ $NIC_FW_UPDATE_PASSED -eq 1 ]; then
				# Reset NIC FW
				fw_reset
				rc=$?
			fi
		fi
	elif [ $ACTIVATE_FLOW -eq 1 ]; then
		fw_reset
		rc=$?
	fi
	return $rc
}
