#!/bin/bash
#
# rdma-detect
#
# A script to detect RDMA/InfiniBand capable devices on a system
# and report their status, vendor, link state, speed, driver, and firmware.

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# PCI IDs database file
PCIIDS_FILE="/usr/share/misc/pci.ids"

# Output file (can be overridden by environment variable)
OUTPUT_FILE="${RDMA_DETECT_OUTPUT_FILE:-/tmp/rdma-detect.status}"

# Check if running as root or with sufficient permissions
if [ ! -r /sys/class/infiniband ] && [ ! -d /sys/class/infiniband ]; then
    echo "Warning: Cannot access /sys/class/infiniband. Some information may be unavailable."
fi

# Function to ensure PCI IDs database is available and updated
ensure_pciids() {
    # Check if update-pciids command exists
    if command -v update-pciids &> /dev/null; then
        if [ ! -f "$PCIIDS_FILE" ] || [ $(find "$PCIIDS_FILE" -mtime +30 2>/dev/null | wc -l) -gt 0 ]; then
            echo "Updating PCI IDs database..."
            if [ "$EUID" -eq 0 ]; then
                update-pciids -q 2>/dev/null || echo "Warning: Failed to update PCI IDs database"
            else
                echo "Note: Run as root to update PCI IDs database"
            fi
        fi
    else
        echo -e "${YELLOW}Warning: 'update-pciids' command not found.${NC}"
        echo "Please install 'pciutils' package for vendor identification."
        echo "  Debian/Ubuntu: apt-get install pciutils"
        echo "  RHEL/CentOS:   yum install pciutils"
        echo ""
    fi
}

# Function to get vendor name from vendor ID using lspci or pci.ids
get_vendor_name() {
    local VENDOR_ID="$1"
    local VENDOR_NAME=""
    
    # Remove 0x prefix if present
    VENDOR_ID="${VENDOR_ID#0x}"
    
    # Try using lspci if available
    if command -v lspci &> /dev/null; then
        VENDOR_NAME=$(lspci -d "${VENDOR_ID}:" -vm 2>/dev/null | grep "^Vendor:" | cut -d: -f2- | sed 's/^[[:space:]]*//' | head -1)
    fi
    
    # If lspci didn't work, try parsing pci.ids file directly
    if [ -z "$VENDOR_NAME" ] && [ -f "$PCIIDS_FILE" ]; then
        VENDOR_NAME=$(grep "^${VENDOR_ID}" "$PCIIDS_FILE" | head -1 | cut -f2- | sed 's/^[[:space:]]*//')
    fi
    
    # Return vendor name or unknown
    if [ -n "$VENDOR_NAME" ]; then
        echo "$VENDOR_NAME"
    else
        echo "Unknown (${VENDOR_ID})"
    fi
}

# Function to get driver information
get_driver_info() {
    local DEV_PATH="$1"
    local DRIVER_NAME=""
    local DRIVER_VERSION=""
    
    # Get driver name from sysfs
    if [ -L "${DEV_PATH}/device/driver" ]; then
        DRIVER_NAME=$(basename $(readlink -f "${DEV_PATH}/device/driver") 2>/dev/null || echo "")
    fi
    
    # Try to get driver version from modinfo
    if [ -n "$DRIVER_NAME" ] && command -v modinfo &> /dev/null; then
        DRIVER_VERSION=$(modinfo "$DRIVER_NAME" 2>/dev/null | grep "^version:" | awk '{print $2}' | head -1)
        if [ -z "$DRIVER_VERSION" ]; then
            # Try srcversion as fallback
            DRIVER_VERSION=$(modinfo "$DRIVER_NAME" 2>/dev/null | grep "^srcversion:" | awk '{print $2}' | head -1)
        fi
    fi
    
    echo "$DRIVER_NAME|$DRIVER_VERSION"
}

# Function to output both to console and file
output() {
    echo -e "$1" | tee -a "$OUTPUT_FILE"
}

output_plain() {
    echo "$1" | tee -a "$OUTPUT_FILE"
}

# Initialize output file
echo "RDMA/InfiniBand Device Detection Report" > "$OUTPUT_FILE"
echo "Generated: $(date)" >> "$OUTPUT_FILE"
echo "Hostname: $(hostname)" >> "$OUTPUT_FILE"
echo "========================================" >> "$OUTPUT_FILE"
echo "" >> "$OUTPUT_FILE"

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}RDMA/InfiniBand Device Detection${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# Check if RDMA tools are available
if ! command -v rdma &> /dev/null; then
    ERROR_MSG="Error: 'rdma' command not found. Please install iproute2 or rdma-core package."
    echo -e "${RED}${ERROR_MSG}${NC}"
    echo "$ERROR_MSG" >> "$OUTPUT_FILE"
    exit 1
fi

# Ensure PCI IDs database is available
ensure_pciids

# Get list of InfiniBand devices
IB_DEVICES=$(ls /sys/class/infiniband 2>/dev/null || echo "")

if [ -z "$IB_DEVICES" ]; then
    WARN_MSG="No RDMA/InfiniBand devices found on this system."
    echo -e "${YELLOW}${WARN_MSG}${NC}"
    echo "$WARN_MSG" >> "$OUTPUT_FILE"
    exit 0
fi

# Counter for devices found
DEVICE_COUNT=0

# Iterate through each device
for DEVICE in $IB_DEVICES; do
    DEVICE_COUNT=$((DEVICE_COUNT + 1))
    output "${GREEN}Device $DEVICE_COUNT: ${DEVICE}${NC}"
    output_plain "----------------------------------------"
    
    # Get device path
    DEV_PATH="/sys/class/infiniband/${DEVICE}"
    
    # Get vendor and device ID
    if [ -f "${DEV_PATH}/device/vendor" ]; then
        VENDOR_ID=$(cat "${DEV_PATH}/device/vendor" 2>/dev/null || echo "N/A")
        VENDOR_NAME=$(get_vendor_name "$VENDOR_ID")
        output_plain "  Vendor:        ${VENDOR_NAME}"
        output_plain "  Vendor ID:     ${VENDOR_ID}"
    fi
    
    if [ -f "${DEV_PATH}/device/device" ]; then
        DEVICE_ID=$(cat "${DEV_PATH}/device/device" 2>/dev/null || echo "N/A")
        output_plain "  Device ID:     ${DEVICE_ID}"
    fi
    
    # Get PCI address
    if [ -L "${DEV_PATH}/device" ]; then
        PCI_ADDRESS=$(basename $(readlink -f "${DEV_PATH}/device") 2>/dev/null || echo "N/A")
        output_plain "  PCI Address:   ${PCI_ADDRESS}"
        
        # Get full device description using lspci
        if command -v lspci &> /dev/null && [ "$PCI_ADDRESS" != "N/A" ]; then
            DEVICE_DESC=$(lspci -s "$PCI_ADDRESS" 2>/dev/null | cut -d: -f3- | sed 's/^[[:space:]]*//' || echo "")
            if [ -n "$DEVICE_DESC" ]; then
                output_plain "  Description:   ${DEVICE_DESC}"
            fi
        fi
    fi
    
    # Get driver information
    DRIVER_INFO=$(get_driver_info "$DEV_PATH")
    DRIVER_NAME=$(echo "$DRIVER_INFO" | cut -d'|' -f1)
    DRIVER_VERSION=$(echo "$DRIVER_INFO" | cut -d'|' -f2)
    
    if [ -n "$DRIVER_NAME" ]; then
        output_plain "  Driver:        ${DRIVER_NAME}"
        if [ -n "$DRIVER_VERSION" ]; then
            output_plain "  Driver Ver:    ${DRIVER_VERSION}"
        fi
    fi
    
    # Get firmware version if available
    if [ -f "${DEV_PATH}/fw_ver" ]; then
        FW_VER=$(cat "${DEV_PATH}/fw_ver" 2>/dev/null || echo "N/A")
        output_plain "  Firmware:      ${FW_VER}"
    fi
    
    # Get board ID if available
    if [ -f "${DEV_PATH}/board_id" ]; then
        BOARD_ID=$(cat "${DEV_PATH}/board_id" 2>/dev/null || echo "N/A")
        output_plain "  Board ID:      ${BOARD_ID}"
    fi
    
    # Get number of ports
    if [ -f "${DEV_PATH}/node_type" ]; then
        NODE_TYPE=$(cat "${DEV_PATH}/node_type" 2>/dev/null || echo "N/A")
        output_plain "  Node Type:     ${NODE_TYPE}"
    fi
    
    # Check each port on the device
    PORT_DIRS=$(ls -d ${DEV_PATH}/ports/* 2>/dev/null || echo "")
    
    if [ -n "$PORT_DIRS" ]; then
        for PORT_PATH in $PORT_DIRS; do
            PORT_NUM=$(basename "$PORT_PATH")
            output_plain ""
            output_plain "  Port ${PORT_NUM}:"
            
            # Get port state
            if [ -f "${PORT_PATH}/state" ]; then
                STATE=$(cat "${PORT_PATH}/state" 2>/dev/null || echo "N/A")
                STATE_NUM=$(echo "$STATE" | awk '{print $1}')
                STATE_NAME=$(echo "$STATE" | awk '{print $2}' | tr -d '()')
                
                if [ "$STATE_NUM" = "4" ] || [ "$STATE_NAME" = "ACTIVE" ]; then
                    output "    Link State:  ${GREEN}${STATE_NAME} (${STATE_NUM})${NC}"
                else
                    output "    Link State:  ${RED}${STATE_NAME} (${STATE_NUM})${NC}"
                fi
            fi
            
            # Get physical state
            if [ -f "${PORT_PATH}/phys_state" ]; then
                PHYS_STATE=$(cat "${PORT_PATH}/phys_state" 2>/dev/null || echo "N/A")
                output_plain "    Phys State:  ${PHYS_STATE}"
            fi
            
            # Get link rate
            if [ -f "${PORT_PATH}/rate" ]; then
                RATE=$(cat "${PORT_PATH}/rate" 2>/dev/null | awk '{print $1}')
                RATE_UNIT=$(cat "${PORT_PATH}/rate" 2>/dev/null | awk '{print $2}')
                output_plain "    Link Speed:  ${RATE} ${RATE_UNIT}"
            fi
            
            # Get link width
            if [ -f "${PORT_PATH}/lid" ]; then
                LID=$(cat "${PORT_PATH}/lid" 2>/dev/null || echo "N/A")
                output_plain "    LID:         ${LID}"
            fi
            
            # Try to find the network interface name using rdma link
            NETDEV=$(rdma link show 2>/dev/null | grep -w "${DEVICE}/${PORT_NUM}" | awk '{print $8}' | head -1 || echo "")
            if [ -n "$NETDEV" ]; then
                output_plain "    Net Device:  ${NETDEV}"
                
                # Get additional info from ip link if network device exists
                if command -v ip &> /dev/null && ip link show "$NETDEV" &> /dev/null; then
                    MAC=$(ip link show "$NETDEV" | grep -oP 'link/\w+ \K[^ ]+' | head -1)
                    if [ -n "$MAC" ]; then
                        output_plain "    MAC Address: ${MAC}"
                    fi
                    
                    # Check if interface is up
                    if ip link show "$NETDEV" | grep -q "state UP"; then
                        output "    Net Status:  ${GREEN}UP${NC}"
                    else
                        output "    Net Status:  ${YELLOW}DOWN${NC}"
                    fi
                fi
            else
                # Try ibdev2netdev if available
                if command -v ibdev2netdev &> /dev/null; then
                    NETDEV=$(ibdev2netdev | grep "^${DEVICE} port ${PORT_NUM}" | awk '{print $5}' | tr -d '()' || echo "")
                    if [ -n "$NETDEV" ] && [ "$NETDEV" != "Down" ]; then
                        output_plain "    Net Device:  ${NETDEV}"
                    fi
                fi
            fi
        done
    fi
    
    output_plain ""
done

output "${BLUE}========================================${NC}"
output "${GREEN}Total RDMA/InfiniBand devices found: ${DEVICE_COUNT}${NC}"
output "${BLUE}========================================${NC}"

# Additional summary using rdma link command
output_plain ""
output_plain "RDMA Link Summary:"
if rdma link show &> /dev/null; then
    rdma link show 2>/dev/null | while read -r LINE; do
        echo "  $LINE" | tee -a "$OUTPUT_FILE"
    done
else
    output_plain "  Unable to retrieve rdma link information"
fi

output_plain ""
output_plain "========================================="
output_plain "Report saved to: ${OUTPUT_FILE}"
output_plain "========================================="

echo ""
echo -e "${GREEN}Report saved to: ${OUTPUT_FILE}${NC}"

exit 0
