#!/bin/bash
#
# Script to generate Xorg configuration for NVIDIA GPU in DGX systems

set -eu

XORG_CONF="/etc/X11/xorg.conf.d/10-nvidia.conf"

get_BDF()
{
    LSPCI_OUT=$1

    # Check if domain is present (format: DDDD:BB:SS.F or BB:SS.F)
    # Count colons to determine format
    COLON_COUNT=$(echo "${LSPCI_OUT}" | tr -cd ':' | wc -c)

    if [ "$COLON_COUNT" -eq 2 ]; then
        # Format with domain: DDDD:BB:SS.F
        DOMAIN=$(echo "${LSPCI_OUT}" | cut -d':' -f1)
        BUS=$(echo "${LSPCI_OUT}" | cut -d':' -f2)
        DEV=$(echo "${LSPCI_OUT}" | cut -d':' -f3 | cut -d'.' -f1)
        FUN=$(echo "${LSPCI_OUT}" | cut -d'.' -f2)
    else
        # Format without domain: BB:SS.F
        DOMAIN="0000"
        BUS=$(echo "${LSPCI_OUT}" | cut -d':' -f1)
        DEV=$(echo "${LSPCI_OUT}" | cut -d':' -f2 | cut -d'.' -f1)
        FUN=$(echo "${LSPCI_OUT}" | cut -d'.' -f2)
    fi

    # Convert BDF to decimal because xorg uses decimal
    DOMAIN=$((0x$DOMAIN))
    BUS=$((0x$BUS))
    DEV=$((0x$DEV))
    FUN=$((0x$FUN))

    # xorg.conf format: PCI:bus@domain:device:function
    BDF="PCI:$BUS@$DOMAIN:$DEV:$FUN"

    echo $BDF
}

NVIDIA_VGA=$(lspci | grep "VGA compatible controller: NVIDIA Corporation" | cut -d' ' -f1)

if [ -z "$NVIDIA_VGA" ]; then
    echo "No NVIDIA VGA controller found. Removing $XORG_CONF if it exists."
    rm -f "$XORG_CONF"
    exit 0
fi

NVIDIA_BDF=$(get_BDF "$NVIDIA_VGA")

rm -f "$XORG_CONF"*
mkdir -p "$(dirname "$XORG_CONF")"
nvidia-xconfig \
	--no-probe-all-gpus \
	--allow-empty-initial-configuration \
	--busid="$NVIDIA_BDF" \
	--output-xconfig="$XORG_CONF"

exit 0
