#!/usr/bin/env bash

# Import functions and environment variables
readonly LIMINE_FUNCTIONS_PATH=/usr/lib/limine/limine-common-functions
# shellcheck disable=SC1090
source "${LIMINE_FUNCTIONS_PATH}" || {
	echo "ERROR: Failed to source '${LIMINE_FUNCTIONS_PATH}'" >&2
	exit 1
}

readonly SCRIPT_NAME="limine-dracut-install"
export HOOK_CALLER="$SCRIPT_NAME"
export HOOK_CMDLINE="$*"

initialize_header || exit 1

all=false
lines=()

# Read stdin and decide whether to rebuild all initramfs or specific ones
while IFS= read -r line; do
	# If the line does not point to a kernel modules directory, trigger a full rebuild
	if [[ "${line}" != usr/lib/modules/* ]]; then
		all=true
		# pacman streams all NeedsTargets entries into stdin and will trigger:
		# "unable to write to pipe (Broken pipe)" if reading stops early.
		# To avoid this, use "continue" to keep reading until stdin is fully drained instead of "break".
		continue
	fi
	# Otherwise record the corresponding kernel directory
	# Make sure there are no duplicate paths added
	if ! $all && [[ " ${lines[*]} " != *" /${line/extramodules\//pkgbase} "* ]]; then
		lines+=("/${line/extramodules\//pkgbase}")
	fi
done

# If a full rebuild was requested, override the list with all kernel modules
if [[ "${all}" == true ]]; then
	lines=(/usr/lib/modules/*/pkgbase)
fi

mutex_lock "${SCRIPT_NAME}"
rc=0
run_boot_hooks pre || rc=$?
if ((rc >= 100)); then
	mutex_unlock
	exit "$rc"
fi

# Process each kernel-pkgbase file
for pkgbase_file in "${lines[@]}"; do
	# Skip if the pkgbase file doesn't exist
	[[ -f "${pkgbase_file}" ]] || continue

	# Skip if pkgbase file is not owned by any package
	if ! pacman -Qqo "${pkgbase_file}" &>/dev/null; then
		continue
	fi

	# Read kernel name from pkgbase file
	kName="$(<"${pkgbase_file}")"
	# Remove all spaces that are invalid in FAT32 names to create kernel file or directory
	kName="${kName// /}"
	if [[ -z "${kName}" ]]; then
		error_msg "kernel name of '${pkgbase_file}' is empty, skipping."
		continue
	fi

	# Read kernel version from the parent directory of the pkgbase file
	kDir="${pkgbase_file%/pkgbase}"
	kVersion="${kDir##*/}"
	kFilePath="${kDir}/vmlinuz"

	if [[ -n "${CUSTOM_UKI_NAME:-}" ]]; then
		if [[ "$CUSTOM_UKI_NAME" =~ ^[a-z0-9]+$ ]]; then
			uki_prefix="$CUSTOM_UKI_NAME"
		else
			warning_msg "CUSTOM_UKI_NAME is invalid. Use only lowercase letters (a–z) and digits (0–9), no spaces or special characters."
			uki_prefix="$MACHINE_ID"
		fi
	else
		uki_prefix="$MACHINE_ID"
	fi

	# UKI paths
	ukiDirPath="${ESP_PATH}/EFI/Linux"
	uki_path="${ukiDirPath}/${uki_prefix}_${kName}.efi"
	uki_fallback_path="${ukiDirPath}/${uki_prefix}_${kName}-fallback.efi"
	tmp_uki_path="/tmp/staging_uki.efi"

	# vmlinuz & initramfs paths
	kDirPath="${ESP_PATH}/${MACHINE_ID}/${kName}"
	initramfs_path="${kDirPath}/initramfs-${kName}"
	initramfs_fallback_path="${kDirPath}/initramfs-${kName}-fallback"
	vmlinuz_path="${kDirPath}/vmlinuz-${kName}"
	tmp_initramfs_path="/tmp/staging_initramfs.img"

	if is_uefi && { [[ "${DRACUT_UKI:-}" == "yes" ]] || [[ "${ENABLE_UKI:-}" == "yes" ]]; }; then
		# Unified Kernel Image (UKI)
		install -dm700 "${ukiDirPath}"

		info_msg "Building UKI for ${kName} (${kVersion})"
		if is_sb_installed && is_snapper_installed; then
			# Note: When Secure Boot is enabled, a signed UKI ignores external cmdline and won't boot into a Btrfs snapshot. To fix this, remove the embedded cmdline from the UKI.
			if ! dracut --force --reproducible --uefi --hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}"; then
				error_msg "dracut failed for kernel ${kVersion}, skipping."
				continue
			fi
		else
			kernel_cmdline=$(limine-entry-tool --get-cmdline "${kName}" --no-mutex --no-hooks | tail -n 1) || {
				error_msg "failed to get kernel cmdline for '${kName}'."
				continue
			}
			if ! dracut --force --reproducible --uefi --hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}" --kernel-cmdline "${kernel_cmdline}"; then
				error_msg "dracut failed for kernel ${kVersion}, skipping."
				continue
			fi
		fi
		# Move the UKI to the boot partition
		mv -f "${tmp_uki_path}" "${uki_path}"
		info_msg "UKI stored in ${uki_path}"
		sb_sign "${uki_path}"
		limine-entry-tool --add-uki "${kName}" "${uki_path}" --comment "Kernel version: ${kVersion}" --no-mutex --no-hooks

		# Building UKI fallback
		if [[ "${DRACUT_FALLBACK:-}" == "yes" ]]; then
			info_msg "Building UKI fallback for ${kName} (${kVersion})"
			if is_sb_installed && is_snapper_installed; then
				# Note: When Secure Boot is enabled, a signed UKI ignores external cmdline and won't boot into a Btrfs snapshot. To fix this, remove the embedded cmdline from the UKI.
				if ! dracut --force --reproducible --uefi --no-hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}"; then
					error_msg "dracut failed for kernel ${kVersion}, skipping."
					continue
				fi
			else
				kernel_cmdline=$(limine-entry-tool --get-cmdline "${kName}-fallback" --no-mutex --no-hooks | tail -n 1) || {
					error_msg "failed to get kernel cmdline for '${kName}-fallback'."
					continue
				}
				if ! dracut --force --reproducible --uefi --no-hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}" --kernel-cmdline "${kernel_cmdline}"; then
					error_msg "dracut failed for kernel ${kVersion}, skipping."
					continue
				fi
			fi
			# Move the UKI to the boot partition
			mv -f "${tmp_uki_path}" "${uki_fallback_path}"
			echo "UKI fallback stored in ${uki_fallback_path}"
			sb_sign "${uki_fallback_path}"
			limine-entry-tool --add-uki "${kName}-fallback" "${uki_fallback_path}" --comment "Kernel version: ${kVersion} fallback" --no-mutex --no-hooks
		else
			# Remove fallback if not needed
			[[ -f "${uki_fallback_path}" ]] && limine-entry-tool --remove-uki "${kName}-fallback" --no-mutex --no-hooks --quiet
		fi

		# Remove legacy initramfs and vmlinuz
		if [[ -d "${kDirPath}" ]]; then
			limine-entry-tool --remove-kernel "${kName}" --no-mutex --no-hooks --quiet
		fi

	else

		# Initramfs and vmlinuz handling
		install -dm700 "${kDirPath}"

		info_msg "Building initramfs for ${kName} (${kVersion})"
		if ! dracut --force --reproducible --no-uefi --hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_initramfs_path}"; then
			error_msg -e "dracut failed for kernel ${kVersion}, skipping."
			continue
		fi
		# Copy the vmlinuz to the boot partition
		install -Dm600 "${kFilePath}" "${vmlinuz_path}"
		info_msg "Kernel stored in ${vmlinuz_path}"
		# Move the initramfs to the boot partition
		mv -f "${tmp_initramfs_path}" "${initramfs_path}"
		info_msg "Initramfs stored in ${initramfs_path}"
		limine-entry-tool --add-kernel "${kName}" "${initramfs_path}" "${vmlinuz_path}" --comment "Kernel version: ${kVersion}" --no-mutex --no-hooks

		# Building initramfs fallback
		if [[ "${DRACUT_FALLBACK:-}" == "yes" ]]; then
			info_msg "Building fallback initramfs for ${kName} (${kVersion})"
			if ! dracut --force --reproducible --no-uefi --no-hostonly --kver "${kVersion}" "${tmp_initramfs_path}"; then
				error_msg "dracut failed for kernel ${kVersion}, skipping."
				continue
			fi
			mv -f "${tmp_initramfs_path}" "${initramfs_fallback_path}"
			info_msg "Fallback initramfs stored in ${initramfs_fallback_path}"
			limine-entry-tool --add-kernel "${kName}" "${initramfs_fallback_path}" "${vmlinuz_path}" "-fallback" --comment "Kernel version: ${kVersion} fallback" --no-mutex --no-hooks
		else
			# Remove fallback if not needed
			[[ -f "${initramfs_fallback_path}" ]] && {
				limine-entry-tool --remove-kernel "${kName}-fallback" --keep-files --no-mutex --no-hooks --quiet
				rm -f "${initramfs_fallback_path}"
			}
		fi

		# Remove existing UKI files.
		if [[ -f "${uki_path}" ]]; then
			limine-entry-tool --remove-uki "${kName}" --no-mutex --no-hooks --quiet
		fi
		if [[ -f "${uki_fallback_path}" ]]; then
			limine-entry-tool --remove-uki "${kName}-fallback" --no-mutex --no-hooks --quiet
		fi
	fi
done

rc=0
run_boot_hooks post || rc=$?
if ((rc >= 100)); then
	mutex_unlock
	exit "$rc"
fi
mutex_unlock

if [[ -f "/var/lib/limine/removed_kernels.list" ]]; then
	info_msg "Clean up kernel entries if they are unused."
	/usr/share/libalpm/scripts/limine-dracut-remove post
fi
