#!/bin/bash

# Function to handle errors
handle_error() {
    echo "Error: $1"
    exit 1
}

# Run the mlxfwmanager command and store the output
output=$(mlxfwmanager 2>&1) || handle_error "$output"

# Filter the output to find the line containing the specific description
device_info=$(echo "$output" | grep -E -B 8 -A 8 "MT_0000000891|MT_0000000937" | grep "PCI Device Name" | cut -d ':' -f 2 | awk '{$1=$1;print}')

# Run mst status -v command and store the output
mst_status_output=$(mst status -v 2>&1) || handle_error "$mst_status_output"

command_output=""
start_recording=0

while IFS= read -r line; do
    # Check if the line starts with "PCI devices:"
    if [[ $line =~ ^PCI\ devices: ]]; then
        start_recording=1
        continue  # Skip to the next iteration
    fi

    # Start recording after the first occurrence of "PCI devices:"
    if [[ $start_recording -eq 1 && ! -z "$line" && $line != *"MST PCI"* && $line != *"MST modules"* ]]; then
        # Extracting the second and third columns from each line
        device_info=$(echo "$line" | awk '{print $2, $3}')
        # Append to command_output variable
        command_output+="$device_info"$'\n'
    fi
done <<< "$mst_status_output"

# Custom sort order
custom_order=(
    "dc:00.0"
    "9a:00.0"
    "ce:00.0"
    "c0:00.0"
    "5e:00.0"
    "18:00.0"
    "4f:00.0"
    "40:00.0"
)

# Array to hold device info
device_info=()

# Extract device info from command output
while read -r line; do
    device_info+=("$line")
done <<< "$command_output"

# Reorder device_info according to custom_order
reordered_device_info=()
for item in "${custom_order[@]}"; do
    for device in "${device_info[@]}"; do
        if [[ $device == *"$item"* ]]; then
            reordered_device_info+=("${device%% *}")
            break
        fi
    done
done

# Initialize previous serial number variable
prev_serial_number=""

# Loop over each device
for device in "${reordered_device_info[@]}"; do
    # Extract the device name from the line
    device_name=$device

    # Run mlxreg command for the current device and capture the output
    secondary_value=$(mlxreg -d "$device_name" --reg_name PMAOS --get -i module=0,slot_index=0 2>&1) || handle_error "$secondary_value"
    secondary_value=$(echo "$secondary_value" | grep secondary | awk '{print $3}')
    # Replace secondary value if it's 0x00000000 with primary
    if [ "$secondary_value" = "0x00000000" ]; then
        secondary_value="Primary\t"
    elif [ "$secondary_value" = "0x00000001" ]; then
        # Replace primary value if it's 0x00000001 with secondary
        secondary_value="Secondary"
    fi

    fw_version=$(flint -d "$device_name" q 2>&1) || handle_error "$fw_version"
    fw_version=$(echo "$fw_version" | grep -i product | awk '{print $3}')

    link_state=$(mlxlink -d "$device_name" 2>&1) || handle_error "$link_state"
    link_state=$(echo "$link_state" | grep -i -m 1 "State" | awk -F ': ' '{print $2}')

    link_speed=$(mlxlink -d "$device_name" 2>&1) || handle_error "$link_speed"
    link_speed=$(echo "$link_speed" | grep -i -m 1 "Speed" | awk -F ': ' '{print $2}')

    pcie_info=$(mlxlink -d "$device_name" --port_type PCIE 2>&1) || handle_error "$pcie_info"
    pcie_info=$(echo "$pcie_info" | grep -A 4 "Depth, pcie index, node            : 0, 0, 0")

    # Extracting and formatting PCIe info
    pcie_speed=$(echo "$pcie_info" | grep -oP 'Link Speed Active \(Enabled\) *: \K[^(]*')
    pcie_width=$(echo "$pcie_info" | grep -oP 'Link Width Active \(Enabled\) *: \K[^(]*' | sed 's/ *$//')
    formatted_pcie_info="$pcie_speed($pcie_width)"

    # Run mlxburn to get serial number
    mlxburn_output=$(mlxburn -d "$device_name" -vpd 2>&1) || handle_error "$mlxburn_output"
    serial_number=$(echo "$mlxburn_output" | grep -i "serial number" | awk '{print $NF}')

    flint_output=$(flint -d "$device_name" q full 2>&1) || handle_error "$flint_output"
    part_number=$(echo "$flint_output" | grep -i "Part Number" | awk '{print $NF}')

    # Check if the current serial number is different from the previous one
    if [ "$serial_number" != "$prev_serial_number" ]; then
        # Output the serial number if it's different from the previous one
        echo
        echo -e "PN:\t$part_number"
        echo -e "SN:\t$serial_number"
        # Output the titles if it's a new serial number
        echo -e "MST Device:\t\t\tConf:\t\tPCIe info:\t\tFW:\t\tLink State:\tLink Speed:"
        # Update the previous serial number to the current one
        prev_serial_number="$serial_number"
    fi

    # Output the device details
    if [[ "$link_speed" == *"N/A"* ]]; then
        echo -e "$device_name\t$secondary_value\t$formatted_pcie_info\t\t$fw_version\t$link_state\t$link_speed"
    else
        echo -e "$device_name\t$secondary_value\t$formatted_pcie_info\t\t$fw_version\t$link_state\t\t$link_speed"
    fi

done
