#!/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

declare this_script_url
this_script_url="${SCRIPT_URL:?}"
postconf_hook="$(dirname "${this_script_url}")"'/zbm_set_new_uefi_boot_entries.sh'

declare zpool_name zfs_arch_dataset_name
zpool_name='zpool'
zfs_arch_dataset_name='archlinux'

declare -A partition_types
partition_types[gpt_zfs]='6a85cf4d-1dd2-11b2-99a6-080020736631'
partition_types[gpt_efi]='c12a7328-f81f-11d2-ba4b-00a0c93ec93b'
partition_types[mbr_zfs]='0xbf'
partition_types[mbr_boot]='0x83'

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

declare zpool_drive efi_drive boot_drive part_schema

function we_are_changerooted () {
    if [ "$(stat -c '%d:%i' '/')" != "$(stat -c '%d:%i' '/proc/1/root/.')" ]; then
        return 0
    else
        return 1
    fi
}

function no_kernel_update_in_iso () {
    #1.1
    sed -ri -e 's'$'\x1''#(IgnorePkg)[^\r\n\f]+'$'\x1''\1 = linux linux-headers'$'\x1''g' /etc/pacman.conf
}

function set_ntp () {
    #1.2
    timedatectl set-ntp true
}

function resize_cow_space () {
    #1.3
    mount -o remount,size='50%' /run/archiso/cowspace
}

function update_pacman_db () {
    #1.4
    printf -- '%s\n' 'Refreshing mirror list ...'
    systemctl start reflector
    # In an ISO and for the minimal number of packages we need we do not
    # care about partial upgrades
    pacman -Syyuu --noconfirm
}

function install_pkgs () {
    #1.5
    printf -- '%s\n' 'Installing packages ...'
    pacman -S --needed --noconfirm "${@}"
}

function install_zfs () {
    #1.6
    declare reset_colors='\033[0m'
    curl -s 'https://raw.githubusercontent.com/eoli3n/archiso-zfs/master/init' | bash
    printf -- "${reset_colors}"
}

function uefi_or_bios () {
    #1.7
    local part_count_linux part_count_efi
    # Select disks with at least one partition. Among them count how many
    # with the given partition type code we have.
    part_count_linux="$(get_partitions | jq --raw-output '[.[][] | select(.children | length > 0) | .children[] | select(.parttype=="'"${partition_types[mbr_boot]}"'") | .path] | length')"
    part_count_efi="$(get_partitions | jq --raw-output '[.[][] | select(.children | length > 0) | .children[] | select(.parttype=="'"${partition_types[gpt_efi]}"'") | .path] | length')"
    if [[ "${part_count_linux}" -eq '1' ]] && [[ "${part_count_efi}" -eq '0' ]]; then
        part_schema='mbr'
        >&3 printf -- '%s\n' \
            'Treating this as an MBR/legacy BIOS machine ...'
    elif [[ "${part_count_linux}" -eq '0' ]] && [[ "${part_count_efi}" -eq '1' ]]; then
        part_schema='gpt'
        >&3 printf -- '%s\n' \
            'Treating this as a GPT/UEFI machine ...'
    else
        >&3 printf -- '%s\n' \
            'We are seeing partitions as follows:' \
            '- Partition type code '"${partition_types[mbr_boot]}"': '"${part_count_linux}" \
            '- Partition type code '"${partition_types[gpt_efi]}"': '"${part_count_efi}" \
            '' \
            'We are instead expecting either 1 and 0 indicating a legacy' \
            'BIOS setup or 0 and 1 indicating UEFI. '"${part_count_linux}"' and '"${part_count_efi}"' is' \
            'neither. Cowardly exiting ...'
        exit 77
    fi
}

function get_partitions () {
    declare partitions_json
    partitions_json="$(lsblk --output PATH,PARTTYPE --json --tree)" || return 1
    printf -- '%s' "${partitions_json}"
    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_parts () {
    local zfs_install_drive part_type_code
    declare parttype parts
    parttype="${1:?}"
    zfs_install_drive="${2:-}"
    case "${parttype}" in
        zfs)
            if [[ "${part_schema}" = 'mbr' ]]; then
                part_type_code="${partition_types[mbr_zfs]}"
            else
                part_type_code="${partition_types[gpt_zfs]}"
            fi
            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | .children[] | select(.parttype=="'"${part_type_code}"'") | .path')" || exit 1
            ;;
        boot)
            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.path=="'"${zfs_install_drive:?}"'") | .children[] | select(.parttype=="'"${partition_types[mbr_boot]}"'") | .path')" || exit 1
            ;;
        efi)
            parts="$(get_partitions | jq --raw-output '.[][] | select(.children | length > 0) | select(.path=="'"${zfs_install_drive:?}"'") | .children[] | select(.parttype=="'"${partition_types[gpt_efi]}"'") | .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 get_drive_id () {
    local drive_id_list drive_id_single
    drive_id_list="$(find -L /dev/disk/by-partuuid -samefile "${1:?}" | sort)"
    drive_id_single="$(head -n1 <<<"${drive_id_list}")"
    if [[ "$(wc -l <<<"${drive_id_single}")" -eq '1' ]] && [[ "$(wc -c <<<"${drive_id_single}")" -gt '1' ]]; then
        printf -- '%s' "${drive_id_single}"
        return 0
    fi
    >&3 printf -- '%s\n' 'No '"'${1:?}'"' partition entry in /dev/disk/by-partuuid, exiting ...'
    exit 77
}

function select_part () {
    local parts enriched_parts enriched_parts_count part_number part_type zfs_install_drive part_type_code specific_part_type
    declare part
    part_type="${1:?}" # 'boot' or 'zfs'
    zfs_install_drive="${2:-}"

    if [[ "${part_type}" = 'boot' ]]; then
        if [[ "${part_schema}" = 'mbr' ]]; then
            specific_part_type='boot'
        else
            specific_part_type='efi'
        fi
    else
        specific_part_type="${part_type}"
    fi

    if [[ "${zfs_install_drive}" ]]; then
        # This is intended to find correct boot/EFI partition
        parts="$(get_parts "${specific_part_type}" "${zfs_install_drive}")"
    else
        parts="$(get_parts "${specific_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)
                if [[ "${part_schema}" = 'mbr' ]]; then
                    part_type_code='bf'
                else
                    part_type_code='BF00'
                fi
                part_type_human_readable='ZFS zpool partition with partition type code '"${part_type_code}"
                ;;
        esac
        >&3 printf -- '%s\n' \
            'It looks as if there is no '"${part_type_human_readable}" \
            'on any of the disks. Did you correctly partition a disk before starting?' \
            'Check https://quico.space/quico-os-setup/arch-zbm#prep. 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 no_zpool_exists () {
    declare zpool_list
    zpool_list="$(zpool list -H)"
    [[ "$(wc -l <<<"${zpool_list}")" -le '1' ]] && [[ "$(wc -c <<<"${zpool_list}")" -le '1' ]] && return 0
    return 1
}

function set_zpool_password () {
    # May or may not have a newline at the end, ZFS doesn't care
    printf -- '%s' 'password' > '/etc/zfs/'"${zpool_name}"'.key'
    chmod '000' '/etc/zfs/'"${zpool_name}"'.key'
}

function import_pool () {
    zpool import -d '/dev/disk/by-partuuid' -R '/mnt' "${zpool_name}" -N -f
    [[ ! "${ARCHZBM_ZFSPROPS_NO_ENCRYPTION}" ]] && zfs load-key "${zpool_name}"
}

function create_pool () {
    # Create a temporary pool that is not cached
    #
    # Add zfsprops 'compression' unless environment variable
    # ARCHZBM_ZFSPROPS_NO_COMPRESSION is set to any value.
    #
    # Add zfsprops 'encryption' along with 'keyformat' and a 'keylocation'
    # unless environment variable ARCHZBM_ZFSPROPS_NO_ENCRYPTION is set to
    # any value.
    zpool create -f                                                     \
                 -o 'ashift=12'                                         \
                 -o 'autotrim=on'                                       \
                 -O 'acltype=posix'                                     \
                 $([[ ! "${ARCHZBM_ZFSPROPS_NO_COMPRESSION}" ]] &&      \
                    printf -- '%s '                                     \
                '-O compression=on')                                    \
                 -O 'relatime=on'                                       \
                 -O 'xattr=sa'                                          \
                 $([[ ! "${ARCHZBM_ZFSPROPS_NO_ENCRYPTION}" ]] &&       \
                    printf -- '%s '                                     \
                '-O encryption=on'                                      \
                '-O keyformat=passphrase'                               \
                '-O keylocation=file:///etc/zfs/'"${zpool_name}"'.key') \
                 -O 'normalization=formD'                               \
                 -O 'mountpoint=none'                                   \
                 -O 'canmount=off'                                      \
                 -O 'devices=off'                                       \
                 -R '/mnt'                                              \
                 "${zpool_name}" "${1:?}"
}

function create_root_dataset () {
    zfs create -o 'mountpoint=none' -o 'canmount=off' "${zpool_name}"'/root'
    # zfs set org.zfsbootmenu:commandline="ro quiet" "${zpool_name}"'/root'
    zfs set org.zfsbootmenu:commandline="ro" "${zpool_name}"'/root'
}

function create_system_dataset () {
    zfs create -o 'mountpoint=/' -o 'canmount=noauto' "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
    zgenhostid
    zpool set bootfs="${zpool_name}"'/root/'"${zfs_arch_dataset_name}" "${zpool_name}"
    zfs mount "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
}

function create_home_dataset () {
    zfs create -o 'mountpoint=/' -o 'canmount=off' "${zpool_name}"'/data'
    zfs create "${zpool_name}"'/data/home'
}

function export_pool () {
    zpool export "${zpool_name}"
}

function setup_zpool () {
    #1.8
    local drive_by_id
    zpool_drive="$(select_part 'zfs')"
    drive_by_id="$(get_drive_id "${zpool_drive}")"

    [[ ! "${ARCHZBM_ZFSPROPS_NO_ENCRYPTION}" ]] && set_zpool_password
    if no_zpool_exists; then
        create_pool "${drive_by_id}"
        create_root_dataset
        create_system_dataset
        create_home_dataset
        export_pool
        import_pool
    else
        >&3 printf -- '%s\n' 'A zpool already exists, that is unexpected. Cowardly exiting 1 ...'
        exit 1
    fi
}

function mount_system () {
    #1.9
    zfs mount "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
    zfs mount -a

    local zfs_parent efi_part boot_part
    zfs_parent="$(get_part_parent "${zpool_drive:?}")"

    if [[ "${part_schema}" = 'mbr' ]]; then
        boot_drive="${zfs_parent}"
        boot_part="$(select_part 'boot' "${zfs_parent:?}")"
        mkdir -p '/mnt/boot/syslinux'
        mount "${boot_part}" '/mnt/boot/syslinux'
    else
        efi_drive="${zfs_parent}"
        efi_part="$(select_part 'boot' "${zfs_parent:?}")"
        mkdir -p '/mnt/efi'
        mount "${efi_part}" '/mnt/efi'
    fi
}

function copy_zpool_cache () {
    #1.10
    mkdir -p '/mnt/etc/zfs'
    zpool set 'cachefile=/etc/zfs/'"${zpool_name}"'.cache' "${zpool_name}"
}

function pacman_dont_check_space () {
    # See pacman bug comment
    # https://bugs.archlinux.org/task/45070#comment142712
    #
    # When we pacstrap onto ZFS pacman incorrectly calculates and
    # overestimates required disk space. We instead assume an installation
    # gets done with at least a 10 GiB drive which is plenty. Skip pacman's
    # size check.
    #
    # We're setting this in Arch Linux ISO CD while we install proper Arch.
    # No need to revert this later as it is ephemeral anyway.
    sed -ri -e 's'$'\x1''^.*?(CheckSpace)([^\r\n\f]*)'$'\x1''#\1\2'$'\x1''g' '/etc/pacman.conf'
}

function install_archlinux () {
    #1.11
    pacman_dl_parallel
    pacman_dont_check_space
    pacstrap /mnt              \
        base                   \
        base-devel             \
        linux                  \
        linux-headers          \
        linux-firmware         \
        amd-ucode              \
        efibootmgr             \
        vim                    \
        git                    \
        iwd                    \
        networkmanager         \
        network-manager-applet \
        dialog                 \
        os-prober              \
        reflector              \
        bluez                  \
        bluez-utils            \
        man-db                 \
        xdg-utils              \
        xdg-user-dirs
}

function gen_fstab () {
    #1.12
    genfstab -U /mnt | grep -v "${zpool_name}" | tr -s '\n' | sed -r -e 's/\/mnt//' -e '/./,$!d' > '/mnt/etc/fstab'
}

function configure_hosts_file () {
    cat > '/mnt/etc/hosts' <<EOF
#<ip-address>   <hostname.domain.org>   <hostname>
127.0.0.1       localhost               ${1}
::1             localhost               ${1}
EOF
}

function set_hostname () {
    #1.13
    declare new_hostname
    install_pkgs 'pwgen'
    new_hostname="$(pwgen --no-numerals --no-capitalize --ambiguous 8)"
    printf -- '%s\n' "${new_hostname}" > '/mnt/etc/hostname'
    configure_hosts_file "${new_hostname}"
}

function set_locale () {
    #1.14
    printf -- '%s\n' \
        'KEYMAP=de-latin1' \
        'FONT=Lat2-Terminus16' \
        'FONT_MAP=8859-1' \
        > '/mnt/etc/vconsole.conf'
    sed -ri -e 's'$'\x1''^(#)(en_US.UTF-8[^\r\n\f]*)'$'\x1''\2'$'\x1''g' '/mnt/etc/locale.gen'
    printf -- '%s\n' 'LANG=en_US.UTF-8' > '/mnt/etc/locale.conf'
}

function add_zfs_hook_to_initramfs () {
    #1.15
    # Add zfs hook, remove fsck hook from initramfs.
    sed -ri \
        -e 's'$'\x1''(HOOKS=)(.*?[\(| ])(filesystems)([\)| ][^\r\n\f]*)'$'\x1''\1\2zfs \3\4'$'\x1''g' \
        -e 's'$'\x1''((\()(fsck)(\)))'$'\x1''\2\4'$'\x1''g' \
        -e 's'$'\x1''(([[:space:]]+)(fsck)|(fsck)([[:space:]]+))'$'\x1'''$'\x1''g' \
        '/mnt/etc/mkinitcpio.conf'
    # Also unless encryption's unwanted add plain text key file into
    # initramfs since it's living inside an encrypted pool anyway.
    [[ ! "${ARCHZBM_ZFSPROPS_NO_ENCRYPTION}" ]] && sed -ri \
        -e 's'$'\x1''^(FILES=)[^\r\n\f]*'$'\x1''\1(/etc/zfs/'"${zpool_name}"'.key)'$'\x1''g' \
        '/mnt/etc/mkinitcpio.conf'
}

function set_initramfs_build_list () {
    #1.16
    # No need to build fallback initramfs, our new fallback is ZFS snapshots
    sed -ri \
        -e '/^#/d' \
        -e '/^$/d' \
        -e '/^fallback/d' \
        -e 's'$'\x1''^(PRESETS=)[^\r\n\f]*'$'\x1''\1('"'"'default'"'"')'$'\x1''g' \
        '/mnt/etc/mkinitcpio.d/linux.preset'

    # Remove any existing fallback initramfs files
    find '/mnt/boot' -type f -regextype posix-extended -iregex '^/mnt/boot/initramfs-.*?-fallback.img' -delete
}

function add_zfs_files_to_new_os () {
    #1.17
    for zfs_file in '/etc/hostid' '/etc/zfs/zpool.cache' $([[ ! "${ARCHZBM_ZFSPROPS_NO_ENCRYPTION}" ]] && printf -- '%s' '/etc/zfs/'"${zpool_name}"'.key'); do
        rsync -av --itemize-changes {'','/mnt'}"${zfs_file}"
    done
}

function create_unpriv_user () {
    account_name="${1:?}"
    full_name="${2:-${account_name}}"
    authorized_keys_abs_path='/home/'"${account_name}"'/.ssh/authorized_keys'
    useradd --create-home --shell '/bin/bash' --comment 'Personal account of '"${full_name}" --user-group "${account_name}"
    chfn --full-name "${full_name}" "${account_name}"
    mkdir -p "$(dirname "${authorized_keys_abs_path}")"
    touch "${authorized_keys_abs_path}"
    chown -R "${account_name}"':' '/home/'"${account_name}"; chmod -R 'u=rwX,go=' "$(dirname "${authorized_keys_abs_path}")"
}

function enter_chroot () {
    #2.1
    arch-chroot /mnt /bin/bash -xe <<EOF
curl --silent '${this_script_url}' | bash
EOF
}

function keep_initiramfs_root_only_rw () {
    #2.3
    declare systemd_local_admin_override_path unit_name
    systemd_local_admin_override_path='/etc/systemd/system'
    unit_name='chmod-initramfs'
    path_unit="${systemd_local_admin_override_path%/}"'/'"${unit_name}"'.path'
    service_unit="${systemd_local_admin_override_path%/}"'/'"${unit_name}"'.service'

    cat > "${path_unit}" <<"EOF"
[Unit]
Description=chmod initramfs to be root-read-writable only

[Path]
PathChanged=/boot/initramfs-linux.img

[Install]
WantedBy=multi-user.target
WantedBy=system-update.target
EOF

    cat > "${service_unit}" <<"EOF"
[Unit]
Description=chmod initramfs to be root-read-writable only

[Service]
Type=oneshot
ExecStart=/usr/bin/chmod 600 /boot/initramfs-linux.img
EOF

    systemctl enable "${path_unit}"
}

function pacman_dl_parallel () {
    #2.4
    # We're setting this in Arch Linux ISO CD while we install proper Arch.
    # No need to revert this later as it is ephemeral anyway.
    sed -ri -e 's'$'\x1''^.*?(ParallelDownloads)[^\r\n\f]*'$'\x1''\1 = 20'$'\x1''g' '/etc/pacman.conf'
}

function unleash_makepkg () {
    #2.5
    local path_prefix
    path_prefix='/mnt'
    if we_are_changerooted; then
        path_prefix=''
    fi
    sed -ri \
        -e 's'$'\x1''^(#?(MAKEFLAGS=))[^\r\n\f]*'$'\x1''\2"-j$(nproc --ignore 1)"'$'\x1''g' \
        "${path_prefix}"'/etc/makepkg.conf'
}

function add_motd_getting_started_msg () {
    #2.6
    cat > '/etc/motd' <<"EOF"

####################

GUI basics:

    paru -S xorg plasma-meta kde-applications-meta sddm
    localectl set-x11-keymap de
    useradd --create-home --shell /bin/bash --user-group --groups wheel <user>
    passwd <user>
    systemctl enable --now sddm.service

####################

EOF
}

function get_aur_helper () {
    #2.7
    create_unpriv_user 'build'
    usermod --append --groups 'wheel' 'build'
    printf -- '%s\n' '%wheel ALL=(ALL:ALL) NOPASSWD: ALL' > '/etc/sudoers.d/10-wheel-group-no-passwd-prompt'
    pushd /tmp
    git clone 'https://aur.archlinux.org/paru.git'
    chown -R 'build:' 'paru'
    pushd 'paru'
    sudo --user 'build' makepkg -si --noconfirm
    popd
    rm -rf 'paru'
    popd
}

function paru_install () {
    declare -a paru_install_packages
    [[ "${1}" ]] && while :; do
        case "${1}" in
            -[[:alnum:]]*)
                >&3 printf -- '%s\n' \
                    'Short-form argument '"'${1}'"' not supported for function '"'${FUNCNAME[0]}()'"'. Only known accepted argument' \
                    'is '"'"'--replace-conflicting'"'"' without a value given. Exiting ...'
                exit 77
                ;;
            --replace-conflicting)
                pacman_force_yes='true'
                shift
                continue
                ;;
            --*)
                >&3 printf -- '%s\n' \
                    'Long-form argument '"'${1}'"' not supported for function '"'${FUNCNAME[0]}()'"'. Only known accepted argument' \
                    'is '"'"'--replace-conflicting'"'"' without a value given. Exiting ...'
                exit 77
                ;;
            '')
                # All arguments processed
                break
                ;;
            *)
                paru_install_packages+=("${1}")
                shift
                ;;
            esac
    done || {
        >&3 printf -- '%s\n' \
            'No argument '"'${1}'"' given for function '"'${FUNCNAME[0]}'"'. Exiting ...'
        exit 77
    }
    if [[ "${pacman_force_yes}" ]]; then
        yes 'y' | sudo --user 'build' paru -S "${paru_install_packages[@]}"
        unset -v pacman_force_yes
    else
        sudo --user 'build' paru -S --noconfirm "${paru_install_packages[@]}"
    fi
}

function configure_syslinux () {
    paru_install 'syslinux'
    cp /usr/lib/syslinux/bios/*.c32 /boot/syslinux
    extlinux --install /boot/syslinux
    cat > /boot/syslinux/syslinux.cfg <<EOF
UI menu.c32
PROMPT 0

MENU TITLE ZFSBootMenu
TIMEOUT 100

DEFAULT zfsbootmenu

LABEL zfsbootmenu
  MENU LABEL ZFSBootMenu
  KERNEL /zfsbootmenu/vmlinuz-bootmenu
  INITRD /zfsbootmenu/initramfs-bootmenu.img
  APPEND zfsbootmenu quiet loglevel=0

LABEL zfsbootmenu-backup
  MENU LABEL ZFSBootMenu (Backup)
  KERNEL /zfsbootmenu/vmlinuz-bootmenu-backup
  INITRD /zfsbootmenu/initramfs-bootmenu-backup.img
  APPEND zfsbootmenu quiet loglevel=0
EOF
    dd bs=440 count=1 conv=notrunc if='/usr/lib/syslinux/bios/mbr.bin' of="${boot_drive}"
}

function configure_zfsbootmenu () {
    #2.9
    paru_install 'zfsbootmenu'

    if [[ "${part_schema}" = 'gpt' ]]; then
        mkdir -p '/etc/zfsbootmenu/posthooks.d'
        cat > '/etc/zfsbootmenu/config.yaml' <<EOF
Global:
  ManageImages: true
  BootMountPoint: /efi
  InitCPIO: true
  PostHooksDir: /etc/zfsbootmenu/posthooks.d
Components:
  Enabled: false
EFI:
  ImageDir: /efi/EFI/ZBM
  Versions: false
  Enabled: true
  Stub: /etc/zfsbootmenu/stub-loader.d/linuxx64.efi.stub
Kernel:
  CommandLine: ro loglevel=0 zbm.import_policy=hostid
  Prefix: vmlinuz
EOF
        get_known_good_stub_loader
    else
        configure_syslinux
        cat > '/etc/zfsbootmenu/config.yaml' <<EOF
Global:
  ManageImages: true
  BootMountPoint: /boot/syslinux
  InitCPIO: true
Components:
  Enabled: true
  Versions: false
  ImageDir: /boot/syslinux/zfsbootmenu
Kernel:
  Prefix: vmlinuz
EOF
    fi

# Up here maybe 'ro quiet' instead of 'ro'. This is ZFSBootMenu's kernel
# command line.

    # Assign cmdline for final kernel start. This is our Arch Linx kernel
    # command line.
    zfs set org.zfsbootmenu:commandline='rw nowatchdog rd.vconsole.keymap=de-latin1' "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
}

function get_known_good_stub_loader () {
    local known_good_stub_loader local_stub_loader_abs
    known_good_stub_loader='https://github.com/zbm-dev/zfsbootmenu/raw/master/testing/stubs/linuxx64.efi.stub'
    local local_stub_loader_abs='/etc/zfsbootmenu/stub-loader.d/linuxx64.efi.stub'
    mkdir -p "$(dirname "${local_stub_loader_abs}")"
    curl --silent --location "${known_good_stub_loader}" --output "${local_stub_loader_abs}"
}

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=="'"${partition_types[gpt_efi]}"'")) and ([select(.children[].parttype=="'"${partition_types[gpt_efi]}"'")] | 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 install_os_in_chroot () {
    #2.2
    ### Reinit keyring
    # As keyring is initialized at boot, and copied to the install dir with pacstrap, and ntp is running
    # Time changed after keyring initialization, it leads to malfunction
    # Keyring needs to be reinitialised properly to be able to sign keys.
    rm -rf '/etc/pacman.d/gnupg'
    pacman-key --init
    pacman-key --populate archlinux
    pacman -S archlinux-keyring --noconfirm

    locale-gen
    source /etc/locale.conf

    keep_initiramfs_root_only_rw                #2.3
    pacman_dl_parallel                          #2.4
    unleash_makepkg                             #2.5
    add_motd_getting_started_msg                #2.6
    get_aur_helper                              #2.7
    paru_install --replace-conflicting 'paru-bin'
    paru_install 'zfs-dkms' 'zfs-utils' 'jq'
    hwclock --systohc
    mkinitcpio -P

    # Install ZFSBootMenu image
    configure_zfsbootmenu                       #2.9
    generate-zbm

    # Yes, we do this twice so we immediately get a functional backup file
    generate-zbm
}

function set_root_pw () {
    #3.2
    printf -- '%s\n' 'root:password' | chpasswd --crypt-method 'SHA512' --root '/mnt'
}

function configure_networking () {
    #3.3
    cat > '/mnt/etc/systemd/network/50-wired.network' <<"EOF"
[Match]
Name=en*

[Network]
DHCP=ipv4
IPForward=yes

[DHCP]
UseDNS=yes
RouteMetric=10
EOF
    systemctl enable 'systemd-networkd' --root='/mnt'
    systemctl disable 'systemd-networkd-wait-online' --root='/mnt'
}

function configure_dns () {
    #3.4
    rm '/mnt/etc/resolv.conf'
    ln -s '/run/systemd/resolve/stub-resolv.conf' '/mnt/etc/resolv.conf'

    # Optionally you may want /etc/systemd/network/50-wired.network to use
    # UseDNS=no and hardcode DNS server(s) here:
    # sed -i 's/^#DNS=.*/DNS=1.1.1.1/' /mnt/etc/systemd/resolved.conf
    systemctl enable 'systemd-resolved' --root='/mnt'
}

function configure_reflector () {
    #3.5
    systemctl enable 'reflector.service' 'reflector.timer' --root='/mnt'
}

function configure_zfs () {
    #3.6
    systemctl enable 'zfs-import-cache' 'zfs-mount' 'zfs-import.target' 'zfs.target' --root='/mnt'
}

function configure_zfs_mount_gen () {
    #3.7
    mkdir -p '/mnt/etc/zfs/zfs-list.cache'
    touch '/mnt/etc/zfs/zfs-list.cache/'"${zpool_name}"
    zfs list -H -o name,mountpoint,canmount,atime,relatime,devices,exec,readonly,setuid,nbmand | sed 's/\/mnt//' > '/mnt/etc/zfs/zfs-list.cache/'"${zpool_name}"
    systemctl enable 'zfs-zed.service' --root='/mnt'
}

function set_new_uefi_boot_entries () {
    #3.8
    declare -a uefi_images
    mapfile -t uefi_images < \
        <(find '/mnt/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}')

    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#/mnt/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
}

function insert_zbm_postconf_hook () {
    #3.9
    declare postconf_target_abs='/mnt/etc/zfsbootmenu/posthooks.d/'"$(basename "${postconf_hook}")"
    curl --silent --location "${postconf_hook}" --output "${postconf_target_abs}"
    chmod +x "${postconf_target_abs}"
}

function umount_all () {
    #3.10
    if [[ "${part_schema}" = 'mbr' ]]; then
        umount '/mnt/boot/syslinux'
    else
        umount '/mnt/efi'
    fi
    zfs umount -a
    zpool export "${zpool_name}"
}

function finalize_os_setup () {
    #3.1
    set_root_pw                                 #3.2
    configure_networking                        #3.3
    configure_dns                               #3.4
    configure_reflector                         #3.5
    configure_zfs                               #3.6
    configure_zfs_mount_gen                     #3.7
    if [[ "${part_schema}" = 'gpt' ]]; then
        set_new_uefi_boot_entries               #3.8
        insert_zbm_postconf_hook                #3.9
    fi
    umount_all                                  #3.10
}

function main () {
    if we_are_changerooted; then
        install_os_in_chroot                    #2.2
    else
        no_kernel_update_in_iso                 #1.1
        set_ntp                                 #1.2
        resize_cow_space                        #1.3
        update_pacman_db                        #1.4
        install_pkgs 'jq'                       #1.5
        install_zfs                             #1.6
        uefi_or_bios                            #1.7
        setup_zpool                             #1.8
        mount_system                            #1.9
        copy_zpool_cache                        #1.10
        install_archlinux                       #1.11
        gen_fstab                               #1.12
        set_hostname                            #1.13
        set_locale                              #1.14
        add_zfs_hook_to_initramfs               #1.15
        set_initramfs_build_list                #1.16
        add_zfs_files_to_new_os                 #1.17
        enter_chroot                            #2.1
        # We're done in chroot
        finalize_os_setup                       #3.1
    fi
}

main