#!/usr/bin/env bash

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

# Default values
EFI_REGISTER=yes
ALLOW_DOWNGRADE=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
        ;;
    --help | -h)
        echo "Usage: $0 [--no-efi-register] [--allow-downgrade]"
        echo "  --no-efi-register | -n    Skip registering Limine in NVRAM."
        echo "  --allow-downgrade | -d    Allow installing older Limine versions."
        exit 0
        ;;
    *)
        echo "Unknown option: $1" >&2
        exit 1
        ;;
    esac
    shift
done

# 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
check_uefi || {
    echo "The system is not using EFI."
    exit 0
}

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

# Function to add Limine to the EFI boot entries
add_efi_entry() {
    local esp_path="$1"
    local boot_label="$2"
    local efi_file_path="$3"

    # 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

    # 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 device: /dev/nvmeXnYpZ -> disk=/dev/nvmeXnY, part=Z
        disk="${BASH_REMATCH[1]}"
        part="${BASH_REMATCH[2]#p}" # Remove leading 'p'
    elif [[ "$source" =~ ^(/dev/[a-z]+)([0-9]+)$ ]]; then
        # Other devices: /dev/sda1 -> disk=/dev/sda, part=1
        disk="${BASH_REMATCH[1]}"
        part="${BASH_REMATCH[2]}"
    else
        echo -e "\033[1;31m Error: Failed to parse disk and partition from source '$source'.\033[0m" >&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
        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
}

# Function to update a specific EFI file if needed
update_efi_file() {
    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
    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 fallback Limine 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, only if fallback exists
    if [[ -f "$target_file" ]]; then
        if ! check_limine_downgrade "$source_file" "$target_file" "Limine (fallback)"; then
            return 0
        fi
    fi

    echo "Updating fallback Limine 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/" "${ESP_PATH}/EFI/Limine/"; then
    echo -e "\033[1;31mError: Failed to create directories in '${ESP_PATH}'.\033[0m" >&2
    exit 1
fi

# Optionally add some bootloaders if they are present on the same ESP
if [[ "${FIND_BOOTLOADERS}" == "yes" ]]; 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}" == "yes" ]] || [[ "${SET_LIMINE_AS_FALLBACK}" == "yes" ]]; then
    update_limine_fallback
fi

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

# Optionally register Limine in UEFI system
if [[ "${EFI_REGISTER}" == "yes" ]]; then
    if add_efi_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
