#!/bin/bash

# Whatever comes in on file descriptor (FD) 3 gets redirected to where file
# descriptor 1 is pointing. File descriptor 1 points to stdout so when we
# output-redirect something into FD 3 it shows up on stdout. We can use this
# to produce arbitrary logging output inside a subshell like so:
#
# function my_func () {
#     some_command "${1:?}"
#     >&3 echo 'A log message'
# }
#
# var="$(my_func arg_1)"
#
# Here "${var}" will only capture the output of some_command "${1:?}". It
# will not capture 'echo' which will instead show up on our stdout/FD 1.
exec 3>&1

# https://unix.stackexchange.com/a/48550
set -E
trap '[ "$?" -ne 77 ] || exit 77' ERR

function get_partitions () {
    declare partitions_json
    partitions_json="$(lsblk --output PATH,PARTTYPE --json --tree)" || return 1
    printf -- '%s' "${partitions_json}"
    return 0
}

function get_parts () {
    local zfs_install_drive
    declare parttype parts
    parttype="${1:?}"
    zfs_install_drive="${2:-}"
    case "${parttype}" in
        zfs)
            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | .children[] | select(.parttype=="6a85cf4d-1dd2-11b2-99a6-080020736631") | .path')" || exit 1
            ;;
        efi)
            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.path=="'"${zfs_install_drive:?}"'") | .children[] | select(.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .path')" || exit 1
            ;;
        *)
            >&3 printf -- '%s\n' 'Unknown partition type '"'"'"${parttype}"'"'"', exiting ...'
            exit 77
            ;;
    esac
    printf -- '%s' "${parts}"
    return 0
}

function we_have_exactly_one_part () {
    local parttype parts_list parts_count
    parttype="${1:?}"
    parts_list="${2:?}"
    parts_count="$(wc -l <<<"${parts_list}")"
    if [[ "$(wc -c <<<"${parts_list}")" -gt '1' ]]; then
        case "${parts_count}" in
            0)
                >&3 printf -- '%s\n' 'No '"${parttype^^}"' partition found. Exiting ...'
                exit 77
                ;;
            1)
                return 0
                ;;
            *)
                return 1
                ;;
        esac
        >&3 printf -- '%s\n' 'Partition list does not look valid. Cowardly exiting ...'
        exit 77
    fi
}

function select_part () {
    local parts enriched_parts enriched_parts_count part_number part_type zfs_install_drive
    declare part
    part_type="${1:?}" # 'efi' or 'zfs'
    zfs_install_drive="${2:-}"
    if [[ "${zfs_install_drive}" ]]; then
        # This is intended to find correct EFI partition
        parts="$(get_parts "${part_type}" "${zfs_install_drive}")"
    else
        parts="$(get_parts "${part_type}")"
    fi

    if [[ ! "${parts}" ]]; then
        case "${part_type}" in
            efi)
                part_type_human_readable='EFI system partition (ESP) with partition type code EF00'
                ;;
            zfs)
                part_type_human_readable='ZFS zpool partition with partition type code BF00'
                ;;
        esac
        >&3 printf -- '%s\n' \
            'It looks as if there is no '"${part_type_human_readable}" \
            'on any of the disks. Is this a chroot? Exiting ...'
        exit 77
    fi

    if we_have_exactly_one_part "${part_type}" "${parts}"; then
        part="${parts}"
    else
        >&3 printf -- '%s\n' 'More than one '"${part_type^^}"' partition to pick for installation. Cowardly exiting ...'
        exit 77
    fi
    printf -- '%s' "${part}"
    return 0
}

function get_part_parent () {
    local child_partition parent_partition
    child_partition="${1:?}"
    parent_partition="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.children[].path=="'"${child_partition:?}"'") | .path')"
    printf -- '%s' "${parent_partition}"
    return 0
}

function get_disks_with_one_efipart () {
    local disks_with_one_efipart
    # Find disks that have exactly one EFI partition and where that EFI
    # partition is partition number 1. We expect exactly one disk to meet
    # these criteria. Anything else and we bail.
    disks_with_one_efipart="$(lsblk --output PATH,PARTTYPE --json --tree | jq --raw-output '.[][] | select(.children | length > 0) | select( any (.children[]; (.path | test("^[^[:digit:]]+1")) and (.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b")) and ([select(.children[].parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b")] | length == 1) ) | .path')"
    if [[ "$(wc -l <<<"${disks_with_one_efipart}")" -eq '1' ]] && [[ "$(wc -c <<<"${disks_with_one_efipart}")" -gt '1' ]]; then
        printf -- '%s' "${disks_with_one_efipart}"
        return 0
    fi
    return 1
}

function set_new_uefi_boot_entries () {
    declare -a uefi_images
    mapfile -t uefi_images < \
        <(find '/efi/EFI/ZBM' -type f -iname '*.efi' -print0 | \
        xargs -0 --no-run-if-empty --max-args '1' stat -c '%Y %n' | \
        sort -V | \
        awk '{print $2}')
    zpool_drive="$(select_part 'zfs')"
    zfs_parent="$(get_part_parent "${zpool_drive:?}")"
    efi_drive="${zfs_parent}"

    if efibootmgr | grep -Piq -- 'ZFSBootMenu'; then
        local -a old_uefi_entries
        mapfile -t old_uefi_entries < \
            <(efibootmgr | \
            grep -Pio -- '(?<=^Boot)[^\*[:space:]]+(?=\*? ZFSBootMenu)')
        for old_uefi_entry in "${old_uefi_entries[@]}"; do
            efibootmgr --bootnum "${old_uefi_entry}" --delete-bootnum &>/dev/null && {
                >&3 printf -- '%s\n' \
                    'EFI boot entry '"${old_uefi_entry}"' deleted.'
            }
        done
    fi

    if ! efibootmgr | grep -Piq -- 'ZFSBootMenu'; then
        local efi_disks_list
        efi_disks_list="$(get_disks_with_one_efipart)"
        if grep -Piq -- '^'"${efi_drive}"'$' <<<"${efi_disks_list}"; then
            for uefi_image in "${uefi_images[@]}"; do
                uefi_image_version="$(basename "${uefi_image%%.EFI}")"
                uefi_image_inverted="${uefi_image#/efi}"
                uefi_image_inverted="${uefi_image_inverted//\//\\}"
                efibootmgr --disk "${efi_drive}" \
                    --part 1 \
                    --create \
                    --label 'ZFSBootMenu '"${uefi_image_version}" \
                    --loader "${uefi_image_inverted}" &>/dev/null && {
                        >&3 printf -- '%s\n' \
                            'EFI boot entry ZFSBootMenu '"${uefi_image_version}"' added.'
            }
            done
        fi
    fi
}

set_new_uefi_boot_entries