#!/usr/bin/env bash

die() {
	echo -e " ${NOCOLOR-\e[1;31m*\e[0m }${*}" >&2
	exit 1
}

einfo() {
	echo -e " ${NOCOLOR-\e[1;32m*\e[0m }${*}" >&2
}

ewarn() {
	echo -e " ${NOCOLOR-\e[1;33m*\e[0m }${*}" >&2
}

add_uefi_entry () (
	# Create label for UEFI entry
	local entry_label
	entry_label="${efi_file_path/.uefibackup} on $partition_label "

	# Add prefix to entry label for easier identification
	if [[ -n $backup_efi ]]; then
		entry_label="UMCB $entry_label"
	else
		entry_label="UMC $entry_label"
	fi

	# Create path to initramfs
	local initramfs_image
	if [[ -z $kernel_version ]]; then
		# Use versionless initramfs image if the kernel efi file is versionless
		initramfs_image="${efi_file_path/${efi_file_path##*/}}initramfs.img"
	else
		initramfs_image="${efi_file_path/${efi_file_path##*/}}initramfs-$kernel_version.img"
	fi

	# Add .old suffix to initramfs in case we are handling kernel with -old suffix
	[[ "$efi_file_path" == *"-old."* ]] && initramfs_image="${initramfs_image/-old}.old"

	# Strip .uefibackup from efi file path
	[[ "$efi_file_path" == *".uefibackup" ]] && efi_file_path=${efi_file_path//.uefibackup/}

	# If shim is present in directory, presume it's used for every kernel in said directory
	local shim
	local adding_kernel_commands
	shim="$(find "$partition_mount""${efi_file_path/${efi_file_path##*/}/}" -maxdepth 1 -iname "*shim*.efi")"
	if [[ -n "$shim" ]]; then
		shim="${shim%%.efi*}.efi"
		adding_kernel_commands="${efi_file_path//\//\\} ${kernel_commands}"
		if [[ -n $backup_efi ]]; then
			einfo "Creating BACKUP UEFI entry \"$bootnum\" for \"$partition_mount$efi_file_path\" using shim \"$shim\" found on \"$partition\""
		else
			einfo "Creating UEFI entry \"$bootnum\" for \"$partition_mount$efi_file_path\" using shim \"$shim\" found on \"$partition\""
		fi
		efi_file_path="${efi_file_path/${efi_file_path##*/}/}${shim/*\//}"
	else
		adding_kernel_commands="${kernel_commands}"
		if [[ -n $backup_efi ]]; then
			einfo "Creating BACKUP UEFI entry \"$bootnum\" for \"$partition_mount$efi_file_path\" found on \"$partition\""
		else
			einfo "Creating UEFI entry \"$bootnum\" for \"$partition_mount$efi_file_path\" found on \"$partition\""
		fi
	fi

	# Check if microcode image exists
	## microcode image can be ignored the same way as kernel image via .ignore suffix
	local microcode_path
	microcode_path="${efi_file_path/${efi_file_path##*/}}"
	if [[ -f "${partition_mount}${microcode_path}amd-uc.img" ]] && [[ ! -f "${partition_mount}${microcode_path}amd-uc.img.ignore" ]]; then
		adding_kernel_commands="${adding_kernel_commands} initrd=${microcode_path//\//\\}amd-uc.img"
	fi
	if [[ -f "${partition_mount}${microcode_path}intel-uc.img" ]] && [[ ! -f "${partition_mount}${microcode_path}intel-uc.img.ignore" ]]; then
		adding_kernel_commands="${adding_kernel_commands} initrd=${microcode_path//\//\\}intel-uc.img"
	fi

	if  [[ -f "$partition_mount$initramfs_image.ignore" ]]; then
	    # Ignore chosen initramfs image
	    ewarn "Ignoring initramfs image \"$partition_mount$initramfs_image\" of \"$partition_mount$efi_file_path\""
	elif [[ -f "$partition_mount$initramfs_image" ]]; then
	    # Check if corresponding initramfs exists
		adding_kernel_commands="${adding_kernel_commands} initrd=${initramfs_image//\//\\}"
	else
		ewarn "No initramfs found for \"$partition_mount$efi_file_path\"."
	fi

	# Find disk name and gpt partition id
	for line in $(lsblk -lno NAME); do
        	if [[ "$partition" == *"$line"* ]] && [[ "$partition" != "$line" ]]; then
                	local disk
                	local partition_id
                	disk="$line"
                	partition_id="${partition//$line}"
                	partition_id="${partition_id//p}"
        	fi
	done

	# Add new entry
	if [[ -n $backup_efi ]]; then
		# When creating backup entry, don't add it to the bootorder
		efibootmgr -q --create-only -b "$bootnum" --disk /dev/"$disk" --part "$partition_id" --label "$entry_label" --loader "${efi_file_path//\//\\}"\
			-u "$adding_kernel_commands" || die "Failed to add BACKUP UEFI entry for \"$efi_file_path\""
	else
		efibootmgr -q --create -b "$bootnum" --disk /dev/"$disk" --part "$partition_id" --label "$entry_label" --loader "${efi_file_path//\//\\}"\
			-u "$adding_kernel_commands" || die "Failed to add UEFI entry for \"$efi_file_path\""
	fi
)

clean_legacy_entries () {
	# Run only if no post 2.0 entries exist
	if [[ "$initial_uefi_state" != *" UMC"* ]]; then
		local IFS=$'\n'
		local entries
		local entry
		entries="$(grep -v -e "BootOrder" -e "Timeout" -e "BootCurrent" <<< "$initial_uefi_state" | grep "$kernel_images")"
		for entry in ${entries}; do
			if [[ "$(printf %d "0x${entry:4:4}" 2> /dev/null)" -gt 255 ]] && [[ "$(printf %d "0x${entry:4:4}" 2> /dev/null)" -lt 513 ]]; then
				efibootmgr -q -B -b "${entry:4:4}"
				ewarn "Removing pre v2.0 legacy entry \"${entry:4:4}\""
			fi
		done
	fi
}

wipe_entries () {
	local initial_uefi_umc_state

	# Get list of all non backup entries made by uefi-mkconfig
	initial_uefi_umc_state="$(printf %s "$initial_uefi_state" | grep " UMC ")"

	local IFS=$'\n'
	for delete_entry in $initial_uefi_umc_state; do
		efibootmgr -q -B -b "${delete_entry:4:4}"
	done
}

main () {
	efi_parttype="c12a7328-f81f-11d2-ba4b-00a0c93ec93b"
	mounted_efi_partitions=$(lsblk -lno MOUNTPOINTS,PARTTYPE | grep "$efi_parttype")
	valid_kernel_commands=
	kernel_prefixes="vmlinuz vmlinux- kernel- bzImage zImage"

	einfo "Running uefi-mkconfig..."

	if [[ ${EUID} -ne 0 ]]; then
		die "Please run uefi-mkconfig as root!"

	elif [[ -z $(lsblk -o PARTTYPE | grep -v -e "PARTTYPE" -e "^$") ]]; then
		die "lsblk can't see PARTTYPE! Please refer to Troubleshooting section in README or open an issue"

	elif [[ -z ${mounted_efi_partitions} ]]; then
		die "No mounted efi partitions!"

	elif [[ ! -x "$(command -v efibootmgr)" ]]; then
		die "Unable to access efibootmgr command!"

	fi

	# Load kernel commands from config files
	if [[ -n "${INSTALLKERNEL_CONF_ROOT}" ]]; then
		if [[ -f "${INSTALLKERNEL_CONF_ROOT}/uefi-mkconfig" ]]; then
			kernel_commands="$(tr -s "${IFS}" ' ' <"${KERNEL_INSTALL_CONF_ROOT}/uefi-mkconfig")"
		fi
	elif [[ -f /etc/default/uefi-mkconfig ]]; then
		kernel_commands="$(tr -s "${IFS}" ' ' </etc/default/uefi-mkconfig)"
		einfo "Using kernel commands from \"/etc/default/uefi-mkconfig\""
	elif [[ -f /etc/kernel/uefi-mkconfig ]]; then
		kernel_commands="$(tr -s "${IFS}" ' ' </etc/kernel/uefi-mkconfig)"
		einfo "Using kernel commands from \"/etc/kernel/uefi-mkconfig\""
	elif [[ -f /usr/lib/kernel/uefi-mkconfig ]]; then
		kernel_commands="$(tr -s "${IFS}" ' ' </usr/lib/kernel/uefi-mkconfig)"
		einfo "Using kernel commands from \"/usr/lib/kernel/uefi-mkconfig\""
	else
		die "No uefi-mkconfig configuration file found!"
	fi

	[[ -z "$kernel_commands" ]] && ewarn "Warning! Loaded empty uefi-mkconfig configuration file!"

	# Get kernel prefix from os-release
	if [[ -f /etc/os-release ]]; then
		. /etc/os-release
		kernel_prefixes="${kernel_prefixes} ${ID}-"
	elif [[ -f /usr/lib/os-release ]]; then
		. /usr/lib/os-release
		kernel_prefixes="${kernel_prefixes} ${ID}-"
	fi

	# Get kernel prefix from entry-token
	if [[ -f /etc/kernel/entry-token ]]; then
		kernel_prefixes="${kernel_prefixes} $(head -n1 /etc/kernel/entry-token)-"
	elif [[ -f /usr/lib/kernel/entry-token ]]; then
		kernel_prefixes="${kernel_prefixes} $(head -n1 /usr/lib/kernel/entry-token)-"
	fi

	# Verify kernel commands
	for check_config in $kernel_commands; do
		# Strip initrd=* from config file
		[[ "$check_config" != *"initrd="* ]] && strip_kernel_commands="$strip_kernel_commands $check_config"
		# Verify if "root=" is present
		[[ "$check_config" == "root="* ]] && valid_kernel_commands=1
	done
	kernel_commands="$strip_kernel_commands"

	[[ -z $valid_kernel_commands ]] && ewarn "Warning! Kernel command \"root=\" is missing from loaded configuration!"

	# Get initial state of UEFI entries
	initial_uefi_state="$(efibootmgr -u)"

	# Remove pre 2.0 legacy entries
	clean_legacy_entries

	# Wipe all managed entries for regeneration
	wipe_entries

	partition_mounts="${mounted_efi_partitions// "$efi_parttype"/}"

	local kernel_images_new=
	local kernel_images_old=

	kernel_images_new="$(find $partition_mounts -type f -name "*.efi" -printf "%f\n" | grep -E ".*(${kernel_prefixes// /\|}).*")"
	kernel_images_old="$(echo "$kernel_images_new" | grep "old.efi")"
	kernel_images_new="$(echo "$kernel_images_new" | grep -v "old.efi")"

	# Sort kernel images
	kernel_images="$(sort -uV <<< "$kernel_images_old") $(sort -uV <<< "$kernel_images_new")"

	local partition_efis=
	# Add path back to the kernel image
	for kernel_image in $kernel_images; do
		# Variable partition_mounts in this case can't be inside of ""
		partition_efis+="$(find $partition_mounts -name "$kernel_image") "
	done

	# Move backup entries to the beginning of the list
	for efi_file in $partition_efis; do
		[[ -f "${efi_file}.uefibackup" ]] && partition_efis="${efi_file}.uefibackup ${partition_efis}"
	done

	for efi_file in $partition_efis; do
		# Get partition mount
		partition_mount="${efi_file/\/${efi_file#\/*\/}/}"

		# Get partition
		partition=$(mount | grep "$partition_mount ")
		partition=${partition// */}
		partition=${partition//*\//}

		# Skip this kernel version if ignore file is found
		if [[ -f "$efi_file.ignore" ]]; then
			ewarn "Ignoring \"$efi_file\" on \"$partition\""
			continue
		fi

		# Get partition label
		partition_label="$(lsblk "/dev/$partition" -lno PARTLABEL)"
		## In case PARTLABEL isn't set, use UUID
		[[ -z "$partition_label" ]] && partition_label="$(lsblk "/dev/$partition" -lno UUID)"

		# Prepare path which will be inserted into efibootmgr
		efi_file_path="${efi_file//$partition_mount}"

		# Get kernel version
		## Remove everything before /
		kernel_version="${efi_file_path//*\//}"
		## Remove everything before first - to remove any prefixes
		if [[ "$kernel_version" == *"-"* ]]; then
			kernel_version="${kernel_version/${kernel_version%%-*}-/}"
		else
			# Disable kernel_version variable if the kernel efi file is versionless
			kernel_version=
		fi
		## Remove .efi suffix if it exists
		[[ "$kernel_version" == *".efi"* ]] && kernel_version="${kernel_version/.efi*/}"

		backup_efi=
		local backup_entry
		# Check for existence of backup entry
		if [[ "$efi_file" == *".uefibackup" ]]; then
			efi_file="${efi_file/.uefibackup/}"
			backup_entry="$(grep "UMCB ${efi_file_path/.uefibackup} on ${partition_label} " <<< "$initial_uefi_state")"
			if [[ -n "$backup_entry" ]]; then
				einfo "Existing BACKUP UEFI entry \"${backup_entry:4:4}\" for \"${efi_file}\" found on \"${partition}\""
				continue
			fi
			backup_efi=1
		fi

		# If bootnum is empty, set it to max ID 0200
		if [[ -z $bootnum ]]; then
			bootnum=512 # 512 is decimal value of 0200
		else
			# If it already has value, convert it to decimal and subtract 1
			bootnum="$(($(printf %d 0x"$bootnum") - 1))"
		fi

		# Find first free hex address smaller than or equal to 0200 for new UEFI boot entry
		while [[ "$(grep -v " UMC " <<< "$initial_uefi_state")" == *"Boot$(printf %04X $bootnum)"* ]]; do
			bootnum=$((bootnum - 1))
			# Die if script exceeds managed range 0200-0100
			[[ $bootnum -lt 256 ]] && die "All Boot IDs within managed range are taken!"
		done

		# Convert chosen entry ID into hex
		bootnum="$(printf %04X $bootnum)"

		# Add efi entry for efi file
		add_uefi_entry

	done

	einfo "Done"

}

main
