#!/usr/bin/env bash

die() {
	echo -e " ${NOCOLOR-\e[1;31m*\e[0m }${*}" >&2
	echo -e " ${NOCOLOR-\e[1;31m*\e[0m }For more information please refer to https://github.com/Biosias/uefi-mkconfig"
	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 () {

	# Substitute %values in entry labels 
	entry_label="${entry_label/\%efi_file_path/${efi_file_path/.uefibackup}}"
	entry_label="${entry_label/\%partition_label/$partition_label}"
	entry_label="${entry_label/\%kernel_version/${kernel_version/-*}}"
	entry_label="${entry_label/\%linux_name/$NAME}"
	entry_label="${entry_label/\%entry_id/$entry_id}"
	entry_label="${entry_label/\%partition/$partition}"
	
	# 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

	[[ $ENTRY_LABEL_LIMIT == 1 ]] && (( $(wc -c <<< "${entry_label}") > 60 )) && die "Entry label length is over 60 characters! Creater shorter one or disable this check. "

	# 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

	#Get partition ID
	partition_id="$(printf "%s\n" "$lsblk_initial_state" | grep "$partition" | cut -d";" -f4)"

	# Add new entry
	if [[ -n $backup_efi ]]; then
		# When creating backup entry, don't add it to the bootorder
		if [[ $UMC_TEST == true ]]; then
			echo "TEST: efibootmgr -q --create-only -b \"$bootnum\" --disk \"/dev/$partition\" --part \"$partition_id\" --label \"$entry_label\" --loader \"${efi_file_path//\//\\}\" -u \"$adding_kernel_commands\""
		else
			efibootmgr -q --create-only -b "$bootnum" --disk /dev/"$partition" --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\""
		fi
	else
		if [[ $UMC_TEST == true ]]; then
			echo "TEST: efibootmgr -q --create -b \"$bootnum\" --disk \"/dev/$partition\" --part \"$partition_id\" --label \"$entry_label\" --loader \"${efi_file_path//\//\\}\" -u \"$adding_kernel_commands\""
		else
			efibootmgr -q --create -b "$bootnum" --disk /dev/"$partition" --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
	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
				if [[ $UMC_TEST == true ]]; then
					echo "TEST: efibootmgr -q -B -b \"${entry:4:4}\""
				else
					efibootmgr -q -B -b "${entry:4:4}"
				fi

				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
		if [[ $UMC_TEST == true ]]; then
			echo "TEST: efibootmgr -q -B -b \"${delete_entry:4:4}\""
		else
			efibootmgr -q -B -b "${delete_entry:4:4}"
		fi
	done
}

load_config () {
	default_label="%entry_id %linux_name Linux %kernel_version"
	local kernel_config_paths
	local kernel_config_path
	local kernel_configs_wip
	local config
	
	kernel_config_paths="/etc/default /etc/kernel /usr/lib/kernel"
	[[ -n "${INSTALLKERNEL_CONF_ROOT}" ]] && kernel_config_paths="${INSTALLKERNEL_CONF_ROOT} ${kernel_config_paths}"

	for kernel_config_path in ${kernel_config_paths}; do
		if [[ -f ${kernel_config_path}/uefi-mkconfig ]]; then
			# Check if found config file is legacy or not
			if [[ "$(grep "KERNEL_CONFIG" "${kernel_config_path}/uefi-mkconfig")" == "" ]]; then
				kernel_configs="${default_label} ; $(tr -s "${IFS}" ' ' <"${kernel_config_path}/uefi-mkconfig")"
				# Turn off label length for legacy config file
				ENTRY_LABEL_LIMIT=0
				ewarn "Using legacy uefi-mkconfig configuration file format!"
			else
				kernel_configs="$(grep "KERNEL_CONFIG" "${kernel_config_path}/uefi-mkconfig" | grep -v "#")"
			fi
			einfo "Using kernel commands from \"${kernel_config_path}/uefi-mkconfig\""
			break
		fi
	done
	
	if [[ -z ${kernel_configs} ]]; then
		ewarn "No configuration found. Creating default one. Don't forget to configure kernel commandline"
		
		echo '# Add your kernel commandline arguments following "; "
# KERNEL_CONFIG="%entry_id %linux_name Linux %kernel_version ; example=test kernel=test cmdline=test arguments=test"
KERNEL_CONFIG="%entry_id %linux_name Linux %kernel_version ; "

# Turn on to only add efi files of the latest kernel version
# ONLY_LATEST=false

# Entry label length limit of 60 characters for better compatibility with UEFI Firmware
# Can be disabled but thorough testing is recommended before using entry labels longer than 60 characters 
# ENTRY_LABEL_LIMIT=true' > "/etc/default/uefi-mkconfig"

		kernel_configs="${default_label} ; "	
	else
		# Check if only the latest efi file should be added
		[[ -n $(grep "ONLY_LATEST=true" "${kernel_config_path}/uefi-mkconfig" | grep -v "#") ]] && ONLY_LATEST=1

		# Check if entry label limit is turned off
		[[ -n $(grep "ENTRY_LABEL_LIMIT=false" "${kernel_config_path}/uefi-mkconfig" | grep -v "#") ]] && ENTRY_LABEL_LIMIT=0
	fi

	local IFS=$'\n'
	for config in ${kernel_configs}; do
		
		local IFS=' '
		# Verify kernel commands
		for check_config in ${config}; 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

		[[ -z $valid_kernel_commands ]] && ewarn "Warning! Kernel command \"root=\" is missing from loaded configuration!"
		
		[[ -z ${config} ]] && ewarn "Warning! Loaded empty uefi-mkconfig configuration!"
	
		kernel_configs_wip="$strip_kernel_commands
$kernel_configs_wip"

		valid_kernel_commands=
		strip_kernel_commands=
		
		local IFS=$'\n'
	done

	kernel_configs_wip=${kernel_configs_wip//'KERNEL_CONFIG="'}
	kernel_configs=${kernel_configs_wip//\"}
}

main () {
	efi_parttype="c12a7328-f81f-11d2-ba4b-00a0c93ec93b"
	
	# Load mocked-up lsblk output for script testing
	if [[ $UMC_TEST == true ]] && [[ $UMC_MOCK == true ]]; then
		lsblk_initial_state="$(sed "s/ * /;/g" "$UMC_TEST_LSBLK")"
	else
		lsblk_initial_state="$(lsblk -lno NAME,MOUNTPOINT,PARTTYPE,PARTN,UUID,PARTLABEL | grep "$efi_parttype" | sed "s/ * /;/g")"
	fi
	
	mounted_efi_partitions="$(printf "%s\n" "$lsblk_initial_state" | grep "/")"
	valid_kernel_commands=
	kernel_prefixes="vmlinuz vmlinux- kernel- bzImage zImage"
	kernel_configs=
	# Entry label limit is on by default
	ENTRY_LABEL_LIMIT=1

	einfo "Running uefi-mkconfig..."

	if [[ ${EUID} -ne 0 ]]; then
		die "Please run uefi-mkconfig as root!"
	
	elif [[ -z $(printf "%s\n" "$mounted_efi_partitions" | cut -d";" -f3) ]]; then
		die "lsblk can't see PARTTYPE!"

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

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

	fi

	# 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

	# Get initial state of UEFI entries
	
	if [[ $UMC_TEST == "true" ]]; then
		initial_uefi_state="$(cat "$UMC_TEST_EFIBOOTMGR")"
	else
		initial_uefi_state="$(efibootmgr -u)"
	fi

	# Remove pre 2.0 legacy entries
	clean_legacy_entries

	# Wipe all managed entries for regeneration
	wipe_entries

	partition_mounts="$(printf "%s\n" "$mounted_efi_partitions" | cut -d";" -f2)"

	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")"

        # Check if any kernel images were found
        if [ -z "${kernel_images_old}" ] && [ -z "${kernel_images_new}" ]; then
                die "No efi kernel images found!"
        fi

	# 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

	load_config

	# If only the latest kernel image shoud be added, remove the rest from the list
	local latest_efi_file
	latest_efi_file=
	[[ -n ${ONLY_LATEST} ]] && latest_efi_file="${partition_efis/*\/}"
	
	# Skip this kernel version if ignore file is found
	for efi_file in $partition_efis; do
		if [[ -f "$efi_file.ignore" ]]; then
			# Get partition mount
			partition_mount="${efi_file/\/${efi_file#\/*\/}/}"
			# Get partition
			partition=$(printf "%s\n" "$mounted_efi_partitions" | grep "$partition_mount;" | cut -d";" -f1)
			
			ewarn "Ignoring \"$efi_file\" on \"$partition\""
			# Remove ignored efi files from the list
			partition_efis="${partition_efis/"${efi_file}"}"
		fi

		# Make sure only the latest efi file is added from all partitions
		[[ -n ${latest_efi_file} ]] && [[ "${efi_file/*\/}" != "${latest_efi_file//\ }" ]] && partition_efis="${partition_efis/"${efi_file}"}"

	done

	# Get number of entries that will be created to be used as unique entry id
	entry_id=$(($(wc -w <<< "$partition_efis") * ($(wc -l <<< "$kernel_configs")-1)))

	for efi_file in $partition_efis; do

		# Get partition mount
		partition_mount="${efi_file/\/${efi_file#\/*\/}/}"

		# Get partition
		partition=$(printf "%s\n" "$mounted_efi_partitions" | grep "$partition_mount;" | cut -d";" -f1)

		# Get partition label
		partition_label="$(printf "%s\n" "$mounted_efi_partitions" | grep "$partition" | cut -d";" -f6)"
		## In case PARTLABEL isn't set, use UUID
		[[ -z "$partition_label" ]] && partition_label="$(printf "%s\n" "$mounted_efi_partitions" | grep "$partition" | cut -d";" -f5)"

		# 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

		# Add efi entry for efi file
		local IFS=$'\n'
		for kernel_config in ${kernel_configs}; do
			
			# Make sure efi_file_path is correct for every iteration
			efi_file_path="${efi_file//$partition_mount}"
			
			entry_label="${kernel_config/\ \;*}"
			entry_label="${entry_label/\ }"
			kernel_commands="${kernel_config/*\ \;}"

			# 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_uefi_entry
			entry_id=$((entry_id-1))
		done

	done

	einfo "Done"

}

main
