#!/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-mkinitcpio-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}"
	vmlinuz_path="${kDirPath}/vmlinuz-${kName}"
	tmp_initramfs_path="/tmp/staging_initramfs.img"

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

		info_msg "Building UKI for ${kName} (${kVersion})"
		read -ra uki_options <<<"${MKINITCPIO_UKI_OPTIONS:-}"
		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 ! /usr/bin/mkinitcpio --kernel "${kVersion}" --no-cmdline --uki "${tmp_uki_path}" "${uki_options[@]}"; then
				error_msg "mkinitcpio failed for kernel ${kVersion}, skipping."
				continue
			fi
		else
			# Get kernel cmdline
			if ! limine-entry-tool --get-cmdline "${kName}" --no-mutex --no-hooks | tail -n 1 >/tmp/kernel-cmdline; then
				error_msg "failed to get kernel cmdline for '${kName}'."
				continue
			fi

			if ! /usr/bin/mkinitcpio --kernel "${kVersion}" --uki "${tmp_uki_path}" --cmdline /tmp/kernel-cmdline "${uki_options[@]}"; then
				error_msg "mkinitcpio failed for kernel ${kVersion}, skipping."
				continue
			fi
			rm /tmp/kernel-cmdline
		fi
		# Move the initramfs to the boot partition
		mv -f "${tmp_uki_path}" "${uki_path}"
		info_msg "UKI stored in ${uki_path}"
		#sb_sign "${uki_path}" # Not needed, as it is already handled by mkinitcpio's sbctl hook
		limine-entry-tool --add-uki "${kName}" "${uki_path}" --comment "Kernel version: ${kVersion}" --no-mutex --no-hooks

		# Remove existing 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 ! /usr/bin/mkinitcpio --kernel "${kVersion}" --no-cmdline --generate "${tmp_initramfs_path}"; then
			error_msg "mkinitcpio 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

		# 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-mkinitcpio-remove post
fi
