#!/usr/bin/env bash
# https://chatgpt.com/c/6828dc89-e524-800e-b978-05b3f27626c9

# IOMMU and Device Enumeration Script
# Description: Enumerates IOMMU groups and lists associated PCI and USB devices
#              in a neatly formatted, column-aligned table.
# Usage:       Run as root or with permissions to read /sys/kernel/iommu_groups,
#              and have lspci and lsusb available.
# Output:      Columns: IOMMU group, Device ID (vendor:device), Flags (e.g., reset support or USB),
#              Address or bus:device, Type (PCI/USB), Description.

# Enable nullglob so patterns that match no files expand to empty
shopt -s nullglob

# Array to collect all output rows
declare -a rows

# Function to queue a row for later formatted output
# queue_row() { rows+=("$1"); }
queue_row() { rows+=("$(printf '%s\t[%s]\t%s\t%s\t%s' "${1}" "${2}" "${3}" "${4}" "${5}")"); }

# Function queue_rowto print all queued rows with column alignment
# print_all_rows() { printf '%s\n' "${rows[@]}" | column -ts $'\t'; }

# Header
rows+=("$(printf '%s\t%s\t%s\t%s\t%s\t%s' "IOMMU"    "Device ID"   "Flags" "Address" "Type"         "Description")");
rows+=("$(printf '%s\t%s\t%s\t%s\t%s\t%s' "--------" "-----------" "-----" "-------" "------------" "--------------------------------------------")");

# Track the last processed IOMMU group ID to control header formatting
last_group_id=""

# Iterate over each IOMMU group directory, sorted numerically
for iommu_group_dir in $(find /sys/kernel/iommu_groups/* -maxdepth 0 -type d | sort -V); do
    # Iterate over each PCI device in this IOMMU group
    for device_dir in "${iommu_group_dir}/devices/"*; do
        # Extract the current group ID from the directory name
        current_group_id="${iommu_group_dir##*/}"

        # Determine group column header for first device in the group
        if [[ "$current_group_id" != "$last_group_id" ]]; then
            group_col="Group ${current_group_id}"
        else
            group_col="    L $current_group_id"
        fi
        last_group_id="$current_group_id"

        # Gather PCI information
        vendor_device_id=$(lspci -nms "${device_dir##*/}" | awk -F'"' '{printf "%s:%s", $4, $6}')
        reset_flag=""
        [[ -e "${device_dir}/reset" ]] && reset_flag="[R]"
        pci_details=$(lspci -mms "${device_dir##*/}" | awk -F'"' '{printf "%s\t%s\t%s", $1, $2, $6}')

        # Queue the PCI row (tab-separated)
        # rows+=("${group_col}	[${vendor_device_id}] ${reset_flag} ${pci_details}")
        # rows+=("$(printf '%s\t[%s]\t[%s]\t%s\t%s' "${group_col}" "${vendor_device_id}" "${reset_flag}" "${pci_details}")")
        queue_row "${group_col}" "${vendor_device_id}" "${reset_flag}" "${pci_details}"

        # Process any USB interfaces for this device
        for usb_dir in "${device_dir}/usb"*/; do
            bus_number=$(<"${usb_dir}/busnum")
            dev_number=$(<"${usb_dir}/devnum")
            group_col="    L $current_group_id"

            # Get the lsusb line and parse vendor:product ID and description
            usb_line=$(lsusb -s "${bus_number}:${dev_number}")
            usb_id=$(awk '{print $6}' <<< "$usb_line")
            usb_desc=${usb_line#* $usb_id }

            # Queue the USB row using printf to embed tabs correctly
            # rows+=("$(printf '\t[%s]\t[%s]\t%s\t%s' "$usb_id" "${bus_number}:${dev_number}" "U" "$usb_desc")")
            queue_row "$group_col" "$usb_id" "USB" "${bus_number}:${dev_number}" "$usb_desc"
        done
    done
done

# Print all collected rows with aligned columns
printf '%s\n' "${rows[@]}" | column -ts $'\t';