#!/usr/bin/env bash

set -euo pipefail

args=('--force' '--no-hostonly-cmdline')
lines=()
status=0
all=false
initramfs_error_file="/var/lib/garuda/initramfs_error"
key_path="/var/lib/garuda/secureboot/keys"

sign_kernel() {
	if [ -f "/etc/garuda/secureboot/enabled" ]; then
		sbsign --key "${key_path}/MOK.key" --cert "${key_path}/MOK.crt" --output "$1" "$1" >/dev/null ||
			(echo -e "\033[1;31mFailed to sign kernel: $1\033[0m" && return 1)
	fi
}

# This function implements a workaround for a bug in grub/btrfs where the vmlinuz file can not be read.
# The workaround is to copy the vmlinuz file with a method that prevents reflinks completely.
# $1 is the source kernel directory
# $2 is the pkgbase
install_kernel() {
	set -euo pipefail

	local src_dir="$1"
	local pkgbase="$2"
	local vmlinuz_file="${src_dir}/vmlinuz"
	local vmlinuz_dest="/boot/vmlinuz-${pkgbase}"

	# Force delete existing vmlinuz file
	rm -f "${vmlinuz_dest}"

	# Sanity check
	if [[ ! -f "${vmlinuz_file}" ]] || [[ "${pkgbase}" == "" ]]; then
		echo -e "\033[1;31m!!Unknown error while installing kernel!!\033[0m"
		return 1
	fi

	# Copy vmlinuz file via dd
	if dd if="${vmlinuz_file}" of="${vmlinuz_dest}" bs=4M status=none; then
		chmod 0644 "${vmlinuz_dest}"
		# Compare files to double check the integrity
		if cmp -s "${vmlinuz_file}" "${vmlinuz_dest}"; then
			sign_kernel "${vmlinuz_dest}" || { echo -e "\033[1;31mSigning the kernel failed: ${vmlinuz_dest}\033[0m"; return 1; }
			return 0
		fi
	fi
	echo -e "\033[1;31mSafe kernel copy failed: ${vmlinuz_file} -> ${vmlinuz_dest}\033[0m"
	if install -Dm0644 "${vmlinuz_file}" "${vmlinuz_dest}"; then
		sign_kernel "${vmlinuz_dest}" || { echo -e "\033[1;31mSigning the kernel failed: ${vmlinuz_dest}\033[0m"; return 1; }
		echo -e "\033[1;33mWarning: Kernel copy with fallback method succeeded: ${vmlinuz_file} -> ${vmlinuz_dest}\033[0m"
		return 0
	fi
	return 1
}

while read -r line; do
	# Check if any changed file is not a kernel image. This means we need to rebuild all initramfs.
	# Triggers when it's a change to dracut files
	# Does not match usr/lib/modules/*
	if [[ "$line" != 'usr/lib/modules/'* ]]; then
		all=true
		# Do not break. Otherwise, we will get a "broken pipe" error
		continue
	fi

	if [[ "$line" == 'usr/lib/modules/'+([^/])'/pkgbase' ]]; then
		lines+=("${line}")
	fi
done

if [[ "$all" == true ]]; then
	echo ":: Rebuilding all initramfs due to changes in dracut files."
	lines=()
	for line in /usr/lib/modules/*/pkgbase; do
		# if pkgbase file does not belong to any package then skip this kernel
		# This likely means it's a restored kernel modules backup and not relevant for the next reboot.
		if ! pacman -Qqo "$line" &>/dev/null; then
			continue
		fi
		# Without leading slash
		lines+=("${line#'/'}")
	done
fi

for line in "${lines[@]}"; do
	read -r pkgbase < "/${line}"
	kver="${line#'usr/lib/modules/'}"
	kver="${kver%'/pkgbase'}"

	if
		echo ":: Building initramfs for ${pkgbase} (${kver})" &&
		install_kernel "/${line%'/pkgbase'}" "$pkgbase" &&
		dracut "${args[@]}" -L 3 --hostonly "/boot/initramfs-${pkgbase}.img" --kver "$kver" &&

		echo ":: Building fallback initramfs for ${pkgbase} (${kver})" &&
		dracut "${args[@]}" -L 1 --no-hostonly -o "network rdma" "/boot/initramfs-${pkgbase}-fallback.img" --kver "$kver"
	then
		if [[ -f "${initramfs_error_file}" ]]; then
			sed -i "/${pkgbase}/d" "${initramfs_error_file}"
			if [[ ! -s "${initramfs_error_file}" ]]; then
				rm -f "${initramfs_error_file}"
			fi
		fi
	else
		# If dracut fails, we capture the error status but continue processing other kernels.
		echo -e "\033[1;31mFailed to build initramfs for ${pkgbase} (${kver})\033[0m"
		echo "${pkgbase}" >> "${initramfs_error_file}"
		status=1
		continue
	fi
done

if [[ $status -eq 0 ]] && [[ "$all" == true ]]; then
	rm -f "${initramfs_error_file}"
fi

exit $status