#!/bin/bash

export LANG=C

COMMAND="${1:?}"
KERNEL_VERSION="${2:?}"
BOOT_DIR_ABS="${3%/*}/0-rescue"
KERNEL_IMAGE="$4"

# If the initrd was provided on the kernel command line, we shouldn't generate our own.
if [[ "$COMMAND" = "add" && "$#" -gt 4 ]]; then
    exit 0
fi

# Do not attempt to create initramfs if the supplied image is already a UKI
if [[ "$KERNEL_INSTALL_IMAGE_TYPE" = "uki" ]]; then
    exit 0
fi

if [[ "$KERNEL_INSTALL_UKI_GENERATOR" = "dracut" ]]; then
    # Rescue images currently not compatible with UKIs
    exit 0
elif [[ "${KERNEL_INSTALL_INITRD_GENERATOR:-dracut}" != "dracut" ]]; then
    # We are not the initrd generator
    exit 0
fi

dropindirs_sort() {
    suffix=$1
    shift
    args=("$@")
    files=$(
        while (($# > 0)); do
            for i in "${1}"/*"${suffix}"; do
                [[ -f $i ]] && echo "${i##*/}"
            done
            shift
        done | sort -Vu
    )

    for f in $files; do
        for d in "${args[@]}"; do
            if [[ -f "$d/$f" ]]; then
                echo "$d/$f"
                continue 2
            fi
        done
    done
}

if [[ -f /etc/os-release ]]; then
    . /etc/os-release
elif [[ -f /usr/lib/os-release ]]; then
    . /usr/lib/os-release
fi

[[ -n $PRETTY_NAME ]] || PRETTY_NAME="Linux $KERNEL_VERSION"

if [[ ${KERNEL_INSTALL_MACHINE_ID+x} ]]; then
    MACHINE_ID=$KERNEL_INSTALL_MACHINE_ID
elif [[ -f /etc/machine-id ]]; then
    read -r MACHINE_ID < /etc/machine-id
fi

if ! [[ $MACHINE_ID ]]; then
    exit 0
fi

if [ -n "$KERNEL_INSTALL_CONF_ROOT" ]; then
    if [ -f "$KERNEL_INSTALL_CONF_ROOT/cmdline" ]; then
        read -r -d '' -a BOOT_OPTIONS < "$KERNEL_INSTALL_CONF_ROOT/cmdline"
    fi
elif [[ -f /etc/kernel/cmdline ]]; then
    read -r -d '' -a BOOT_OPTIONS < /etc/kernel/cmdline
elif [[ -f /usr/lib/kernel/cmdline ]]; then
    read -r -d '' -a BOOT_OPTIONS < /usr/lib/kernel/cmdline
else
    declare -a BOOT_OPTIONS

    read -r -d '' -a line < /proc/cmdline
    for i in "${line[@]}"; do
        [[ ${i#initrd=*} != "$i" ]] && continue
        BOOT_OPTIONS+=("$i")
    done
fi

if [[ -d ${BOOT_DIR_ABS%/*} ]]; then
    BOOT_DIR="/${MACHINE_ID}/0-rescue"
    BOOT_ROOT=${BOOT_DIR_ABS%"$BOOT_DIR"}
    LOADER_ENTRY="$BOOT_ROOT/loader/entries/${MACHINE_ID}-0-rescue.conf"
    KERNEL="linux"
    INITRD="initrd"
else
    BLS_DIR="/boot/loader/entries"
    BOOT_DIR_ABS="/boot"
    LOADER_ENTRY="$BLS_DIR/${MACHINE_ID}-0-rescue.conf"
    KERNEL="vmlinuz-0-rescue-${MACHINE_ID}"
    INITRD="initramfs-0-rescue-${MACHINE_ID}.img"
fi

ret=0

case "$COMMAND" in
    add)
        if [[ -f $LOADER_ENTRY ]] && [[ -f "$BOOT_DIR_ABS/$KERNEL" ]] \
            && [[ -f "$BOOT_DIR_ABS/$INITRD" ]]; then
            [[ $KERNEL_INSTALL_VERBOSE == 1 ]] \
                && echo "Skipping, there is already a rescue image generated with the same input parameters"
            exit 0
        fi

        # source our config dir
        for f in $(dropindirs_sort ".conf" "/etc/dracut.conf.d" "/usr/lib/dracut/dracut.conf.d"); do
            if [[ -e $f ]]; then
                # shellcheck disable=SC1090
                . "$f"
            fi
        done

        if [[ $dracut_rescue_image != "yes" ]]; then
            [[ $KERNEL_INSTALL_VERBOSE == 1 ]] \
                && echo "Skipping, 'dracut_rescue_image' not set to 'yes' in any dracut configuration file"
            exit 0
        fi

        [[ -d $BOOT_DIR_ABS ]] || mkdir -p "$BOOT_DIR_ABS"

        if ! cp --reflink=auto "$KERNEL_IMAGE" "$BOOT_DIR_ABS/$KERNEL"; then
            echo "Can't copy '$KERNEL_IMAGE to '$BOOT_DIR_ABS/$KERNEL'!" >&2
        fi

        if [[ ! -f "$BOOT_DIR_ABS/$INITRD" ]]; then
            # shellcheck disable=SC2046
            dracut -f --no-hostonly --no-uefi \
                -a "rescue" \
                $([[ $KERNEL_INSTALL_VERBOSE == 1 ]] && echo --verbose) \
                --kver "$KERNEL_VERSION" \
                "$BOOT_DIR_ABS/$INITRD"
            ((ret += $?))
        fi

        [[ $KERNEL_INSTALL_VERBOSE == 1 ]] && echo "Creating $LOADER_ENTRY"
        if [[ ${BOOT_DIR_ABS} != "/boot" ]]; then
            {
                echo "title      $PRETTY_NAME - Rescue Image"
                echo "version    $KERNEL_VERSION"
                echo "machine-id $MACHINE_ID"
                echo "options    ${BOOT_OPTIONS[*]} rd.auto=1"
                echo "linux      $BOOT_DIR/linux"
                echo "initrd     $BOOT_DIR/initrd"
            } > "$LOADER_ENTRY"
        else
            if [[ -e "${BLS_DIR}/${MACHINE_ID}-${KERNEL_VERSION}.conf" ]]; then
                cp -aT "${BLS_DIR}/${MACHINE_ID}-${KERNEL_VERSION}.conf" "$LOADER_ENTRY"
            else
                cp -aT "${KERNEL_IMAGE%/*}/bls.conf" "$LOADER_ENTRY"
            fi
            sed -i "s/${KERNEL_VERSION}/0-rescue-${MACHINE_ID}/" "$LOADER_ENTRY"
        fi

        ((ret += $?))
        ;;

    remove)
        exit 0
        ;;

esac

exit $ret