#! /bin/bash
#
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This software product is a proprietary product of Mellanox Technologies Ltd.
# (the "Company") and all right, title, and interest in and to the software
# product, including all associated intellectual property rights, are and
# shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
#

# Supported Mellanox PCI device IDs (vendor 15b3):
#   a2d2 = ConnectX-5
#   a2d6 = ConnectX-5 Ex
#   a2dc = ConnectX-6 Dx / BlueField-3 (DPA)
#   a2df = BlueField-4 (DPA)
#   101d = ConnectX-7
#   1021 = ConnectX-8 (DPA)
#   1025 = ConnectX-9 (DPA)
DEVICE_ID_REGEX='a2d[26cf]|101d|1021|1025'

# DPA-capable devices (need dpa.provider generation)
DPA_DEVICE_REGEX='a2d[cf]|1021|1025'

# BlueField-4, used to detect a Socket Direct setup (see below)
BF4_DEVICE_ID='a2df'

# Helper functions for JSON manipulation
# Prefer jq if available, fallback to sed/awk for compatibility

# Read a JSON field value
json_get() {
	local file="$1"
	local field="$2"
	local default="$3"

	if [[ ! -f "$file" ]]; then
		echo "${default:-null}"
		return
	fi

	# Use jq if available (robust, handles all JSON edge cases)
	if command -v jq &> /dev/null; then
		local value=$(jq -r ".${field} // empty" "$file" 2>/dev/null)
		if [[ -z "$value" ]]; then
			echo "${default:-null}"
		else
			echo "$value"
		fi
		return
	fi

	# Fallback: Use grep/sed/awk to extract JSON field value
	local value=$(grep -o "\"${field}\"[[:space:]]*:[[:space:]]*[^,}]*" "$file" 2>/dev/null | \
		     sed 's/.*:[[:space:]]*//' | \
		     sed 's/^"\(.*\)"$/\1/' | \
		     sed 's/[[:space:]]*$//')

	if [[ -z "$value" ]]; then
		echo "${default:-null}"
	else
		echo "$value"
	fi
}

# Set a JSON field value
json_set() {
	local file="$1"
	local field="$2"
	local value="$3"
	local numeric="$4"   # optional: non-empty => write value as a bare JSON number/bool

	# Render the value as it appears in JSON: a quoted string by default, or a
	# bare token (number/bool) when "numeric" is set. Everything below uses
	# $jval so string and numeric callers share one code path.
	local jval="\"${value}\""
	[[ -n "$numeric" ]] && jval="${value}"

	# Use jq if available (robust, handles all JSON edge cases)
	if command -v jq &> /dev/null; then
		if [[ -f "$file" ]]; then
			# Update existing file
			local tmp=$(mktemp)
			jq ".${field} = ${jval}" "$file" > "$tmp" && mv "$tmp" "$file"
		else
			# Create new file
			echo "{}" | jq ".${field} = ${jval}" > "$file"
		fi
		return
	fi

	# Fallback: Use sed/awk for JSON manipulation
	if [[ -f "$file" ]]; then
		# Check if field already exists
		if grep -q "\"${field}\"[[:space:]]*:" "$file"; then
			# Update existing field (value may be a quoted string or a bare number)
			sed -i "s/\"${field}\"[[:space:]]*:[[:space:]]*\(\"[^\"]*\"\|[0-9][0-9]*\)/\"${field}\": ${jval}/g" "$file"
		else
			# Remove any trailing whitespace and ensure proper line ending
			sed -i 's/[[:space:]]*$//' "$file"

			# Check if file is empty object {} - handle specially to avoid trailing comma
			# Normalize the file content for comparison (remove all whitespace)
			local file_content=$(cat "$file" | tr -d '[:space:]')
			if [[ "$file_content" == "{}" ]]; then
				# Empty object: write new content directly without comma
				echo "{" > "$file"
				echo "  \"${field}\": ${jval}" >> "$file"
				echo "}" >> "$file"
			else
				# Check if this is a single-line JSON or multi-line
				# Handle case where file has no trailing newline (wc -l returns 0)
				local line_count=$(wc -l < "$file")
				if [[ $line_count -le 1 ]]; then
					# Single line JSON (including files with no trailing newline): add comma and field
					sed -i 's/}[[:space:]]*$/,\n  "'"${field}"'": '"${jval}"'\n}/' "$file"
				else
					# Multi-line JSON: add comma and field with proper indentation
					sed -i '$s/}[[:space:]]*$/,\n  "'"${field}"'": '"${jval}"'\n}/' "$file"
				fi
			fi
		fi
	else
		# Create new file
		echo "{\"${field}\": ${jval}}" > "$file"
	fi
}

# Delete a JSON field if present
json_del() {
	local file="$1"
	local field="$2"

	[[ -f "$file" ]] || return

	if command -v jq &> /dev/null; then
		local tmp=$(mktemp)
		jq "del(.${field})" "$file" > "$tmp" && mv "$tmp" "$file"
		return
	fi

	# Fallback: remove the field line (with optional trailing comma) via sed
	if grep -q "\"${field}\"[[:space:]]*:" "$file"; then
		sed -i "/\"${field}\"[[:space:]]*:/d" "$file"
		# Drop a dangling comma left on the previous line (common with multi-line JSON)
		sed -i ':a;N;$!ba;s/,\([[:space:]]*\)}/\1}/g' "$file"
	fi
}

mftconfig=mstconfig
if [ -x /usr/bin/mlxconfig ]; then
	mftconfig=mlxconfig
fi

devlink=devlink
if [ -x /opt/mellanox/iproute2/sbin/devlink ]; then
	devlink=/opt/mellanox/iproute2/sbin/devlink
fi

pci=`lspci -nD -d 15b3: | grep -E "$DEVICE_ID_REGEX" | head -n 1`
emu_manager=`echo $pci | cut -d ' ' -f 1`

if [[ -f /opt/mellanox/mlnx_virtnet/virtnet.conf ]]; then
	_emu_manager=$(json_get "/opt/mellanox/mlnx_virtnet/virtnet.conf" "ib_dev_p0" "")
	if [[ ! -z $_emu_manager ]] && [[ $_emu_manager != "null" ]]; then
		# Check if the infiniband device path exists
		if [[ -L "/sys/class/infiniband/$_emu_manager/device" ]]; then
			device_path=$(readlink "/sys/class/infiniband/$_emu_manager/device")
			if [[ ! -z "$device_path" ]]; then
				emu_manager=`basename "$device_path"`
			fi
		fi
	fi
fi

if [[ -z $emu_manager ]]; then
	echo "Error: no valid emulation manager is defined"
	exit 1
fi

config_file="/opt/mellanox/mlnx_virtnet/virtnet.conf"

# BF4 Socket Direct: a single RDMA device is reached over two PCI links. lspci
# then reports 4 BF4 functions (two per link) and the links themselves are the
# ".0" ones. Both links are recorded so the controller creates a port context
# each. Two bond devices would mean each link has its own RDMA device and a
# port context each, therefore no need to record the PCI addresses.
socket_direct=0
bond_ibdevs=$(ls /sys/class/infiniband/ 2>/dev/null | grep -c '^mlx5_bond_')
bf4_funcs=$(lspci -D -d 15b3:${BF4_DEVICE_ID} 2>/dev/null | awk '{print $1}')
bf4_links=$(echo "$bf4_funcs" | grep '\.0$' | sort)
if [[ $bond_ibdevs -eq 1 && \
      $(echo "$bf4_funcs" | grep -c .) -eq 4 && \
      $(echo "$bf4_links" | grep -c .) -eq 2 ]]; then
	socket_direct=1
	pci_addr_p0=$(echo "$bf4_links" | sed -n '1p')
	pci_addr_p1=$(echo "$bf4_links" | sed -n '2p')
fi

# LAG auto-config (bond-only): when no ib device is configured and the first
# emulation-capable function maps to a bond, record it as a LAG (is_lag=1 +
# bond ib_device) so the controller drives the single bonded device instead of
# failing the second-port init. Non-bond setups are left untouched.
_p0=$(json_get "$config_file" "ib_dev_p0" "")
_lag=$(json_get "$config_file" "ib_dev_lag" "")
if [[ ( -z $_p0 || $_p0 == "null" ) && ( -z $_lag || $_lag == "null" ) ]]; then
	ib_dev=$(ls /sys/bus/pci/devices/$emu_manager/infiniband/ 2>/dev/null | head -n 1)
	if [[ $ib_dev == mlx5_bond_* ]]; then
		echo "LAG emulation manager $ib_dev detected; setting is_lag=1, ib_dev_lag, ib_dev_for_static_pf"
		json_set "$config_file" "ib_dev_lag" "$ib_dev"
		json_set "$config_file" "ib_dev_for_static_pf" "$ib_dev"
		json_set "$config_file" "is_lag" "1" num
	fi
fi

# Add pci_addr_* when Socket Direct is first enabled (is_lag already set), and
# remove them when Socket Direct is no longer detected so a prior SD config
# does not linger.
_is_lag=$(json_get "$config_file" "is_lag" "0")
if [[ $socket_direct -eq 1 && $_is_lag == "1" ]]; then
	echo "BF4 Socket Direct detected; setting pci_addr_p0=$pci_addr_p0, pci_addr_p1=$pci_addr_p1"
	json_set "$config_file" "pci_addr_p0" "$pci_addr_p0"
	json_set "$config_file" "pci_addr_p1" "$pci_addr_p1"
else
	_pci0=$(json_get "$config_file" "pci_addr_p0" "")
	_pci1=$(json_get "$config_file" "pci_addr_p1" "")
	if [[ ( -n $_pci0 && $_pci0 != "null" ) || \
	      ( -n $_pci1 && $_pci1 != "null" ) ]]; then
		echo "Not a Socket Direct LAG setup; removing pci_addr_p0, pci_addr_p1"
		json_del "$config_file" "pci_addr_p0"
		json_del "$config_file" "pci_addr_p1"
	fi
fi

p2=`$mftconfig -d ${emu_manager} -e q LINK_TYPE_P2`
if [[ -n `echo "$p2" 2> /dev/null | grep "Unknown Parameter"` ]]; then
	echo "Single port NIC"
	json_set "$config_file" "single_port" "1"
fi

if [[ -z `modinfo mlxdevm 2>&1 | grep -E "ERROR.*mlxdevm not found"` ]]; then
	if [[ -f $config_file ]]; then
		current_sf_provider=$(json_get "$config_file" "sf_provider" "")
		if [[ "$current_sf_provider" != "devlink" ]]; then
			echo "mlxdevm as SF provider"
			json_set "$config_file" "sf_provider" "mlxdevm"
		else
			echo "mlxdevm available, keeping existing devlink SF provider"
		fi
	fi
else
	echo "devlink as SF provider"
	json_set "$config_file" "sf_provider" "devlink"
fi

dpa_device=`echo $pci | grep -E "$DPA_DEVICE_REGEX"`

if [[ -n $dpa_device ]]; then
	if [[ ! -e /opt/mellanox/mlnx_virtnet/providers/dpa.provider ]]; then
		echo "Generating default dpa.provider for DPA-capable device"
		echo -e "Provider=libprovider-dpa\nScore=200" > \
		      /opt/mellanox/mlnx_virtnet/providers/dpa.provider
	fi
fi

recovery_dir="/opt/mellanox/mlnx_virtnet/recovery"

is_lag_cur=0
if [[ -f $config_file ]]; then
	is_lag_cur=$(json_get "$config_file" "is_lag" "0")
fi

is_lag_cur_valid=1
if [[ "$is_lag_cur" != "0" && "$is_lag_cur" != "1" ]]; then
	echo "Invalid is_lag value (${is_lag_cur}) in $config_file; expected 0 or 1. Skipping recovery purge/update; controller will fail config validation."
	is_lag_cur_valid=0
fi

mkdir -p ${recovery_dir}
if [[ -f "${recovery_dir}/.virtnet_conf_save" ]]; then
	is_lag_save=$(json_get "${recovery_dir}/.virtnet_conf_save" "is_lag" "0")
	is_lag_save_valid=1
	if [[ "$is_lag_save" != "0" && "$is_lag_save" != "1" ]]; then
		is_lag_save_valid=0
	fi

	if [[ $is_lag_cur_valid -eq 1 && $is_lag_save_valid -eq 1 && "$is_lag_save" != "$is_lag_cur" ]]; then
		cd $recovery_dir
		find . -type f -not \( -name ".virtnet_conf_save" -or -name ".mlxconfig_save" \) -delete
	fi
fi

if [[ -f $config_file && $is_lag_cur_valid -eq 1 ]]; then
	cp "$config_file" "${recovery_dir}/.virtnet_conf_save"
elif [[ ! -f $config_file ]]; then
	echo '{"is_lag": 0}' > "${recovery_dir}/.virtnet_conf_save"
fi

curr_map_count=`cat /proc/sys/vm/max_map_count`

#Define map threshold value to 131060 (65530 * 2) to support upto 2K devices
max_map_threshold=131060

if [[ "$curr_map_count" -lt "$max_map_threshold" ]]; then
	echo $max_map_threshold > /proc/sys/vm/max_map_count
fi

fields=(PF_BAR2_ENABLE \
	HIDE_PORT2_PF \
	PER_PF_NUM_SF \
	SRIOV_EN \
	PF_SF_BAR_SIZE \
	PF_TOTAL_SF)

fields_emu_per_pf=(VIRTIO_NET_EMU_MNG_ENABLE \
	PCI_SWITCH_EMU_MNG_ENABLE \
	PCI_SWITCH_EMU_MNG_NUM_PORT \
	VIRTIO_NET_EMU_MNG_NUM_VF \
	VIRTIO_NET_EMU_MNG_NUM_PF \
	VIRTIO_NET_EMU_MNG_NUM_MSIX)

fields_emu_global=(VIRTIO_NET_EMULATION_ENABLE \
	PCI_SWITCH_EMULATION_ENABLE \
	PCI_SWITCH_EMULATION_NUM_PORT \
	VIRTIO_NET_EMULATION_NUM_VF \
	VIRTIO_NET_EMULATION_NUM_PF \
	VIRTIO_NET_EMULATION_NUM_MSIX)

is_emu_per_pf=$($mftconfig -d ${emu_manager} -e q VIRTIO_NET_EMU_MNG_ENABLE |
	grep -o "VIRTIO_NET_EMU_MNG_ENABLE.*" | awk '{print $3}' | grep True)
if [[ -z $is_emu_per_pf ]]; then
	fields=("${fields[@]}" "${fields_emu_global[@]}")
else
	fields=("${fields[@]}" "${fields_emu_per_pf[@]}")
fi
configs=$($mftconfig -d ${emu_manager} -e q "${fields[@]}")

fields_str=""
for f in ${fields[@]};
do
	if [[ $fields_str != "" ]]; then
		fields_str="${fields_str}|${f}"
	else
		fields_str="${f}"
	fi
done

mlxconfig_cur=`echo "$configs" | grep -E $fields_str | tr -s ' '`

if [[ ! -f "${recovery_dir}/.mlxconfig_save" ]]; then
	mkdir -p ${recovery_dir}
	echo "$mlxconfig_cur" > ${recovery_dir}/.mlxconfig_save
	exit 0
fi

mlxconfig_save=`cat ${recovery_dir}/.mlxconfig_save | grep -E $fields_str | tr -s ' '`

for f in ${fields[@]};
do
	val_save=`echo "$mlxconfig_save" | grep -o "$f.*" | awk '{print $3}'`
	val_cur=`echo "$mlxconfig_cur" | grep -o "$f.*" | awk '{print $3}'`
	if [[ "$val_save" != "$val_cur" ]]; then
		echo "$mlxconfig_cur" > ${recovery_dir}/.mlxconfig_save
		cd $recovery_dir
		find . -type f -not \( -name ".virtnet_conf_save" -or -name ".mlxconfig_save" \) -delete
		exit 0
	fi
done

exit 0
