#!/bin/bash

set -e

get_BDF()
{
	LSPCI_OUT=$1
	BUS=$(echo ${LSPCI_OUT} | cut -c 1,2 | head -n 1)
	DEV=$(echo ${LSPCI_OUT} | cut -c 4,5 | head -n 1)
	FUN=$(echo ${LSPCI_OUT} | cut -c 7,8 | head -n 1)

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

	echo $BDF
}

err_and_exit_zero()
{
	echo $@ && exit 1
}

edit_xconf()
{
	BDF=$1
	DRIVER=$2
	VENDOR_NAME=$3
	XORG_FILE=$4
	XORG_FILE_FULL="/etc/X11/xorg.conf.d/${XORG_FILE}"

	#edit xorg.conf.
	# <DO NOT try to format below lines.>
	sed -i "/Driver         \"nvidia\"/c\    Driver         \"${DRIVER}\"" \
		${XORG_FILE_FULL} || err_and_exit_zero "sed failed"

	# If file contains BusID, replace it
	if grep -q 'BusID          \"PCI:.*\"' ${XORG_FILE_FULL}; then
		sed -i "/BusID          \"PCI:.*\"/c \    BusID          \"PCI:${BDF}\"" \
			${XORG_FILE_FULL} || err_and_exit_zero "sed failed"
	else
		# Otherwise, append BusID
		sed -i "/Driver         \"${DRIVER}\"/a \    BusID          \"PCI:${BDF}\"" \
			${XORG_FILE_FULL} || err_and_exit_zero "sed failed"
	fi

	sed -i "/VendorName     \"NVIDIA Corporation\"/c\    VendorName     \"${VENDOR_NAME}\"" \
		${XORG_FILE_FULL} || err_and_exit_zero "sed failed"

	sed -i '/Section "ServerLayout"/i \
Section "ServerFlags" \
    Option         "AutoAddGPU" "off" \
EndSection\
' ${XORG_FILE_FULL} || err_and_exit_zero "sed failed"
#xorg.conf <DONE>

}


configure_aspeed()
{
	. /etc/os-release

	ASPEED_VGA=$1
	ASPEED_BDF=$(get_BDF "$ASPEED_VGA")
	DRIVER="ast"
	mv /etc/X11/xorg.conf /etc/X11/xorg.conf.d/xorg-aspeed.conf || err_and_exit_zero "mv failed"

	# For RHEL-like OSes, use modesetting
	if is_rhel_like_os; then
		DRIVER="modesetting"
	fi

	VENDOR_NAME="ASPEED Technology, Inc."
	XORG_FILE="xorg-aspeed.conf"

	edit_xconf "$ASPEED_BDF" "$DRIVER" "$VENDOR_NAME" "$XORG_FILE"
}

configure_nvidia()
{
	NVIDIA_VGA=$1
	NVIDIA_BDF=$(get_BDF "$NVIDIA_VGA")
	mv /etc/X11/xorg.conf /etc/X11/xorg.conf.d/xorg-nvidia.conf || err_and_exit_zero "mv failed"
	DRIVER="nvidia"
	VENDOR_NAME="NVIDIA Corporation"
	XORG_FILE="xorg-nvidia.conf"

	edit_xconf "$NVIDIA_BDF" "$DRIVER" "$VENDOR_NAME" "$XORG_FILE"
}


# Use common APIs to figure out platform
plat_funcs="/usr/local/sbin/nv_scripts/plat_funcs.bash"
. ${plat_funcs}
general_funcs="/usr/local/sbin/nv_scripts/general_funcs.bash"
. ${general_funcs}

if plat_needs_initial_nvidia_xconfig; then
    if [[ ! -f /etc/X11/xorg.conf ]]; then
        /usr/bin/nvidia-xconfig --allow-empty-initial-configuration \
        || err_and_exit_zero "nvidia-xconfig failed"
    fi
fi

if plat_needs_adaptive_nvidia_xconfig; then
	mkdir -p /etc/X11/xorg.conf.d || err_and_exit_zero "mkdir failed"

	# Remove xorg-nvidia.conf or xorg-aspeed.conf if exists
	# We dont take backups as these files are dynamically generated everytime
	# nvidia-conf-xconfig.service starts.
	rm  /etc/X11/xorg.conf.d/xorg-nvidia.conf > /dev/null 2>&1 || true
	rm  /etc/X11/xorg.conf.d/xorg-aspeed.conf > /dev/null 2>&1 || true

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

	if [ -z "${NVIDIA_VGA}" ]  && [ -z "${ASPEED_VGA}" ]; then
		err_and_exit_zero "Neither Nvidia nor ASPEED onboard VGA compatible controller found. Exiting without setup.."
	fi

	# Retrieve BIOS VGA settings
	# e.g
	# 	$dmidecode -t 11 | grep "OnBrd/Ext VGA Select"
	#	String 1: OnBrd/Ext VGA Select = Auto
	BIOS_VGA=`dmidecode -t 11 | grep "OnBrd/Ext VGA Select" | awk '{print $NF}'`

	if [ "${BIOS_VGA}" = "Auto" ] ;then
		# if no NVIDIA VGA device found then fallback to ASPEED
		 if [ -z "${NVIDIA_VGA}" ]; then
			echo "Nvidia VGA compatible controller not found"
			echo "Configuring ASPEED onboard VGA controller for display"
			configure_aspeed "$ASPEED_VGA"
		else
			configure_nvidia "$NVIDIA_VGA"
		fi

	elif [ "${BIOS_VGA}" = "Onboard" ] ;then
		if [ -z "${ASPEED_VGA}" ]; then
			echo "ASPEED VGA compatible controller not found for some reason"
			echo "Configuring Nvidia VGA controller for display"
			configure_nvidia "$NVIDIA_VGA"
		else
			configure_aspeed "$ASPEED_VGA"
		fi

	else
		# This means BIOS_VGA="External"
		configure_nvidia "$NVIDIA_VGA"
	fi

	echo "New X configuration file written under /etc/X11/xorg.conf.d/"
	echo "If you restart this service, please reboot the system or restart the GNOME Display Manager (gdm3) service for these changes to take effect"
fi
exit 0