#!/bin/sh

# flag to capture valid port_guids
capture_flag=false

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

# prefixed port guid to be passed to FM and NVLSM launch commands
prefix_port_guid=""

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

# default mode for running this script
fm_script_mode="interactive"

# arguments to run fabric manager
fm_program_args=""

# FM config file environment variable
FM_CONFIG_FILE=${FM_CONFIG_FILE:-/usr/share/nvidia/nvswitch/fabricmanager.cfg}

# FM pid file environment variable
FM_PID_FILE=${FM_PID_FILE:-/var/run/nvidia-fabricmanager/nv-fabricmanager.pid}

# NVLSM config file environment variable
NVLSM_CONFIG_FILE=${NVLSM_CONFIG_FILE:-/usr/share/nvidia/nvlsm/nvlsm.conf}

# NVLSM pid file environment variable
NVLSM_PID_FILE=${NVLSM_PID_FILE:-/var/run/nvidia-fabricmanager/nvlsm.pid}

# show help message for this script
show_help() {
    echo "NVIDIA Fabric Manager Start Script"
    echo "Warning: If you are not a fabric manager developer, please avoid running this script directly."
    echo "Please either start or stop using the nvidia-fabricmanager systemd service,"
    echo "or refer to the fabric manager user guide for how to run the fabric manager program."
    echo
    echo "Usage: $0 [script options] --args [arguments to run fabric manager]"
    echo "Script options:"
    echo "  -h, --help              Show this help message"
    echo "  -d, --debug             Enable debug output for this script"
    echo "  -m, --mode              Mode to run this script (start|stop|precheck|interactive)"
    echo "  --fm-config-file        Path to fabric manager configuration file"
    echo "  --fm-pid-file           Path to fabric manager PID file"
    echo "  --nvlsm-config-file     Path to NVLink Subnet Manager configuration file"
    echo "  --nvlsm-pid-file        Path to NVLink Subnet Manager PID file"
    echo "  --args                  Trailing arguments to run fabric manager"
}

# Check if the required kernel module is loaded
check_required_module() {
    local required_module_name="ib_umad"
    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
}

# precheck for which environment FM would be running in
fm_environment_precheck() {
    local timeout=60
    local delay=2
    local ib_device_dir="/sys/class/infiniband"

    # Collect PCI BDFs for Mellanox devices with SW_MNG in their VPD
    local sw_mng_bdfs=""
    for bdf in $(lspci -D | grep "Mellanox" | awk '{print $1}'); do
        local vpd_file="/sys/bus/pci/devices/$bdf/vpd"
        if [ -f "$vpd_file" ] && grep -q "SW_MNG" "$vpd_file" 2>/dev/null; then
            sw_mng_bdfs="$sw_mng_bdfs $bdf"
        fi
    done

    if [ -z "$sw_mng_bdfs" ]; then
        echo "No CX Bridge devices detected!"
        return 0
    fi

    local sw_mng_bdfs_list=""
    for bdf in $sw_mng_bdfs; do
        for ib_dev_path in "$ib_device_dir"/*/device; do
            [ -e "$ib_dev_path" ] || continue
            local resolved_bdf
            resolved_bdf=$(basename "$(readlink -f "$ib_dev_path")")
            if [ "$resolved_bdf" = "$bdf" ]; then
                local dev_name="${ib_dev_path%/device}"
                dev_name="${dev_name##*/}"
                sw_mng_bdfs_list="$sw_mng_bdfs_list $dev_name($bdf)"
            fi
        done
    done
    echo "Found CX Bridge devices: $sw_mng_bdfs_list"

    echo "Checking whether any of the CX Bridge devices are usable for NVLink subnet management..."

    local start_time=$(date +%s)
    while true; do
        cx_bridge_device_usable=false
        for bdf_pair in $sw_mng_bdfs_list; do
            check_port_physical_link_up $bdf_pair
            if [ $? -eq 0 ]; then
                echo "CX Bridge device $bdf_pair is usable for NVLink subnet management."
                cx_bridge_device_usable=true
            fi
        done
        if [ "$cx_bridge_device_usable" = true ]; then
            break
        fi
        local current_time=$(date +%s)
        if [ $((current_time - start_time)) -ge "$timeout" ]; then
            echo "CX Bridge devices detected ($sw_mng_bdfs), but links physical state did not transition to the LinkUp state within $timeout seconds"
            exit 1
        fi
        sleep "$delay"
    done
}

check_port_physical_link_up() {
    local bdf_pair="$1"
    local dev_name=$(echo "$bdf_pair" | cut -d'(' -f1)
    local bdf=$(echo "$bdf_pair" | cut -d'(' -f2 | cut -d')' -f1)
    for port_dir in "$ib_device_dir/$dev_name/ports"/*/; do
        [ -d "$port_dir" ] || continue
        local port_num
        port_num=$(basename "$port_dir")
        local cap_mask_file="${port_dir}cap_mask"
        [ -f "$cap_mask_file" ] || continue
        local cap_mask
        cap_mask=$(cat "$cap_mask_file" 2>/dev/null) || continue
        # Skip ports with isSMDisabled (bit 10) set
        if [ $((cap_mask & 0x400)) -ne 0 ]; then
            echo "Skipping device $dev_name (BDF $bdf) as port $port_num is not NVLink subnet management capable"
            continue
        fi
        local phys_state_file="${port_dir}phys_state"
        echo "Device $dev_name (BDF $bdf) Port $port_num is $(cat "$phys_state_file" 2>/dev/null)."
        if [ -f "$phys_state_file" ] && grep -q "LinkUp" "$phys_state_file" 2>/dev/null; then
            return 0
        fi
    done
    return 1
}

# ibstat output capture in input_file for further processing
launch_ibstat() {
    local 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() {
    local fm_binary="/usr/bin/nv-fabricmanager"
    if [ -z "$fm_program_args" ]; then
        if [ -z $prefix_port_guid ]; then
            fm_program_args="$fm_binary -c $FM_CONFIG_FILE"
        else
            # if port guid is captured, pass it to FM
            fm_program_args="$fm_binary -c $FM_CONFIG_FILE -g $prefix_port_guid"
        fi
    fi
    $fm_program_args
    local ret="$?"
    if [ "$ret" -ne 0 ]; then
        echo "\"$fm_program_args\" failed! Exit code: $ret"
    else
        echo "Started \"Nvidia Fabric Manager\""
    fi
    return $ret
}

stop_fm_proc() {
    if [ -f "$FM_PID_FILE" ] && [ -s "$FM_PID_FILE" ]; then
        kill "$(cat "$FM_PID_FILE")"
        echo "Stopped \"Nvidia Fabric Manager\""
    fi
}

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() {
    local 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

    # check if prefixed port guid is set
    if [ -z "$prefix_port_guid" ]; then
        echo "No port guid was captured. Please check the IB devices and ensure they are NVL5+ capable."
        exit 1
    fi

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

stop_nvlsm_proc() {
    if [ -f "$NVLSM_PID_FILE" ] && [ -s "$NVLSM_PID_FILE" ]; then
        kill "$(cat "$NVLSM_PID_FILE")"
        echo "Stopped \"Nvidia NVLink Subnet Manager\""
    fi
}

# process input_file to find port_guids that have has_smi bit set
capture_portguids() {
    # Parse input file line by line.  Input will be a single interface block from ibstat should look something like:
    # CA 'mlx5_7'
    #        CA type: MT4133
    #        Number of ports: 1
    #        Firmware version: 82.97.9814
    #        Hardware version: 0
    #        Node GUID: 0x4cbb4703004a2433
    #        System image GUID: 0x4cbb4703004a2430
    #        Port 1:
    #                State: Initializing
    #                Physical state: LinkUp
    #                Rate: 400
    #                Base lid: 65535
    #                LMC: 0
    #                SM lid: 1
    #                Capability mask: 0xa741ec48
    #                Port GUID: 0x4cbb4703004a2433
    #                Link layer: InfiniBand
    local has_capability=false
    local physical_state_ready=false
    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
            local hex_value=$(echo "$line" | sed -En 's/.*Capability mask: 0x([a-fA-F0-9]+).*/\1/p')
            local prefix_hex_value="0x$hex_value"

            # Check if isSMDisabled bit is unset. Bit position starts at 0
            # @note avoiding a nasty syntax highlighting bug in both vscode and vim
            # bit_position=10
            # mask=$((1<<bit_position))
            local mask=1024
            if [ $((prefix_hex_value & mask)) -eq 0 ]; then
                has_capability=true
            fi
        elif echo "$line" | grep -q "Physical state:"; then
            local physical_port_state=$(echo "$line" | sed -En 's/.*Physical state: (.+).*/\1/p')
            if [ "$physical_port_state" = "LinkUp" ]; then
                physical_state_ready=true
            fi
        elif echo "$line" | grep -q "Port GUID:" && [ "$has_capability" = true ] && [ "$physical_state_ready" = true ]; then
            # Extract port guid from the line
            local port_guid=$(echo "$line" | sed -En 's/.*Port GUID: 0x([a-fA-F0-9]+).*/\1/p')
            capture_flag=true
            prefix_port_guid="0x$port_guid"
            return 0
        fi
    done < "$input_file"
}

# detect whether we are using NVL5+ system
detect_nvl5_system() {
    # Loop through each Infiniband device directory
    for dir in /sys/class/infiniband/*/device; do
        # Define the path to the VPD file
        local 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
                # Extract the Infiniband device name using parameter expansion
                local 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
                # if port guid was captured, return since we are only grabbing first GUID
                if [ "$capture_flag" = "true" ]; then
                    post_nvl5=true
                    echo "Using device $device_name, port $prefix_port_guid"
                    break
                fi
            fi
        fi
    done
}

# quick and dirty argument parser
while [ "$#" -gt 0 ]; do
    case "$1" in
        -h|--help)
            show_help
            exit 0
            ;;
        -d|--debug)
            set -x
            shift
            ;;
        -m|--mode)
            fm_script_mode="$2"
            shift 2
            ;;
        --fm-config-file)
            FM_CONFIG_FILE="$2"
            shift 2
            ;;
        --fm-pid-file)
            FM_PID_FILE="$2"
            shift 2
            ;;
        --nvlsm-config-file)
            NVLSM_CONFIG_FILE="$2"
            shift 2
            ;;
        --nvlsm-pid-file)
            NVLSM_PID_FILE="$2"
            shift 2
            ;;
        --args)
            shift
            fm_program_args="$@"
            break
            ;;
        *)
            echo "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
done

# specific logic for precheck mode; script should immmediately exit after this
if [ "$fm_script_mode" = "precheck" ]; then
    fm_environment_precheck
    exit $?
fi

# in interactive mode, install a trap for signals
if [ "$fm_script_mode" = "interactive" ]; then
    trap 'echo; echo "Caught signal in interactive mode, continuing this script..."' HUP INT TERM QUIT
fi

# start FM logic, only executed in start and interactive mode
if [ "$fm_script_mode" = "start" ] || [ "$fm_script_mode" = "interactive" ]; then
    detect_nvl5_system
    if [ "$post_nvl5" = "true" ]; then
        echo "Detected NVL5+ system"
        check_required_module
        launch_nvlsm_proc
        launch_fm_proc
    else
        echo "Detected Pre-NVL5 system"
        launch_fm_proc
    fi

    # if not interactive mode, the bash script needs to exit with the same code
    fm_exit_code=$?
    if [ "$fm_script_mode" = "start" ]; then
        exit $fm_exit_code
    fi
fi

# if FM is running in daemon mode, ask for interrupt to stop it
if [ "$fm_script_mode" = "interactive" ] && [ "$fm_exit_code" -eq 0 ]; then
    # $daemonize is 0 or 1
    daemonize=$(sed -n 's/^[[:space:]]*DAEMONIZE=\([01]\)$/\1/p' "$FM_CONFIG_FILE")

    if [ "$daemonize" -eq 1 ]; then
        echo "Fabric Manager is running in daemon mode. Press ctrl-c to stop it..."
        sleep infinity
    fi
fi

# stop FM logic, only executed in stop and interactive mode
if [ "$fm_script_mode" = "stop" ] || [ "$fm_script_mode" = "interactive" ]; then
    stop_fm_proc
    stop_nvlsm_proc
fi

# cleanup logic for interactive mode only
if [ "$fm_script_mode" = "interactive" ]; then
    # wait for 2 seconds, which is more than enough for a normal shutdown
    echo "Waiting for processes to cleanup in interactive mode..."
    sleep 2

    # interactive means we are not using systemd, so need to remove PID file
    if [ -f "$FM_PID_FILE" ] || [ -f "$NVLSM_PID_FILE" ]; then
        echo "Removing fabricmanager and nvlsm PID files: $FM_PID_FILE, $NVLSM_PID_FILE"
        rm -f "$FM_PID_FILE" "$NVLSM_PID_FILE"
    fi

    # kill -9 all dangling processes that did not exit upon kill signal
    # @note systemd takes care of this by killing all processes in the cgroup upon stopping service
    dangling_prog_names="nv-fabricmanager nvlsm"
    for prog_name in $dangling_prog_names; do
        dangling_pids=$(pidof "$prog_name")
        for pid in $dangling_pids; do
            cmd=$(cat /proc/$pid/cmdline | tr '\0' ' ')
            echo "Killing dangling process with PID $pid: $cmd"
            kill -9 "$pid"
        done
    done
fi
