#!/usr/bin/env bash

set -euo pipefail

# Import functions and environment variables
LIMINE_FUNCTIONS_PATH=/usr/lib/limine/limine-common-functions
# shellcheck disable=SC1090
source "${LIMINE_FUNCTIONS_PATH}" || {
	echo -e "\033[1;31m Error: Failed to source '${LIMINE_FUNCTIONS_PATH}'.\033[0m" >&2
	exit 1
}

initialize_header || exit 1

all=false
lines=()

# Read stdin and decide whether to rebuild all initramfs or specific ones
while read -r line; do
	# If the line does not point to a pkgbase file, trigger a full rebuild
	if [[ "${line}" != */pkgbase ]]; then
		all=true
		continue
	fi
	# Otherwise record the corresponding kernel directory
	lines+=("/${line}")
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 "limine-dracut-install"
reset_enroll_config

# 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
		echo -e "\033[1;31m Error: kernel name of '${pkgbase_file}' is empty, skipping.\033[0m" >&2
		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
			echo -e "\033[1;33m Warning: CUSTOM_UKI_NAME is invalid. Use only lowercase letters (a–z) and digits (0–9), no spaces or special characters.\033[0m" >&2
			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 check_uefi && { [[ "${DRACUT_UKI:-}" == "yes" ]] || [[ "${ENABLE_UKI:-}" == "yes" ]]; }; then
		# Unified Kernel Image (UKI)
		install -dm700 "${ukiDirPath}"

		echo "Building UKI for ${kName} (${kVersion})"
		if are_snapper_and_sb_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
				echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
				continue
			fi
		else
			kernel_cmdline=$(limine-entry-tool --get-cmdline "${kName}" --no-mutex | tail -n 1) || {
				echo -e "\033[1;31m Error: failed to get kernel cmdline for '${kName}'.\033[0m" >&2
				continue
			}
			if ! dracut --force --reproducible --uefi --hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}" --kernel-cmdline "${kernel_cmdline}"; then
				echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
				continue
			fi
		fi
		# Move the UKI to the boot partition
		mv -f "${tmp_uki_path}" "${uki_path}"
		echo "UKI stored in: ${uki_path}"
		sb_sign "${uki_path}"
		limine-entry-tool --add-uki "${kName}" "${uki_path}" --comment "${kVersion}" --no-mutex

		# Building UKI fallback
		if [[ "${DRACUT_FALLBACK:-}" == "yes" ]]; then
			echo "Building UKI fallback for ${kName} (${kVersion})"
			if are_snapper_and_sb_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
					echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
					continue
				fi
			else
				kernel_cmdline=$(limine-entry-tool --get-cmdline "${kName}-fallback" --no-mutex | tail -n 1) || {
					echo -e "\033[1;31m Error: failed to get kernel cmdline for '${kName}-fallback'.\033[0m" >&2
					continue
				}
				if ! dracut --force --reproducible --uefi --no-hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_uki_path}" --kernel-cmdline "${kernel_cmdline}"; then
					echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
					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 "${kVersion} fallback" --no-mutex
		else
			# Remove fallback if not needed
			[[ -f "${uki_fallback_path}" ]] && limine-entry-tool --remove-uki "${kName}-fallback" --no-mutex --quiet
		fi

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

	else

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

		echo "Building initramfs for ${kName} (${kVersion})"
		if ! dracut --force --reproducible --no-uefi --hostonly --no-hostonly-cmdline --kver "${kVersion}" "${tmp_initramfs_path}"; then
			echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
			continue
		fi
		# Copy the vmlinuz to the boot partition
		install -Dm600 "${kFilePath}" "${vmlinuz_path}"
		echo "Kernel stored in: ${vmlinuz_path}"
		# Move the initramfs to the boot partition
		mv -f "${tmp_initramfs_path}" "${initramfs_path}"
		echo "Initramfs stored in: ${initramfs_path}"
		limine-entry-tool --add "${kName}" "${initramfs_path}" "${vmlinuz_path}" --comment "${kVersion}" --no-mutex

		# Building initramfs fallback
		if [[ "${DRACUT_FALLBACK:-}" == "yes" ]]; then
			echo "Building fallback initramfs for ${kName} (${kVersion})"
			if ! dracut --force --reproducible --no-uefi --no-hostonly --kver "${kVersion}" "${tmp_initramfs_path}"; then
				echo -e "\033[1;31m Error: dracut failed for kernel ${kVersion}, skipping.\033[0m" >&2
				continue
			fi
			mv -f "${tmp_initramfs_path}" "${initramfs_fallback_path}"
			echo "Fallback initramfs stored in: ${initramfs_fallback_path}"
			limine-entry-tool --add "${kName}" "${initramfs_fallback_path}" "${vmlinuz_path}" "-fallback" --comment "${kVersion} fallback" --no-mutex
		else
			# Remove fallback if not needed
			[[ -f "${initramfs_fallback_path}" ]] && {
				limine-entry-tool --remove-entry "${kName}-fallback" --no-mutex --quiet
				rm -f "${initramfs_fallback_path}"
			}
		fi

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

enroll_config
mutex_unlock

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