#!/usr/bin/env bash

readonly SCRIPT_NAME="limine-mkinitcpio-install"
# 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
}

export HOOK_CALLER="$SCRIPT_NAME"
export HOOK_CMDLINE="$*"

initialize_header || exit 1

# Initialize global variables
declare -a PKGBASE_FILES=()
KERNEL_NAME=""
KERNEL_VERSION=""
KERNEL_DIR=""
BOOT_KERNEL_DIR=""
INITRAMFS_FALLBACK_PATH=""
UKI_PATH=""
UKI_FALLBACK_PATH=""
CMDLINE_FILE=""
BOOT_PATH="$ESP_PATH"
UKI_PREFIX="$MACHINE_ID"
read -ra UKI_OPTIONS <<<"${MKINITCPIO_UKI_OPTIONS:-}"
TMP_DIR=$(mktemp -d /tmp/limine-mkinitcpio.XXXXXX) || {
	echo "ERROR: failed to create temporary directory" >&2
	exit 1
}
readonly TMP_DIR
cleanup_tmp_dir() {
	[[ -d "$TMP_DIR" && "$TMP_DIR" == /tmp/limine-mkinitcpio.* ]] && rm -r "$TMP_DIR"
	return 0
}
trap cleanup_tmp_dir EXIT

resolve_uki_prefix() {
	UKI_PREFIX="$MACHINE_ID"
	[[ -z ${CUSTOM_UKI_NAME:-} ]] && return

	if [[ $CUSTOM_UKI_NAME =~ ^[a-z0-9]+$ ]]; then
		UKI_PREFIX=$CUSTOM_UKI_NAME
	else
		warning_msg warning_msg "CUSTOM_UKI_NAME is invalid. Only lowercase letters (a-z) and digits (0-9) are allowed."
	fi
}

resolve_uki_prefix

# Collect pkgbase files from pacman hook stdin.
collect_pkgbase_files() {
	local line path
	local rebuild_all=false
	local -A seen_paths=()
	while IFS= read -r line; do
		if [[ "$line" != usr/lib/modules/* ]]; then
			rebuild_all=true
			# Do not change "continue" to "break", as it causes:
			# "unable to write to pipe (Broken pipe)" if reading stops early.
			continue
		fi
		path="/${line/extramodules\//pkgbase}"
		if [[ -z "${seen_paths[$path]:-}" ]]; then
			PKGBASE_FILES+=("$path")
			seen_paths[$path]=1
		fi
	done
	if $rebuild_all; then
		PKGBASE_FILES=(/usr/lib/modules/*/pkgbase)
	fi
}

# Set the global kernel context for the current pkgbase.
set_kernel_context() {
	local pkgbase_file="$1"
	# Read kernel name from pkgbase file
	KERNEL_NAME="$(<"$pkgbase_file")"
	# Remove all spaces that are invalid in FAT32 names to create kernel file or directory
	KERNEL_NAME="${KERNEL_NAME// /}"
	if [[ -z "${KERNEL_NAME}" ]]; then
		error_msg "kernel name of '${pkgbase_file}' is empty, skipping."
		return 1
	fi
	# Read kernel version from the parent directory of the pkgbase file
	KERNEL_DIR="${pkgbase_file%/pkgbase}"
	KERNEL_VERSION="${KERNEL_DIR##*/}"
	# kernel locations
	BOOT_KERNEL_DIR="${BOOT_PATH}/${MACHINE_ID}/${KERNEL_NAME}"
	INITRAMFS_FALLBACK_PATH="${BOOT_KERNEL_DIR}/initramfs-fallback"

	local uki_install_dir="${BOOT_PATH}/EFI/Linux"
	UKI_PATH="${uki_install_dir}/${UKI_PREFIX}_${KERNEL_NAME}.efi"
	UKI_FALLBACK_PATH="${uki_install_dir}/${UKI_PREFIX}_${KERNEL_NAME}-fallback.efi"
	return 0
}

# Sets CMDLINE_FILE, or leaves it empty if no cmdline should be embedded.
resolve_cmdline_file() {
	local kernel_name="$1"
	CMDLINE_FILE=""
	if is_sb_installed && is_snapper_installed; then
		# Signed UKIs with embedded cmdlines cannot boot Btrfs snapshots.
		# Omit the embedded cmdline when Secure Boot and Snapper are used.
		return 0
	fi
	local cmdline_file="$TMP_DIR/cmdline-${kernel_name}"
	if ! limine-entry-tool --get-cmdline "$kernel_name" --no-mutex --no-hooks | tail -n 1 >"$cmdline_file"; then
		error_msg "failed to get kernel cmdline for '${kernel_name}'."
		return 1
	fi
	CMDLINE_FILE="$cmdline_file"
	return 0
}

# build_uki OUTPUT IS_FALLBACK CMDLINE_FILE
build_uki() {
	local output="$1" is_fallback="$2" cmdline_file="$3"
	local args=(--kernel "$KERNEL_VERSION")
	if [[ -n "$cmdline_file" ]]; then
		args+=(--cmdline "$cmdline_file")
	else
		args+=(--no-cmdline)
	fi
	args+=(--uki "$output")
	if [[ "$is_fallback" == true ]]; then
		args+=(-S autodetect)
	fi
	args+=("${UKI_OPTIONS[@]}")
	/usr/bin/mkinitcpio "${args[@]}"
}

# build_initramfs OUTPUT IS_FALLBACK
build_initramfs() {
	local output="$1" is_fallback="$2"
	local args=(--kernel "$KERNEL_VERSION" --no-cmdline --generate "$output")
	if [[ "$is_fallback" == true ]]; then
		args+=(-S autodetect)
	fi
	/usr/bin/mkinitcpio "${args[@]}"
}

# install_uki KERNEL_NAME UKI_FILE COMMENT
install_uki() {
	local name="$1" file="$2" comment="$3"
	limine-entry-tool --add-uki "$name" "$file" \
		--comment "$comment" --no-mutex --no-hooks
}

# install_kernel KERNEL_NAME INITRAMFS VMLINUZ COMMENT [-fallback]
install_kernel() {
	local name="$1" initramfs="$2" vmlinuz="$3" comment="$4" suffix="${5:-}"
	limine-entry-tool --add-kernel "$name" "$initramfs" "$vmlinuz" ${suffix:+"$suffix"} \
		--comment "$comment" --no-mutex --no-hooks
}

remove_uki_if_present() {
	local name="$1" path="$2"
	[[ -f $path ]] || return 0
	limine-entry-tool --remove-uki "$name" --no-mutex --no-hooks --quiet
	return 0
}

cleanup_legacy_uki_entries() {
	remove_uki_if_present "$KERNEL_NAME" "$UKI_PATH"
	remove_uki_if_present "${KERNEL_NAME}-fallback" "$UKI_FALLBACK_PATH"
}

cleanup_legacy_regular_entry() {
	[[ -d "$BOOT_KERNEL_DIR" ]] && limine-entry-tool --remove-kernel "$KERNEL_NAME" --no-mutex --no-hooks --quiet
}

# build_install_uki KERNEL_NAME TMP_PATH IS_FALLBACK COMMENT
build_install_uki() {
	local kernel_name="$1" tmp_uki="$2" is_fallback="$3" comment="$4"

	resolve_cmdline_file "$kernel_name" || return 1

	if ! build_uki "$tmp_uki" "$is_fallback" "$CMDLINE_FILE"; then
		error_msg "mkinitcpio failed for kernel ${KERNEL_VERSION}, skipping."
		return 1
	fi
	#sb_sign "$tmp_uki" # Not needed, as it is already handled by mkinitcpio's sbctl hook
	install_uki "$kernel_name" "$tmp_uki" "$comment"
	return 0
}

process_uki_kernel() {
	info_msg "Building UKI for ${KERNEL_NAME} (${KERNEL_VERSION})"
	local tmp_uki="$TMP_DIR/${KERNEL_NAME}.efi"
	build_install_uki "$KERNEL_NAME" "$tmp_uki" false "Kernel version: ${KERNEL_VERSION}" || return 1

	if [[ "${MKINITCPIO_FALLBACK:-}" == "yes" || "${MKINITCPIO_FALLBACK:-}" == "$KERNEL_NAME" ]]; then
		info_msg "Building UKI fallback for ${KERNEL_NAME} (${KERNEL_VERSION})"
		local tmp_uki_fb="$TMP_DIR/${KERNEL_NAME}-fallback.efi"
		build_install_uki "${KERNEL_NAME}-fallback" "$tmp_uki_fb" true "Kernel version: ${KERNEL_VERSION} fallback" || return 1
	else
		remove_uki_if_present "${KERNEL_NAME}-fallback" "$UKI_FALLBACK_PATH"
	fi

	cleanup_legacy_regular_entry
	return 0
}

process_regular_kernel() {
	info_msg "Building initramfs for ${KERNEL_NAME} (${KERNEL_VERSION})"
	local tmp_initramfs="$TMP_DIR/initramfs"
	if ! build_initramfs "$tmp_initramfs" false; then
		error_msg "mkinitcpio failed for kernel ${KERNEL_VERSION}, skipping."
		return 1
	fi
	# Cleanup the legacy kernel directory after confirming it contains legacy files.
	if [[ -f "$BOOT_KERNEL_DIR/initramfs-${KERNEL_NAME}" ]]; then
		rm -r "$BOOT_KERNEL_DIR"
	fi
	install_kernel "$KERNEL_NAME" "$tmp_initramfs" "${KERNEL_DIR}/vmlinuz" "Kernel version: ${KERNEL_VERSION}"

	if [[ "${MKINITCPIO_FALLBACK:-}" == "yes" || "${MKINITCPIO_FALLBACK:-}" == "$KERNEL_NAME" ]]; then
		info_msg "Building fallback initramfs for ${KERNEL_NAME} (${KERNEL_VERSION})"
		local tmp_initramfs_fb="$TMP_DIR/initramfs-fallback"
		if ! build_initramfs "$tmp_initramfs_fb" true; then
			error_msg "mkinitcpio failed for kernel ${KERNEL_VERSION}, skipping."
			return 1
		fi
		install_kernel "$KERNEL_NAME" "$tmp_initramfs_fb" "${KERNEL_DIR}/vmlinuz" "Kernel version: ${KERNEL_VERSION} fallback" "-fallback"
	elif [[ -f "$INITRAMFS_FALLBACK_PATH" ]]; then
		limine-entry-tool --remove-kernel "${KERNEL_NAME}-fallback" --keep-files --no-mutex --no-hooks --quiet
		rm -f "$INITRAMFS_FALLBACK_PATH"
	fi

	cleanup_legacy_uki_entries
	return 0
}

# Process a kernel
process_kernel() {
	local pkgbase_file="$1"

	[[ -f "$pkgbase_file" ]] || return 0
	# Skip if pkgbase file is not owned by any package
	pacman -Qqo "$pkgbase_file" &>/dev/null || return 0
	set_kernel_context "$pkgbase_file" || return 0

	if [[ "${ENABLE_UKI:-}" == "yes" ]] && is_uefi; then
		process_uki_kernel || return 0
	else
		process_regular_kernel || return 0
	fi
}

# Main logic
main() {
	collect_pkgbase_files

	mutex_lock "$SCRIPT_NAME"
	if ! run_boot_hooks pre; then
		mutex_unlock
		exit 2
	fi

	for pkgbase_file in "${PKGBASE_FILES[@]}"; do
		process_kernel "$pkgbase_file" || true
	done

	if ! run_boot_hooks post; then
		mutex_unlock
		exit 3
	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
}

main
