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

	if [[ "${ENABLE_UKI:-}" == "yes" ]] && check_uefi; 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 ! /usr/bin/mkinitcpio --kernel "${kVersion}" --no-cmdline --uki "${tmp_uki_path}"; then
				echo -e "\033[1;31m==> ERROR: mkinitcpio failed for kernel ${kVersion}, skipping.\033[0m" >&2
				continue
			fi
		else
			# Get kernel cmdline
			if ! limine-entry-tool --get-cmdline "${kName}" --no-mutex | tail -n 1 >/tmp/kernel-cmdline; then
				echo -e "\033[1;31m==> ERROR: failed to get kernel cmdline for '${kName}'.\033[0m" >&2
				continue
			fi

			if ! /usr/bin/mkinitcpio --kernel "${kVersion}" --uki "${tmp_uki_path}" --cmdline /tmp/kernel-cmdline; then
				echo -e "\033[1;31m==> ERROR: mkinitcpio failed for kernel ${kVersion}, skipping.\033[0m" >&2
				continue
			fi
			rm /tmp/kernel-cmdline
		fi
		# Move the initramfs to the boot partition
		mv -f "${tmp_uki_path}" "${uki_path}"
		echo "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 "${kVersion}" --no-mutex

		# Remove existing 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 ! /usr/bin/mkinitcpio --kernel "${kVersion}" --no-cmdline --generate "${tmp_initramfs_path}"; then
			echo -e "\033[1;31m==> ERROR: mkinitcpio 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

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