#!/bin/bash

DO_JSON=0
MY_PATH=$(dirname $0)
BOOT_DRIVES="unknown"

dgx1_boot_drives()
{
    my_device=$(get_sd_from_list "${1}" 1)
    echo ${my_device}
}

dgx2_boot_drives()
{
    # Taken from nv-pre-install.sh
    BDF_LIST="0000:01:00.0 0000:05:00.0"
    BOOTDEVS=""
    SYSFS_PATH="/sys/bus/pci/devices"

    #
    # Find all M.2 drives attached to the BDF
    #
    for BDF in $(echo ${BDF_LIST}); do
        NVME_PATH="${SYSFS_PATH}/${BDF}/nvme*/nvme*/uevent"

        # uevent has the following fields: MAJOR, MINOR, DEVNAME
        # DEVNAME is the variable we are interested in. It points
        # nvme controller instance. Grab the controller instance
        # and append namespace "n1" to it.
        if [ -f ${NVME_PATH} ]; then
            NVME_NAME=$(cat ${NVME_PATH} | grep DEVNAME | cut -d= -f2)
            BOOTDEVS="${BOOTDEVS} /dev/${NVME_NAME}n1"
        fi
    done

    echo ${BOOTDEVS} | sed -e 's/^[ ]*//'
    return
}

# Taken from DCS nv-pre-install.sh
dcs_boot_drives()
{
    ODM=$(dmidecode --string system-product-name | grep -o "920-23887-2500-000")
    if [ "$ODM" = "920-23887-2500-000" ]; then
        BIOS_VERSION=$(dmidecode --string bios-version | cut -c5-8)
        OUT=$(printf "\n$BIOS_VERSION\n2A07" | sort)
        OUT_1=$(echo $OUT | cut -d ' ' -f1)

        # if BIOS_VERSION is >= 2A07
        if [ "$OUT_1" = '2A07' ]; then
            # This section taken from nvidia-nvme-bdf
            M2_BDF=$(dcs_get_m2_bdfs)
            BOOTDEVS=$(bdf_to_nvme_map ${M2_BDF})
        else
            # use model number to identify M.2
            BOOTDEVS=$(dcs_get_m2_by_model)
        fi
    # if ODM is NOT QT then use model number to identify the M.2
    else
        # use model number to identify M.2
        BOOTDEVS=$(dcs_get_m2_by_model)
    fi
    echo ${BOOTDEVS}
}

dcs_legacy_boot_drives()
{
    BOOTDEVS=$(dcs_get_m2_by_model)
    echo ${BOOTDEVS}
}

c2_ovx_boot_drives()
{
    BOOTDEVS=$(c2_get_m2_by_model)
    echo ${BOOTDEVS}
}

c2_io_boot_drives()
{
    BOOTDEVS=$(c2_get_m2_by_model)
    echo ${BOOTDEVS}
}

#
#  The following function is used for A100, A800, H100, H200, H800, B200, and B300
dgx_xx00_boot_drives()
{
    BOOT_DRIVES=""
    DGX_XX00_BDF_LIST=$(dgx_xx00_get_M2_BDF)
    SYSFS_PATH="/sys/bus/pci/devices"

    #
    # Find all M.2 drives attached to the BDF
    #
    for BDF in `echo ${DGX_XX00_BDF_LIST}`; do
        NVME_PATH="${SYSFS_PATH}/${BDF}/nvme*/nvme*/uevent"

        # uevent has the following fields: MAJOR, MINOR, DEVNAME
        # DEVNAME is the variable we are interested in. It points
        # nvme controller instance. Grab the controller instance
        # and append namespace "n1" to it.
        if [ -f ${NVME_PATH} ]; then
            NVME_NAME=`cat ${NVME_PATH} | grep DEVNAME | cut -d= -f2`
            BOOT_DRIVES="${BOOT_DRIVES} /dev/${NVME_NAME}n1"
        fi
    done

    echo ${BOOT_DRIVES} | sed -e 's/^[ ]*//'
}

dgx_a100_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_h100_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_h200_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_b200_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_b300_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_gb200_boot_drives()
{
    l4t_nv_get_dual_boot_drives "1" "2"
}

dgx_gb300_boot_drives()
{
    l4t_nv_get_dual_boot_drives "1" "2"
}

dgx_gb300ws_boot_drives()
{
    l4t_nv_get_dual_boot_drives "A" "B"
}

dgx_a800_boot_drives()
{
    dgx_xx00_boot_drives
}

dgx_h800_boot_drives()
{
    dgx_xx00_boot_drives
}

dcs_get_m2_by_model()
{
    NVME_DEV_LIST="nvme0n1 nvme1n1 nvme2n1 nvme3n1"
    M2_LIST="MZ1LB960HAJQ-00007 MZ1LW960HMJP-00003 SSDPEKKA010T7 SSDPEKKA010T8"
    U2_LIST="MZQLB3T8HALS-00007 SSDPE2KX040T7 MTFDHAL3T8TCT"
    M2_DRIVES=""
    SYSFS_PATH="/sys/block"
    #
    # Find M2
    #
    for DEV in $(echo ${NVME_DEV_LIST}); do
        NVME_PATH="${SYSFS_PATH}/${DEV}/device/model"

        if [ -f ${NVME_PATH} ]; then
           NVME_MODEL=$(cat ${NVME_PATH})
           for MODEL in $(echo ${M2_LIST}); do
               if echo "${NVME_MODEL}" | grep -q "${MODEL}"; then
                    M2_DRIVES="${M2_DRIVES} /dev/${DEV}"
               fi
           done

        fi
    done

    M2_DRIVES=$(echo ${M2_DRIVES} | sed 's/^ *//')
    echo "${M2_DRIVES}"
}

try_get_nvme_name_from_bdf()
{
    local BDF_LIST="${1}"
    local SYSFS_PATH="/sys/bus/pci/devices"
    local BOOT_DRIVES=""

    for BDF in $(echo ${BDF_LIST}); do
        NVME_PATH="${SYSFS_PATH}/${BDF}/nvme*/nvme*/uevent"

        # uevent has the following fields: MAJOR, MINOR, DEVNAME
        # DEVNAME is the variable we are interested in. It points
        # nvme controller instance. Grab the controller instance
        # and append namespace "n1" to it.
        if [ -f ${NVME_PATH} ]; then
            NVME_NAME=$(cat ${NVME_PATH} | grep DEVNAME | cut -d= -f2)
            BOOT_DRIVES="${BOOT_DRIVES} /dev/${NVME_NAME}n1"
        fi
    done

    echo ${BOOT_DRIVES} | sed -e 's/^[ ]*//' | sed 's:/dev/::g'
}

get_nvme_name_by_slot()
{
    local SLOTNO="${1}"
    local BDF_DIRECT=$(get_nvme_bdf_direct_by_slot "${SLOTNO}")
    local BDF_DOWNSTREAM=$(get_nvme_bdf_downstream_by_slot "${SLOTNO}")

    # Try searching for NVME devices directly first
    local BOOT_DRIVES=$(try_get_nvme_name_from_bdf "${BDF_DIRECT}")

    # Try searching for NVME devices downstream of port
    if [[ -z "${BOOT_DRIVES}" ]]; then
        BOOT_DRIVES=$(try_get_nvme_name_from_bdf "${BDF_DOWNSTREAM}")
    fi

    echo "${BOOT_DRIVES}"
}

# Returns BDF of M.2
get_nvme_bdf_direct_by_slot()
{
    local SLOTNO="${1}"
    local DMIOUT=$(dmidecode --type 9 | grep -E "Designation.* Slot ${SLOTNO}|Bus Address" | grep -A1 'M2\|M\.2' | grep -v 'M2\|M\.2')
    local SLOTN_BDF=$(echo ${DMIOUT} | sed 's/Bus Address://g' | awk '{$1=$1};1' | cut -d' ' -f1)
    echo "${SLOTN_BDF}"
}

# Can return list of M.2 BDFs downstream from a bridge/switch
get_nvme_bdf_downstream_by_slot()
{
    local SLOTNO="${1}"
    local DMIOUT=$(dmidecode --type 9 | grep -E "Designation.* Slot ${SLOTNO}|Bus Address" | grep -A1 'M2\|M\.2' | grep -v 'M2\|M\.2')
    local SLOTN_BDF=$(echo ${DMIOUT} | sed 's/Bus Address://g' | awk '{$1=$1};1' | cut -d' ' -f1)
    local M2_LIST=""
    local BDF_LIST=""

    for DEV in $(echo ${SLOTN_BDF}); do
        SYSFS_PATH="/sys/bus/pci/devices/${DEV}/"
        BDF_LIST=$(find "${SYSFS_PATH}" -maxdepth 1 -type l -or -type d -name '0*' 2> /dev/null | xargs -n1 basename 2> /dev/null)

        for NVME_BDF in $(echo ${BDF_LIST}); do
            DIR="${SYSFS_PATH}/${NVME_BDF}/nvme"
            if [ -d "${DIR}" ]; then
                M2_LIST="${M2_LIST} ${NVME_BDF}"
            fi
        done
    done

    # remove any leading whitespace
    M2_LIST=$(echo ${M2_LIST} | sed 's/^ *//')
    echo "${M2_LIST}"
}

l4t_ut2_1_boot_drives()
{
    slot1_name=$(get_nvme_name_by_slot "1")
    echo "/dev/${slot1_name}"
}

l4t_keystone_boot_drives()
{
    l4t_nv_get_dual_boot_drives "1" "2"
}

l4t_oberon_boot_drives()
{
    l4t_nv_get_dual_boot_drives "1" "2"
}

l4t_nv_get_dual_boot_drives()
{
    slot1_no="${1}"
    slot2_no="${2}"
    slot1_name=$(get_nvme_name_by_slot "${slot1_no}")
    slot2_name=$(get_nvme_name_by_slot "${slot2_no}")
    retval=""

    for i in $(echo "${slot1_name} ${slot2_name}"); do
        if [[ ! -z "${i}" ]]; then
            retval="${retval} /dev/${i}"
        fi
    done

    echo "${retval}"
}

l4t_smcmgx_boot_drives()
{
    slot_names=$(smc_get_nvme_names)
    retval=""

    for i in $(echo "${slot_names}"); do
        if [[ ! -z "${i}" ]]; then
            retval="${retval} /dev/${i}"
        fi
    done

    echo "${retval}"
}

smc_get_nvme_names()
{
    local nvme_bdfs=$(smc_get_nvme_bdfs_from_dmi)
    local nvme_names=$(try_get_nvme_name_from_bdf "${nvme_bdfs}")

    echo "${nvme_names}"
}

smc_get_nvme_bdfs_from_dmi()
{
    # SMC DMI type 9 example:
    #
    # Handle 0x0034, DMI type 9, 24 bytes
    # System Slot Information
    #     Designation: C1U3-2
    #     Type: x4 M.2 Socket 3
    #     Current Usage: In Use
    #     Length: Other
    #     Characteristics: None
    #     Bus Address: 0007:01:00.0
    #     Data Bus Width: 10
    #     Peer Devices: 0
    #     Slot Physical Width: x4
    #     Height: Other
    #
    # Handle 0x0040, DMI type 9, 24 bytes
    # System Slot Information
    #     Designation: C2U3-2
    #     Type: x4 M.2 Socket 3
    #     Current Usage: In Use
    #     Length: Other
    #     Characteristics: None
    #     Bus Address: 0017:01:00.0
    #     Data Bus Width: 10
    #     Peer Devices: 0
    #     Slot Physical Width: x4
    #     Height: Other
    local DMIOUT=$(sudo dmidecode --type 9 | grep -E "Type:.*M.2|Bus Address" | grep -A1 'M2\|M\.2' | grep -v 'M2\|M\.2' | grep -v -- "^--$")
    # Now we should be left with (separate lines, with leading tabs):
    #	Bus Address: 0007:01:00.0
    #	Bus Address: 0017:01:00.0

    local cleaned_dmi=$(echo "${DMIOUT}" | sed 's/^\tBus Address://g')
    # Now we should be left with (separate lines, with leading spaces):
    # 0007:01:00.0
    # 0017:01:00.0

    local retval=""

    for token in $(echo "${cleaned_dmi}"); do
        retval+=$(echo "${token} ")
    done

    echo "${retval}"
}

dcs_get_m2_bdfs()
{
    PCI_BUS_LIST=$(sudo dmidecode --type 9 | grep 'Designation:\|Bus Address:' | grep -A1 M.2 | grep Bus | cut -d' ' -f3)
    M2_LIST=""

    for DEV in $(echo ${PCI_BUS_LIST}); do
        BDF_LIST=$(ls -al /sys/bus/pci/devices/$DEV/0* -d | cut -d '/' -f7)
        for NVME_BDF in $(echo ${BDF_LIST}); do
            DIR=/sys/bus/pci/devices/$DEV/$NVME_BDF/nvme
            if [ -d "$DIR" ]; then
                M2_LIST="$M2_LIST $NVME_BDF"
            fi
        done
    done

    # remove any leading whitespace
    M2_LIST=$(echo ${M2_LIST} | sed 's/^ *//')
    echo "${M2_LIST}"
}

c2_get_m2_by_model()
{
    NVME_DEV_LIST="nvme0n1 nvme1n1 nvme2n1 nvme3n1"
    M2_LIST="MZ1LB960HAJQ-00007 MZ1LW960HMJP-00003 SSDPEKKA010T7 SSDPEKKA010T8"
    U2_LIST="MZQL27T6HBLA-00A07"
    M2_DRIVES=""
    SYSFS_PATH="/sys/block"
    #
    # Find M2
    #
    for DEV in $(echo ${NVME_DEV_LIST}); do
        NVME_PATH="${SYSFS_PATH}/${DEV}/device/model"

        if [ -f ${NVME_PATH} ]; then
           NVME_MODEL=$(cat ${NVME_PATH})
           for MODEL in $(echo ${M2_LIST}); do
               if echo "${NVME_MODEL}" | grep -q "${MODEL}"; then
                    M2_DRIVES="${M2_DRIVES} /dev/${DEV}"
               fi
           done

        fi
    done

    M2_DRIVES=$(echo ${M2_DRIVES} | sed 's/^ *//')
    echo "${M2_DRIVES}"
}

dgx_xx00_get_M2_BDF()
{
    # Retrive the M.2 SLOT information and get the Bus Address
    PCI_BUS_LIST=`dmidecode --type 9 | grep -E 'Designation|Bus Address' | sed -n '/M2\|M\.2/{n; p}' | cut -d ' ' -f3`

    M2_LIST=""
    for DEV in `echo ${PCI_BUS_LIST}`; do
        BDF_LIST=`ls -al /sys/bus/pci/devices/$DEV/0* -d | cut -d '/' -f7`
        for NVME_BDF in `echo ${BDF_LIST}`; do
            DIR=/sys/bus/pci/devices/$DEV/$NVME_BDF/nvme
            if [ -d "$DIR" ]; then
                M2_LIST="$M2_LIST $NVME_BDF"
            fi
        done
    done

    # remove any leading whitespace
    M2_LIST=`echo ${M2_LIST} | sed 's/^ *//'`
    echo "${M2_LIST}"
}

bdf_to_nvme_map()
{
    BDF_LIST=$@
    SYSFS_PATH="/sys/bus/pci/devices"

    #
    # Find all M.2 drives attached to the BDF
    #
    for BDF in `echo ${BDF_LIST}`; do
        NVME_PATH="${SYSFS_PATH}/${BDF}/nvme*/nvme*/uevent"

        # uevent has the following fields: MAJOR, MINOR, DEVNAME
        # DEVNAME is the variable we are interested in. It points
        # nvme controller instance. Grab the controller instance
        # and append namespace "n1" to it.
        if [ -f ${NVME_PATH} ]; then
            NVME_NAME=`cat ${NVME_PATH} | grep DEVNAME | cut -d= -f2`
            BOOTDEVS="${BOOTDEVS} /dev/${NVME_NAME}n1"
        fi
    done

    BOOTDEVS=`echo ${BOOTDEVS} | sed 's/^ *//'`
    echo "${BOOTDEVS}"
}

dmi_type9_get_m2()
{
    # 1) Look for lines containing "Designation" or "Bus Address", leaving us with:
    #   Bus Address: 0000:c2:00.0
    #   Designation: PCIE7
    #   Bus Address: 0000:81:00.0
    #   Designation: OCU1
    #   Designation: OCU2
    #   Designation: M2_1
    #   Designation: M2_2
    #   Bus Address: 0000:41:00.0
    #   Designation: M2_3
    #   Bus Address: 0000:42:06.0
    #
    # 2) Look for lines containing "M2" or "M.2"; if found, also include the
    #    following line.  This leaves us with:
    #   Designation: M2_1
    #   Designation: M2_2
    #   Bus Address: 0000:41:00.0
    #   Designation: M2_3
    #   Bus Address: 0000:42:06.0
    #
    # 3) Remove lines containing "M2" or "M.2" so all we have left are the
    #    "Bus Address" lines for M.2 drives:
    #   Bus Address: 0000:41:00.0
    #   Bus Address: 0000:42:06.0
    DMIOUT=$(dmidecode --type 9 | grep -E 'Designation|Bus Address' | grep -A1 'M2\|M\.2' | grep -v 'M2\|M\.2')

    # Finally, remove the "Bus Address:" string.  The final awk is to
    # strip leading/trailing spaces, and squeeze all white-space into
    # a single space:
    #
    # 0000:41:00.0 0000:42:06.0
    #
    echo ${DMIOUT} | sed 's/Bus Address://g' | awk '{$1=$1};1'
}

dgxstation_a100_boot_drives()
{
    BOOT_DRIVES=""
    M2_BDF_LIST=$(dmi_type9_get_m2)
    SYSFS_PATH="/sys/bus/pci/devices"

    #
    # Find all M.2 drives attached to the BDF
    #
    for BDF in ${M2_BDF_LIST}; do
        NVME_PATH="${SYSFS_PATH}/${BDF}/nvme*/nvme*/uevent"

        # uevent has the following fields: MAJOR, MINOR, DEVNAME
        # DEVNAME is the variable we are interested in. It points
        # nvme controller instance. Grab the controller instance
        # and append namespace "n1" to it.
        if [ -f ${NVME_PATH} ]; then
            NVME_NAME=`cat ${NVME_PATH} | grep DEVNAME | cut -d= -f2`
            BOOT_DRIVES="${BOOT_DRIVES} /dev/${NVME_NAME}n1"
        fi
    done

    echo ${BOOT_DRIVES} | sed -e 's/^[ ]*//'
}

dgxstation_a800_boot_drives()
{
    dgxstation_a100_boot_drives
}

dgxstation_boot_drives()
{
    my_device=$(get_nvme_from_list "${1}" 1)

    if [ -z "${my_device}" ]; then
        my_device=$(get_sd_from_list "${1}" 1)
    fi

    echo ${my_device}
}

kvm_boot_drives()
{
    get_first_from_list "${1}"
}

xen_boot_drives()
{
    get_first_from_list "${1}"
}

vbox_boot_drives()
{
    my_device=$(get_nvme_from_list "${1}" 2)

    if [ -z "${my_device}" ]; then
        my_device=$(get_sd_from_list "${1}" 2)
    fi

    echo ${my_device}
}

other_boot_drives()
{
    get_first_from_list "${1}"
}

l4t_cg4_boot_drives()
{
    other_boot_drives "${1}"
}

l4t_c2_boot_drives()
{
    other_boot_drives "${1}"
}

l4t_cg1_boot_drives()
{
    other_boot_drives "${1}"
}

dgx_spark_boot_drives()
{
    other_boot_drives "${1}"
}

output_results_json()
{
    BOOT_DRIVES_LIST=$(list_to_json "${BOOT_DRIVES}")
    echo "{"
    echo "  \"boot_drives\": ${BOOT_DRIVES_LIST}"
    echo "}"
}

output_results()
{
    do_json=${1}

    if [ ${do_json} -eq 1 ]; then
        output_results_json
    else
        BOOT_DRIVES_LIST=$(list_to_human_readable "${BOOT_DRIVES}")
        echo "boot_drives:${BOOT_DRIVES_LIST}"
    fi
}

get_nvme_from_list()
{
    device_list=${1}
    num_devs=${2}
    ret_dev=""
    count=0

    for dev in ${device_list}; do
        if echo ${dev} | grep -q '/dev/nvme'; then
            ret_dev="${ret_dev} ${dev}"
            count=$((count + 1))

            if [ $count == ${num_devs} ]; then
                break
            fi
        fi
    done

    echo ${ret_dev}
}

get_sd_from_list()
{
    device_list=${1}
    num_devs=${2}
    ret_dev=""
    count=0

    for dev in ${device_list}; do
        if echo ${dev} | grep -q '/dev/sd'; then
            ret_dev="${ret_dev} ${dev}"
            count=$((count + 1))

            if [ $count == ${num_devs} ]; then
                break
            fi
        fi
    done

    echo ${ret_dev}
}

get_first_from_list()
{
    device_list=${1}
    echo ${device_list} | cut -d' ' -f1
}

##### MAIN #####
source ${MY_PATH}/general_funcs.bash
source ${MY_PATH}/device_funcs.bash
source ${MY_PATH}/plat_funcs.bash
must_run_as_root
PRODUCT_NAME=$(get_system_product_name)
PLAT_SHORT="other"
ALL_DISKS="unknown"

arg_has_json "$@"
if [ $? -eq 0  ]; then
    DO_JSON=1
fi

ALL_DISKS=$(list_devices disk | sort | tr '\n' ' ')
PLAT_SHORT=$(get_platform_short "${PRODUCT_NAME}")
if [[ -n "$PLAT_SHORT" ]]; then
    BOOT_DRIVES=$(${PLAT_SHORT}_boot_drives "${ALL_DISKS}")
else
    BOOT_DRIVES=$(other_boot_drives "${ALL_DISKS}")
fi

output_results ${DO_JSON}

exit 0