#!/bin/bash

PARAM=$#
FBARG=""
FBRESOLUTION=""
GRUBDEFAULT="/etc/default/grub"
UPDATEGRUBCMD="/usr/sbin/update-grub"
MKINITCPIOCONF="/etc/mkinitcpio.conf"



check_system() {
	local missing_file=""

	[ ! -e "${GRUBDEFAULT}" ] && missing_file="${GRUBDEFAULT}"
	[ ! -e "${UPDATEGRUBCMD}" ] && missing_file="${UPDATEGRUBCMD}"
	[ ! -e "${MKINITCPIOCONF}" ] && missing_file="${MKINITCPIOCONF}"

	if [ "${missing_file}" != "" ]; then
		echo "warning: '${missing_file}' does not exists! Abording..."
		echo ""
		echo "note: mhwd-fb couldn't find grub installed on your system. mhwd-fb sets kernel arguments to tell the mhwd-fb initramfs hook which framebuffer should be loaded. Check the wiki how to add this kernel arguments manually."
		exit 1
	fi

	# Check if mhwd-fb hook is enabled
	. "${MKINITCPIOCONF}"
	local found="";

	for arg in ${HOOKS}
	do
		if [ "${arg}" == "mhwd-fb" ]; then
			found="true"
			break
		fi 
	done

	if [ "${found}" != "true" ]; then
		echo "error: mhwd-fb hook isn't enabled! Abording..."
		exit 1
	fi
}



get_framebuffer_resolution() {
	# Find fitting resolution if not set
	[ "${FBRESOLUTION}" == "" ] && FBRESOLUTION="$(LANG=C LC_MESSAGES=C xrandr | grep \* | cut -d' ' -f4 | head -n 1)-32"
	[ "${FBRESOLUTION}" == "" ] && FBRESOLUTION="$(LANG=C LC_MESSAGES=C hwinfo --framebuffer | grep -w 'Mode' | cut -d: -f2 | sed 's/^ *//g' | sed 's/ .*//g' | sort -nr | head -n 1)-32"

	# Check if resolution is valid
	if ! [[ "$(echo "${FBRESOLUTION}" | sed 's/-//' | sed 's/x//')" =~ ^[0-9]+$ ]] || [[ "${FBRESOLUTION}" != *[0-9]"x"[0-9]* ]]; then
		echo "error: '${FBRESOLUTION}' isn't a valid framebuffer resolution!"
		exit 1
	fi
}



# param 1: "kms", "uvesafb" or "none"
set_framebuffer() {
	echo ":: setup mhwd-fb"

	. "${GRUBDEFAULT}"
	local NEW_GRUB_CMDLINE_LINUX=""

	for arg in ${GRUB_CMDLINE_LINUX}
	do
		if [[ "${arg}" != "mhwdfb="* ]]; then
			NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}${arg} "
		fi
	done

	if [ "$1" == "kms" ]; then
		NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=kms"
		echo "  using kms framebuffer"

		# Enable GRUB_GFXPAYLOAD_LINUX
		sed -i '/#[]*GRUB_GFXPAYLOAD_LINUX[ ]*=.*/ s/^.//' "${GRUBDEFAULT}"
	elif [ "$1" == "uvesafb" ]; then
		get_framebuffer_resolution

		if [ ${FBRESOLUTION} == "" ]; then
			NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=uvesafb"
			echo "  using uvesafb framebuffer"
		else
			NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=uvesafb:${FBRESOLUTION}"
			echo "  using uvesafb framebuffer with resolution: ${FBRESOLUTION}"
		fi

		# Disable GRUB_GFXPAYLOAD_LINUX
		sed -i "/#/!s|GRUB_GFXPAYLOAD_LINUX[ ]*=.*|#&|g" "${GRUBDEFAULT}"
	elif [ "$1" == "none" ]; then
		echo "  disabling mhwd-fb framebuffer option"

		# Enable GRUB_GFXPAYLOAD_LINUX
		sed -i '/#[]*GRUB_GFXPAYLOAD_LINUX[ ]*=.*/ s/^.//' "${GRUBDEFAULT}"
	else
		echo "error: invalid function argument!"
		exit 1
	fi
	
	# Adjust grub default file
	sed -i "/#/!s|GRUB_CMDLINE_LINUX[ ]*=.*|GRUB_CMDLINE_LINUX=\"${NEW_GRUB_CMDLINE_LINUX}\"|g" "${GRUBDEFAULT}"

	# Update grub configuration file
	echo ""
	echo ":: update-grub output:"
	${UPDATEGRUBCMD}
	if [ $? != 0 ]; then
		echo "error: grub-update failed!"
		exit 1
	fi
}



print_status() {
	. "${GRUBDEFAULT}"
	local FB_DRIVER=""
	local FB_RESOLUTION=""

	for arg in ${GRUB_CMDLINE_LINUX}
	do
		if [[ "${arg}" == "mhwdfb="* ]]; then
			arg="$(echo "${arg}" | cut -d= -f2)"
			FB_DRIVER="$(echo "${arg}" | cut -d: -f1)"
			FB_RESOLUTION="$(echo "${arg}:" | cut -d: -f2)" # additional ':' if f2 isn't set -> empty string
			break;
		fi
	done

	if [ "${FB_DRIVER}" != "kms" ] && [ "${FB_DRIVER}" != "uvesafb" ] && [ "${FB_DRIVER}" != "" ]; then
		echo "warning: set mhwd-fb driver is invalid: '${FB_DRIVER}'"
		exit 1
	elif [ "${FB_DRIVER}" == "" ]; then
		FB_DRIVER="disabled"
		FB_RESOLUTION="-"
	else
		[ "${FB_RESOLUTION}" == "" ] && FB_RESOLUTION="auto"
	fi

	echo ":: current mhwd-fb state"
	echo -e "  framebuffer:\t${FB_DRIVER}"
	echo -e "  resolution:\t${FB_RESOLUTION}"
}



print_help() {
	echo "mhwd-fb [OPTION] [...]"
	echo ""
	echo "   --help                          show help"
	echo "   --status                        show current status"
	echo "   --setfb [kms/uvesafb/none]      set framebuffer"
	echo "   --setres [resolution]           set uvesafb resolution. ex.: 1440x900 or 1440x900-32"
	echo ""
	echo "note: mhwd-fb sets kernel arguments to tell the mhwd-fb initramfs hook which framebuffer should be loaded. mhwd-fb requires an installed grub which is configured to this system to work. Otherwise check the wiki how to add mhwd-fb kernel arguments manually."
	echo ""
}

if [ "${PARAM}" -lt 1 ]; then
	print_status
	exit 0
fi

for (( I=1; $I <= $PARAM; I++ ))
do
	case "$1" in
		--help)
			print_help
			exit 0
			;;
		--setfb)
			shift
			FBARG="$1"
			;;
		--setres)
			shift
			FBRESOLUTION="$1"
			;;
		--status)
			print_status
			exit 0
			;;
		"") ;;
		*)
			echo "error: invalid argument: $1"
			echo ""
			print_help
			exit 1
			;;
	esac

	shift
done


# Check root
if [[ $EUID -ne 0 ]]; then
   echo "error: you cannot perform this operation unless you are root." 1>&2
   exit 1
fi

# Check system
check_system

# Set framebuffer
if [ "${FBARG}" != "" ]; then
	if [ "${FBARG}" != "kms" ] && [ "${FBARG}" != "uvesafb" ] && [ "${FBARG}" != "none" ]; then
		echo "error: invalid argument '${FBARG}'"
		exit 1
	fi

	set_framebuffer "${FBARG}"
fi

exit 0