NIC_FW_UPDATE_DONE=0
FORCE_NIC_FW_UPDATE=${FORCE_NIC_FW_UPDATE:-"no"}
NIC_FW_RESET_REQUIRED=0
NIC_FW_FOUND=0
FW_UPDATER=/opt/mellanox/mlnx-fw-updater/mlnx_fw_updater.pl
FW_DIR=/opt/mellanox/mlnx-fw-updater/firmware/
is_nic_mode=${is_nic_mode:-0}

CHROOT=

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

provided_nic_fw()
{
	if [ ! -z "${PSID}" ]; then
		$CHROOT ${FW_DIR}/mlxfwmanager_sriov_dis_aarch64_${cx_dev_id} --list 2> /dev/null | grep -w "${PSID}" | awk '{print $4}'
	fi
}

fw_update()
{
	if [[ -x ${FW_UPDATER} && -d ${FW_DIR} ]]; then
		NIC_FW_FOUND=1
	else
		if [[ ! -d /mnt/root && ! -z "${root_device}" ]]; then
			mount -t $ROOTFS /dev/${root_device} /mnt
			bind_partitions
		fi

		if [[ -x /mnt/${FW_UPDATER} && -d /mnt/${FW_DIR} ]]; then
			CHROOT="chroot /mnt"
			NIC_FW_FOUND=1
		fi
	fi

	if [ $NIC_FW_FOUND -eq 1 ]; then
		if [ "$(running_nic_fw)" == "$(provided_nic_fw)" ]; then
			if [ "${FORCE_NIC_FW_UPDATE}" == "yes" ]; then
				log "INFO: Installed NIC Firmware is the same as provided. FORCE_NIC_FW_UPDATE is set."
			else
				log "INFO: Installed NIC Firmware is the same as provided. Skipping NIC Firmware update."
				return
			fi
		fi

		log "INFO: Updating NIC firmware..."
		$CHROOT ${FW_UPDATER} --log /tmp/mlnx_fw_update.log -v \
			--force-fw-update \
			--fw-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
		if [ $rc -ne 0 ] || (grep -q '\-E- Failed' /tmp/mlnx_fw_update.log); then
			NIC_FW_UPDATE_PASSED=0
			log "INFO: NIC firmware update failed"
		else
			NIC_FW_UPDATE_PASSED=1
			log "INFO: NIC firmware update done: $(provided_nic_fw)"
		fi
		NIC_FW_UPDATE_DONE=1
	else
		log "WARNING: NIC Firmware files were not found"
	fi
}

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

	/sbin/modprobe -a mlx5_core ib_umad
	run_mlnx_bf_configure=0
	if [ -x /mnt/sbin/mlnx_bf_configure ]; then
		CHROOT="chroot /mnt"
		run_mlnx_bf_configure=1
	elif [ -x /sbin/mlnx_bf_configure ]; then
		run_mlnx_bf_configure=1
	else
		run_mlnx_bf_configure=0
	fi

	if [ $run_mlnx_bf_configure -eq 1 ]; then
		ilog "Running mlnx_bf_configure:"
		ilog "$($CHROOT /sbin/mlnx_bf_configure 2>&1)"
	fi

	MLXFWRESET_TIMEOUT=${MLXFWRESET_TIMEOUT:-180}
	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 "INFO: NIC Firmware reset is not supported. Host power cycle is required"
			return
		fi
		sleep 1
	done

	log "INFO: Running NIC Firmware reset"
	save_log
	if [ "X$mode" == "Xmanufacturing" ]; then
		log "INFO: Rebooting..."
	fi
	# 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)
	if [ $? -ne 0 ]; then
		log "INFO: NIC Firmware reset failed"
		log "INFO: $msg"
	else
		log "INFO: NIC Firmware reset done"
	fi
}

update_nic_firmware()
{
	if [ $NIC_FW_UPDATE_DONE -eq 0 ]; then
		fw_update
	fi
}

reset_nic_firmware()
{
	if [ $NIC_FW_UPDATE_DONE -eq 1 ]; then
		if [ $NIC_FW_UPDATE_PASSED -eq 1 ]; then
			# Reset NIC FW
			mount -t ${ROOTFS} ${ROOT_PARTITION} /mnt
			bind_partitions
			fw_reset
			unmount_partitions
		fi
	fi
}
