#!/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"

# 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|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
}

# 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
        fm_program_args="$fm_binary -c $FM_CONFIG_FILE"
    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

    $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\""
}

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
}

clear_guid_entries() {
    sed -i "/^FM_SM_MGMT_PORT_GUID=0x[a-fA-F0-9]\\+$/d" $FM_CONFIG_FILE
    if [ -f "$NVLSM_CONFIG_FILE" ]; then
        sed -i "/^guid 0x[a-fA-F0-9]\\+$/d" $NVLSM_CONFIG_FILE
    fi
    sync
}

# 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
            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
                # 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
            local 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"
}

# 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
                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
}

# 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

# 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
        # 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

    # 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
    clear_guid_entries
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" ] && [ -s "$FM_PID_FILE" ]; then
        echo "Removing fabricmanager PID file: $FM_PID_FILE"
        rm -f "$FM_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

