#!/bin/sh

# flag to capture valid port_guids
capture_flag=false

# flag to indicate pre/post nvl5 generation
post_nvl5=false

# Temp input file
input_file="/tmp/nvl_fmsm_ibstat_input.txt"

# FM config file: /usr/share/nvidia/nvswitch/fabricmanager.cfg
fm_config_file=$1

# FM pid file: /var/run/nvidia-fabricmanager/nvidia-fabricmanager.pid
fm_pid_file=$2

# NVLSM config file: /usr/share/nvidia/nvlsm/nvlsm.cfg
nvlsm_config_file=$3

# NVLSM pid file: /var/run/nvidia-fabricmanager/nvlsm.pid
nvlsm_pid_file=$4

# Required kernel module for fabric manager
required_module_name="ib_umad"

# Check if the required kernel module is loaded
check_required_module() {
    if ! lsmod | grep -q "^$required_module_name"; then
        echo "Kernel module '$required_module_name' has not been loaded, fabric manager cannot be started"
        echo "Please run 'modprobe $required_module_name' before starting fabric manager"
        exit 1
    fi
}

# ibstat output capture in input_file for further processing
launch_ibstat() {
    ibstat_binary=$(which ibstat)
    if [ -z "$ibstat_binary" ]; then
        echo "\"ibstat\" command not found! Please install \"ibstat\"."
        exit 1
    fi
    $ibstat_binary $1 > $input_file
    local ret=$?
    if [ $ret -ne 0 ]; then
        echo "\"$ibstat_binary $1\" failed with error $ret"
        return $ret
    fi
}

launch_fm_proc() {
    fm_binary="/usr/bin/nv-fabricmanager"
    $fm_binary -c $fm_config_file
    local ret="$?"
    if [ "$ret" -ne 0 ]; then
        echo "\"$fm_binary\" failed! Exit code: $ret"
        return $ret
    fi
    echo "Started \"Nvidia Fabric Manager\""
    return $ret
}

remove_cruft_files() {
    if [ -f "$input_file" ]; then
        rm $input_file
    fi
}

# Use -J pidfile and -B daemonize option in config file
launch_nvlsm_proc() {
    nvlsm_binary="/opt/nvidia/nvlsm/sbin/nvlsm"

    if [ ! -f "$nvlsm_binary" ]; then
        echo "\"$nvlsm_binary\" does not exist. Please make sure to install \"Nvidia NVLink Subnet Manager\"(nvlsm) package."
        exit 1
    fi

    if [ ! -x "$nvlsm_binary" ]; then
        echo "\"$nvlsm_binary\" does not have execute permission. Please provide execute permissions or reinstall \"Nvidia NVLink Subnet Manager\"(nvlsm) package."
        exit 1
    fi

    # Check NVLSM config file
    if [ ! -f "$nvlsm_config_file" ]; then
        echo "\"$nvlsm_config_file\" does not exist. Please make sure to install \"Nvidia NVLink Subnet Manager\"(nvlsm) package."
        exit 1
    fi

    $nvlsm_binary -F $nvlsm_config_file -B --pid_file $nvlsm_pid_file
    local ret="$?"
    if [ "$ret" -ne 0 ]; then
        echo "\"$nvlsm_binary\" failed! Exit code: $ret"
        exit $ret
    fi
    echo "Started \"Nvidia NVLink Subnet Manager\""
}

clear_guid_entries() {
    sed -i "/^FM_SM_MGMT_PORT_GUID=0x[a-fA-F0-9]\\+$/d" $fm_config_file && sed -i "/^guid 0x[a-fA-F0-9]\\+$/d" $nvlsm_config_file
}

# process input_file to find port_guids that have has_smi bit set
capture_portguids() {
    # Clear any GUID entries if any present
    clear_guid_entries

    # Parse input file line by line
    while IFS= read -r line; do
        # For each line, check if it contains "Capability mask"
        if echo "$line" | grep -q "Capability mask:"; then
            # Isolate mask value
            hex_value=$(echo "$line" | sed -En 's/.*Capability mask: 0x([a-fA-F0-9]+).*/\1/p')
            prefix_hex_value="0x$hex_value"

            # Check if isSMDisabled bit is unset. Bit position starts at 0
            bit_position=10
            mask=$((1<<bit_position))
            if [ $((prefix_hex_value & mask)) -eq 0 ]; then
                # Set flag to true to parse and grab next
                # occurence of port_guid
                capture_flag=true
            fi
        elif echo "$line" | grep -q "Port GUID:" && [ "$capture_flag" = true ]; then
            # Extract port guid from the line
            port_guid=$(echo "$line" | sed -En 's/.*Port GUID: 0x([a-fA-F0-9]+).*/\1/p')
            prefix_port_guid="0x$port_guid"
            echo "FM_SM_MGMT_PORT_GUID=$prefix_port_guid" >> "$fm_config_file"
            echo "guid $prefix_port_guid" >> "$nvlsm_config_file"
            return 0
        fi
    done < "$input_file"
}

# Loop through each Infiniband device directory
for dir in /sys/class/infiniband/*/device; do
    # Define the path to the VPD file
    vpd_file="$dir/vpd"

    # Check if the VPD file exists
    if [ -f "$vpd_file" ]; then
        # Search for 'SW_MNG' in the VPD file
        if grep -q "SW_MNG" "$vpd_file"; then
            echo "Detected NVL5+ system"
            # Extract the Infiniband device name using parameter expansion
            device_name="${dir%/device}"  # Removes '/device' from the end of $dir
            device_name="${device_name##*/}"  # Extracts the part after the last '/'
            echo "Checking $device_name..."
            launch_ibstat $device_name
            if [ $? -ne 0 ]; then
                remove_cruft_files
                exit 1
            fi
            capture_portguids
            remove_cruft_files
            post_nvl5=true
            # if port guid was captured, return since we are only grabbing first GUID
            if [ "$capture_flag" = "true" ]; then
                echo "Using device $device_name, port $prefix_port_guid"
                break
            fi
        fi
    fi
done

if [ "$post_nvl5" = "true" ]; then
    check_required_module
    launch_nvlsm_proc
    # SM takes few mins to start the GRPC service, hence lets wait before starting FM
    sleep 5
    launch_fm_proc
else
    echo "Detected Pre-NVL5 system"
    launch_fm_proc
fi

# capture FM exit code, the bash script needs to exit with the same code
fm_exit_code=$?

# now exit with the code we have previously captured
exit $fm_exit_code
