#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-2.0-only

set -e

args=()
rebuildwithhooks=()
declare -i package=0 extramodules=0
declare -A hooktriggers=(
    ['usr/bin/cryptsetup']='encrypt'
    ['usr/bin/lvm']='lvm2'
    ['usr/bin/mdadm']='mdadm_udev'
    ['usr/bin/pdata_tools']='lvm2'
    ['usr/lib/libcryptsetup.so']='sd-encrypt'
    ['usr/lib/libp11-kit.so']='sd-encrypt-opensc'
    ['usr/lib/libpcsclite.so']='sd-encrypt-opensc'
    ['usr/lib/modprobe.d/']='modconf'
    ['usr/lib/systemd/systemd']='systemd'
    ['usr/lib/systemd/systemd-udevd']='udev'
)

install_preset() {
    local pkgbase="$1" preset="$2"

    if [[ ! -e "$preset" ]]; then
        if [[ -e "$preset.pacsave" ]]; then
            # move the pacsave to the template
            mv -- "${preset}.pacsave" "$preset"
        else
            # create the preset from the template
            sed "s|%PKGBASE%|${pkgbase}|g" /usr/share/mkinitcpio/hook.preset \
                | install -Dm644 /dev/stdin "$preset"
        fi
    fi
}

generate_presets() {
    local preset pkgbase pkgbase_path

    # only make presets for kernels with a pkgbase
    for pkgbase_path in /usr/lib/modules/*/pkgbase; do
        if read -r pkgbase &>/dev/null <"$pkgbase_path"; then
            preset="/etc/mkinitcpio.d/${pkgbase}.preset"
            install_preset "$pkgbase" "$preset"
        fi
    done
}

remove_preset() {
    if [[ -n "$pkgbase" && -e "$preset" ]]; then
        if ! cmp "$preset" <(sed "s|%PKGBASE%|${pkgbase}|g" /usr/share/mkinitcpio/hook.preset) &>/dev/null; then
            if [[ ! -e "$preset.pacsave" ]]; then
                # save the preset as pacsave
                mv -- "$preset" "$preset.pacsave" && return 0
            fi
        else
            # remove the preset
            rm -- "$preset" && return 0
        fi
    fi
}

# is_kernelcopy checks whether a `*_kver` value from a .preset file is a
# filename that the .preset expects us to copy the kernel file to.
is_kernelcopy() {
    local preset_kver="$1"

    # Check that this is a filepath (e.g. "/boot/vmlinuz-${pkgbase}") rather
    # than a version number (e.g. "6.5.2-arch-1").
    #
    # mkinitcpio:resolve_kernver() does this by checking whether the string
    # starts with a "/", so do the same check here.
    if [[ "${preset_kver}" != /* ]]; then
        return 1
    fi

    # Check that this isn't the filepath of the original
    # "/usr/lib/modules/${kver}-${pkgbase}/vmlinuz" file that we'd be copying
    # from... or really any other file that is actually owned by a package.
    if pacman -Qoq -- "${preset_kver}" &>/dev/null; then
        return 1
    fi

    return 0
}

read_preset() {
    local pkgbase="$1" p preset_image preset_uki preset_kver
    local unsorted_filelist=() unsorted_kernellist=()

    if [[ -v PRESETS ]]; then
        for p in "${PRESETS[@]}"; do
            declare -n preset_image="${p}_image" preset_uki="${p}_uki" preset_kver="${p}_kver"

            if [[ -v preset_image ]]; then
                unsorted_filelist+=("${preset_image}")
            elif [[ -v ALL_image ]]; then
                unsorted_filelist+=("${ALL_image}")
            fi
            if [[ -v preset_uki ]]; then
                unsorted_filelist+=("${preset_uki}")
            elif [[ -v ALL_uki ]]; then
                unsorted_filelist+=("${ALL_uki}")
            fi
            if [[ -v preset_kver ]]; then
                if is_kernelcopy "${preset_kver}"; then
                    unsorted_filelist+=("${preset_kver}")
                fi
                unsorted_kernellist+=("${preset_kver}")
            elif [[ -v ALL_kver ]]; then
                if is_kernelcopy "${ALL_kver}"; then
                    unsorted_filelist+=("${ALL_kver}")
                fi
                unsorted_kernellist+=("${ALL_kver}")
            fi
        done
    else
        unsorted_filelist+=("/boot/vmlinuz-${pkgbase}"
            "/boot/initramfs-${pkgbase}.img"
            "/boot/initramfs-${pkgbase}-fallback.img"
            "/efi/EFI/Linux/arch-${pkgbase}.efi"
            "/efi/EFI/Linux/arch-${pkgbase}-fallback.efi")
        unsorted_kernellist+=("/boot/vmlinuz-${pkgbase}")
    fi

    # Deduplicate file lists
    filelist=()
    while IFS='' read -r -d '' arrayelement; do
        filelist+=("$arrayelement")
    done < <(printf '%s\0' "${unsorted_filelist[@]}" | LC_ALL=C.UTF-8 sort -uz)
    kernellist=()
    while IFS='' read -r -d '' arrayelement; do
        kernellist+=("$arrayelement")
    done < <(printf '%s\0' "${unsorted_kernellist[@]}" | LC_ALL=C.UTF-8 sort -uz)
}

install_kernel() {
    local pkgbase="$1"
    local kernel preset="/etc/mkinitcpio.d/${pkgbase}.preset"

    install_preset "$pkgbase" "$preset"
    (
        # source the preset to get the kernel and image locations
        # shellcheck disable=SC1090
        [[ -s "$preset" ]] && . "$preset"
        read_preset "$pkgbase"

        # always install the kernel
        for kernel in "${kernellist[@]}"; do
            if is_kernelcopy "${kernel}"; then
                install -Dm644 -- "${line}" "${kernel}"
            fi
        done
    )

    add_pkgbase_to_args "$pkgbase"
}

remove_kernel() {
    local pkgbase="$1"
    local preset="/etc/mkinitcpio.d/${pkgbase}.preset"

    # subshell to avoid namespace pollution
    (
        # source the preset to get the kernel and image locations
        # shellcheck disable=SC1090
        [[ -s "$preset" ]] && . "$preset"
        read_preset "$pkgbase"

        # access all the files to trigger any potential automounts
        stat -- /boot/ /efi/ "${filelist[@]}" &>/dev/null || :

        # remove the actual kernel and images for the package being removed
        rm -f -- "${filelist[@]}"
    ) || return
    # remove the preset
    remove_preset "$pkgbase" "$preset"
}

add_pkgbase_to_args() {
    local pkgbase="$1"

    # do not add the same pkgbase twice
    # shellcheck disable=SC2076
    if [[ " ${args[*]} " =~ " -p ${pkgbase} " ]]; then
        return 0
    fi

    # compound args for each kernel
    args+=(-p "$pkgbase")
}

add_presets_with_hooks() {
    local -a rebuild_hooks=("$@") rebuild_presets=()
    local PRESETS HOOKS hook preset_name preset_file preset_config confd_files pkgbase

    (( ${#rebuild_hooks[@]} )) || return 0

    for preset_file in $(compgen -G /etc/mkinitcpio.d/'*.preset'); do
        [[ -e "$preset_file" ]] || continue
        pkgbase="${preset_file##/etc/mkinitcpio.d/}"
        pkgbase="${pkgbase%.preset}"

        (
            # shellcheck disable=SC1090
            . "$preset_file" &>/dev/null
            [[ -v PRESETS ]] || exit 1
            for preset_name in "${PRESETS[@]}"; do
                preset_config="${preset_name}_config"
                (
                    if [[ -n "${!preset_config:-$ALL_config}" ]]; then
                        # shellcheck disable=SC1090
                        . "${!preset_config:-$ALL_config}" &>/dev/null
                    else
                        if [[ -d /etc/mkinitcpio.conf.d ]]; then
                            mapfile -d '' confd_files < <(LC_ALL=C.UTF-8 find /etc/mkinitcpio.conf.d -maxdepth 1 -xtype f -name '*.conf' -print0 | LC_ALL=C.UTF-8 sort -zVu)
                        fi
                        # shellcheck disable=SC1091 source=mkinitcpio.conf
                        source <(cat /etc/mkinitcpio.conf "${confd_files[@]}") &>/dev/null
                    fi

                    [[ -v HOOKS ]] || exit 1
                    for hook in "${rebuild_hooks[@]}"; do
                        # shellcheck disable=SC2076
                        [[ " ${HOOKS[*]} " =~ " ${hook} " ]] && exit 0
                    done
                    exit 1
                ) && exit 0
            done
            exit 1
        ) && rebuild_presets+=("$pkgbase")

    done

    (( ${#rebuild_presets[@]} )) || return 0
    for pkgbase in "${rebuild_presets[@]}"; do
        add_pkgbase_to_args "$pkgbase"
    done
}

while read -r line; do
    if [[ "${line%%-git}" == "mkinitcpio" && "$1" == "install" ]]; then
        # generate presets for each kernel on package install
        generate_presets
        package=1
        continue
    fi

    if [[ "$line" != */vmlinuz && "$line" != 'usr/lib/modules/'*'/extramodules/' ]]; then
        if ! [[ -v hooktriggers["$line"] ]]; then
            # triggers when it's a change to usr/lib/initcpio/*
            package=1
        else
            rebuildwithhooks+=("${hooktriggers["$line"]}")
        fi
        continue
    fi

    if [[ "$line" == 'usr/lib/modules/'*'/extramodules/' ]]; then
        line="${line%/extramodules/}"
        extramodules=1
    else
        extramodules=0
    fi

    if ! read -r pkgbase &>/dev/null <"${line%/vmlinuz}/pkgbase"; then
        # if the kernel has no pkgbase, we skip it
        continue
    fi

    case "$1" in
        install)
            if (( extramodules )); then
                add_pkgbase_to_args "$pkgbase"
                continue
            fi
            install_kernel "$pkgbase"
            ;;
        remove) remove_kernel "$pkgbase" ;;
    esac
done

if (( package )) && compgen -G /etc/mkinitcpio.d/"*.preset" > /dev/null; then
    case "$1" in
        install)
            # change to use all presets
            args=(-P)
            ;;
        remove)
            shopt -s nullglob
            for preset in /etc/mkinitcpio.d/*.preset; do
                pkgbase=${preset##*/}
                pkgbase=${pkgbase%.preset}
                remove_preset "$pkgbase" "$preset"
            done
            shopt -u nullglob
            ;;
    esac
fi

if (( ! package && ${#rebuildwithhooks} )) && compgen -G /etc/mkinitcpio.d/'*.preset' > /dev/null; then
    add_presets_with_hooks "${rebuildwithhooks[@]}"
fi

if [[ "$1" == "install" ]] && (( ${#args[@]} )); then
    mkinitcpio "${args[@]}"
fi
