#!/usr/bin/env bash

set -euo pipefail

# Default values
EFI_REGISTER=yes
ALLOW_DOWNGRADE=no
SET_LIMINE_AS_FALLBACK=no

# Parse command-line arguments
while [[ $# -gt 0 ]]; do
	case "$1" in
	--no-efi-register | -n)
		EFI_REGISTER=no
		;;
	--allow-downgrade | -d)
		ALLOW_DOWNGRADE=yes
		;;
	--fallback | -f)
		SET_LIMINE_AS_FALLBACK=yes
		;;
	--skip-uefi | -s)
		SKIP_UEFI=yes
		;;
	--help | -h | *)
		echo "Usage: $0 [options]"
		echo "  --no-efi-register | -n    Skip registering Limine in NVRAM."
		echo "  --allow-downgrade | -d    Allow installing older Limine versions."
		echo "  --fallback        | -f    Make Limine the default fallback bootloader."
		echo "  --skip-uefi       | -s    Skip UEFI check and registration for non-compliant UEFI boards."
		exit 0
		;;
	esac
	shift
done

if [[ $EUID -ne 0 ]]; then
	echo -e "\033[91m limine-install must be run with root privileges.\033[0m" >&2
	exit 1
fi

# 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

# Skip UEFI check on some MSI UEFI boards with broken EFI implementation.
if [[ "${SKIP_UEFI:-no}" == "yes" ]]; then
	SET_LIMINE_AS_FALLBACK=yes
else
	check_uefi || {
		echo "The system is not using UEFI."
		SET_LIMINE_AS_FALLBACK=yes
		SKIP_UEFI=yes
	}
fi

is_x64 || {
	echo "The system is not x86_64."
	exit 0
}

# Function that registers Limine in the UEFI boot manager
register_uefi_entry() {
	local esp_path="$1"
	local boot_label="$2"
	local efi_file_path="$3"

	# Check if efibootmgr exists
	if ! command -v efibootmgr &>/dev/null; then
		echo -e "\033[1;31m Error: efibootmgr command not found.\033[0m" >&2
		return 1
	fi

	# Get the disk and partition information
	local source
	if ! read -r source < <(findmnt -n -o SOURCE "$esp_path"); then
		echo -e "\033[1;31m Error: Unable to determine the source device for '$esp_path'.\033[0m" >&2
		return 1
	fi

	local gpt_uuid
	if ! read -r gpt_uuid < <(findmnt -n -o PARTUUID "$esp_path"); then
		echo -e "\033[1;31m Error: Unable to determine GPT UUID for '$esp_path'.\033[0m" >&2
		return 1
	fi

	# Validate GPT UUID
	if [[ -z "$gpt_uuid" ]]; then
		echo -e "\033[1;31m Error: Empty GPT UUID detected for '$esp_path'.\033[0m" >&2
		return 1
	fi
	if [[ ${#gpt_uuid} -ne 36 ]]; then
		echo -e "\033[1;31m Error: Invalid GPT UUID length (${#gpt_uuid}): '$gpt_uuid' for '$esp_path'.\033[0m" >&2
		return 1
	fi

	# Check if the Limine EFI entry already exists
	if efibootmgr | grep -Fi "${gpt_uuid}" | grep -Fqi "/${efi_file_path}"; then
		#echo "Info: EFI boot entry '$boot_label' already exists. No action needed."
		return 0
	fi

	# Extract disk and partition
	local disk part
	if [[ "$source" =~ ^(/dev/nvme[0-9]+n[0-9]+)p([0-9]+)$ ]]; then
		# NVMe: /dev/nvmeXnYpZ -> disk=/dev/nvmeXnY, part=Z
		disk="${BASH_REMATCH[1]}"
		part="${BASH_REMATCH[2]}"
	elif [[ "$source" =~ ^(/dev/mmcblk[0-9]+)p([0-9]+)$ ]]; then
		# MMC/SD: /dev/mmcblkXpY -> disk=/dev/mmcblkX, part=Y
		disk="${BASH_REMATCH[1]}"
		part="${BASH_REMATCH[2]}"
	elif [[ "$source" =~ ^(/dev/[a-z]+)([0-9]+)$ ]]; then
		# /dev/sdaX, /dev/vdaX, /dev/xvdaX -> disk=/dev/[sda|vda|xvda], part=X
		disk="${BASH_REMATCH[1]}"
		part="${BASH_REMATCH[2]}"
	else
		printf '\033[1;31m Error: Failed to parse disk and partition from source %q.\033[0m\n' "$source" >&2
		return 1
	fi

	# Add the EFI entry using efibootmgr
	if efibootmgr --create \
		--disk "${disk}" \
		--part "${part}" \
		--label "${boot_label}" \
		--loader "${efi_file_path}" \
		--unicode; then
		echo "EFI boot entry '${boot_label}' for '${efi_file_path}' added successfully."
	else
		echo -e "\033[1;31m Error: UEFI NVRAM may already be full.\033[0m" >&2
		printf "\033[1;31m Error: Failed to add EFI boot entry: disk=%s, part=%s, label=%s, loader=%s\033[0m\n" "${disk}" "${part}" "${boot_label}" "${efi_file_path}" >&2
		return 1
	fi
}

# Check for downgrade when an older Limine version is detected
# Returns:
#   0 -> OK (no downgrade or downgrade allowed)
#   1 -> Downgrade detected but not allowed
#   2 -> Invalid version or comparison error
check_limine_downgrade() {
	local src_file="$1"
	local tgt_file="$2"
	local name="${3:-Limine}"

	compare_limine_versions "$src_file" "$tgt_file"
	local cmp_result=$?

	case "$cmp_result" in
	0 | 1)
		# Source is equal or newer → proceed
		return 0
		;;
	2)
		# Downgrade detected
		local src_ver tgt_ver
		src_ver=$(get_limine_version "$src_file") || return 2
		tgt_ver=$(get_limine_version "$tgt_file") || return 0

		echo "Info: Downgrade detected for $name from version $tgt_ver to $src_ver."

		if [[ "$ALLOW_DOWNGRADE" != "yes" ]]; then
			echo "Info: Downgrade skipped; '--allow-downgrade' not set for limine-install."
			return 1
		fi

		return 0
		;;
	*)
		# Invalid version or comparison error
		return 2
		;;
	esac
}

# Check if a Limine binary is within the supported major version range.
# Args: <limine binary file> <min major version> <max major version> <name>
# Returns:
#   0 -> Supported version
#   1 -> Version newer than $max_version (unsupported future major version)
#   2 -> Version older than $min_version (too old / unsupported)
#   3 -> Failed to retrieve version (get_limine_version failed)
check_limine_upgrade() {
	local src_file="$1"
	local min_major_version="$2"
	local max_major_version="$3"
	local name="${4:-Limine}"
	local src_version major_version

	src_version=$(get_limine_version "$src_file") || return 3

	# Extract the major version number
	major_version=$(cut -d. -f1 <<<"$src_version")

	if ((major_version > max_major_version)); then
		echo "Info: $name version $major_version exceeds the supported maximum version ($max_major_version) for limine-entry-tool; skipping upgrade."
		return 1
	fi

	if ((major_version < min_major_version)); then
		echo -e "\033[1;33mWarning: $name version $major_version is below the supported minimum version ($min_major_version) for limine-entry-tool; skipping installation.\033[0m" >&2
		return 2
	fi
}

# Check for conflicting Limine config files in common ESP locations.
# Warn if any are found, since Limine v10.3.0+ will ignore the default ${ESP_PATH}/limine.conf
check_limine_config_conflicts() {
	local paths=(
		"${ESP_PATH}/EFI/limine/limine.conf"
		"${ESP_PATH}/EFI/BOOT/limine.conf"
		"${ESP_PATH}/boot/limine.conf"
		"${ESP_PATH}/limine/limine.conf"
	)
	local found=0

	for config_path in "${paths[@]}"; do
		if [[ -f "$config_path" ]]; then
			echo -e "\033[1;33mDetected conflicting config: $config_path\033[0m" >&2
			found=1
		fi
	done

	if ((found)); then
		echo -e "\033[1;33mWarning: One or more conflicting Limine config files were detected.\033[0m" >&2
		echo -e "\033[1;33m         Limine will load one of these files based on its search order, ignoring the default ${ESP_PATH}/limine.conf.\033[0m" >&2
		echo -e "\033[1;33m         To avoid unexpected behavior, please remove the conflicting Limine config files manually.\033[0m" >&2
	fi
}

# Function to update a specific EFI file if needed
update_limine_efi() {
	local source_file="${BINARY_SOURCE_PATH}"
	local backup_file="${LIMINE_DIR_PATH}${LIMINE_BACKUP_FILE}"
	local target_file="${LIMINE_DIR_PATH}${LIMINE_EFI_FILE}"

	# Verify the original Limine binary file exists
	if [[ ! -f "$source_file" ]]; then
		echo -e "\033[1;31mError: Limine EFI file '$source_file' is not available!\033[0m" >&2
		return 1
	fi

	# If a backup archive exists, compare hashes
	if [[ -f "$backup_file" ]]; then
		local source_hash
		local backup_hash

		source_hash=$(b2sum "$source_file" | awk '{print $1}')
		backup_hash=$(tar --to-command="b2sum" -xf "$backup_file" | awk '{print $1}')

		if [[ "$source_hash" == "$backup_hash" ]]; then
			# File is identical to backup → no update needed
			return 0
		fi
	fi

	# Check Limine version only if a target file already exists
	if [[ -f "$target_file" ]]; then
		if ! check_limine_downgrade "$source_file" "$target_file" "Limine EFI"; then
			return 0
		fi
		if ! check_limine_upgrade "$source_file" 8 10 "Limine EFI"; then
			return 0
		fi
	fi

	echo "Updating $target_file..."
	cp "$source_file" "$target_file" || return 1

	echo "Creating a backup $backup_file..."
	tar -cf "$backup_file" --directory="${LIMINE_DIR_PATH}" "${LIMINE_EFI_FILE}" || {
		echo -e "\033[1;31mError: Failed to create backup at $backup_file.\033[0m" >&2
		return 1
	}
}

# Function to update the Limine fallback EFI binary if needed
update_limine_fallback() {
	local source_file="${BINARY_SOURCE_PATH}"
	local target_file="${BINARY_FALLBACK_PATH}"

	# Verify the original Limine binary file exists
	if [[ ! -f "$source_file" ]]; then
		echo -e "\033[1;31mError: Limine EFI file '$source_file' is not available!\033[0m" >&2
		return 1
	fi

	# If target already exists, compare hashes
	if [[ -f "$target_file" ]]; then
		local source_hash
		local target_hash

		source_hash=$(b2sum "$source_file" | awk '{print $1}')
		target_hash=$(b2sum "$target_file" | awk '{print $1}')

		if [[ "$source_hash" == "$target_hash" ]]; then
			# No update needed
			return 0
		fi
	fi

	# Check for downgrade and upgrade, only if fallback exists
	if [[ -f "$target_file" ]]; then
		if ! check_limine_downgrade "$source_file" "$target_file" "Limine (fallback)"; then
			return 0
		fi
		if ! check_limine_upgrade "$source_file" 8 10 "Limine (fallback)"; then
			return 0
		fi
	fi

	echo "Updating Limine fallback EFI..."
	cp -f "$source_file" "$target_file" || {
		echo -e "\033[1;31mError: Failed to copy fallback Limine EFI.\033[0m" >&2
		return 1
	}
}

# Validate ESP_PATH
# shellcheck disable=SC2153
if [[ ! -d "${ESP_PATH}" ]]; then
	echo -e "\033[1;31mError: '${ESP_PATH}' is invalid. Please set ESP_PATH correctly.\033[0m" >&2
	exit 1
fi

# Create necessary directories
if ! install -dm 700 "${ESP_PATH}/EFI/BOOT/" "${LIMINE_DIR_PATH}"; then
	echo -e "\033[1;31mError: Failed to create directories in '${ESP_PATH}'.\033[0m" >&2
	exit 1
fi

# Check for any Limine config conflicts on the same boot partition and warn users if found.
check_limine_config_conflicts

# Optionally add some bootloaders if they are present on the same ESP
if [[ "${FIND_BOOTLOADERS:-no}" == "yes" && "${SKIP_UEFI:-no}" == "no" ]]; then
	bootloader_path="${ESP_PATH}/EFI/systemd/systemd-bootx64.efi"
	if [[ -f "${bootloader_path}" ]]; then
		limine-entry-tool --add-efi "Systemd-boot" "${bootloader_path}" --comment "Systemd bootloader" --priority 20 --quiet
	fi

	bootloader_path="${ESP_PATH}/EFI/refind/refind_x64.efi"
	if [[ -f "${bootloader_path}" ]]; then
		limine-entry-tool --add-efi "rEFInd" "${bootloader_path}" --comment "rEFInd bootloader" --priority 20 --quiet
	fi

	bootloader_path="${ESP_PATH}/EFI/BOOT/BOOTX64.EFI"
	if [[ -f "${bootloader_path}" ]]; then
		limine-entry-tool --add-efi "EFI fallback" "${bootloader_path}" --comment "Default EFI loader" --priority 10 --quiet
	fi
fi

# Optionally set Limine as fallback bootloader
if [[ "${ENABLE_LIMINE_FALLBACK:-no}" == "yes" ]] || [[ "${SET_LIMINE_AS_FALLBACK}" == "yes" ]] || [[ ! -f "${BINARY_FALLBACK_PATH}" ]]; then
	update_limine_fallback
fi

mutex_lock "limine-install"
# Update Limine EFI file
if ! update_limine_efi; then
	mutex_unlock
	exit 1
fi
enroll_config
mutex_unlock

# Skip UEFI registration on some MSI UEFI boards with broken EFI implementation.
if [[ "${SKIP_UEFI:-no}" == "yes" ]]; then
	echo "Skipping UEFI registration."
	exit 0
fi

# Optionally register Limine in UEFI system
if [[ "${EFI_REGISTER}" == "yes" ]]; then
	if register_uefi_entry "${ESP_PATH}" "Limine" "${LIMINE_EFI_PATH}"; then
		echo "Limine EFI install completed successfully."
	else
		echo "Limine EFI install failed."
		exit 1
	fi
else
	echo "Limine EFI update completed successfully."
fi
