#!/bin/bash

# this script is based on NVIDIA's documentation and scripts

ENABLED_VALUE="0x00000107"
DISABLED_VALUE="0x00000000"

logger () {
    echo "[NVME interrupt coalescing] $*"
}

set_coalescing() {
    local device=${1}
    local value=${2}
    local vendor_id=${3}
    logger "Setting interrupt coalescing for ${device} (vendor: ${vendor_id}) to ${value}"

    if ! nvme set-feature "${device}" -f 8 --value "${value}";
    then
        logger "Error: Failed to set interrupt coalescing for ${device}"
        return 1
    fi

    # Verify the setting
    local current_value
    current_value=$(nvme get-feature "${device}" -f 8 | awk -F':' '{print $3}')
    if [ "${current_value}" != "${value}" ]; then
        logger "Error: Verification failed for ${device}. Expected: ${value}, Got: ${current_value}"
        return 1
    fi
    return 0
}

# Find all NVMe controller devices
for device in $(ls -d /dev/nvme[0-9]* 2>/dev/null | grep -E '^/dev/nvme[0-9]+$'); do
    device_name=$(basename "${device}")
    vendor_id=$(cat "/sys/class/nvme/${device_name}/device/vendor" 2>/dev/null)
    
    case "${vendor_id}" in "0x144d"|"0x1e0f"|"0x1344") # Samsung, Kioxia, or Micron
        set_coalescing "${device}" "${ENABLED_VALUE}" "${vendor_id}"
        ;;
    *)
        logger "Skipping unsupported device ${device_name} (vendor: ${vendor_id})"
        ;;
    esac
done
