#!/bin/bash

declare this_script_remote_repo_raw_url zpool_name zfs_arch_dataset_name
this_script_remote_repo_raw_url='https://quico.space/...'
zpool_name='zpool'
zfs_arch_dataset_name='archlinux'

function set_ntp () {
    timedatectl set-ntp true
}

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

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

function install_zfs () {
    curl -s 'https://raw.githubusercontent.com/eoli3n/archiso-zfs/master/init' | bash
}

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

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

function we_have_exactly_one_part () {
    declare 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)
                printf -- '%s\n' 'No '"${parttype}"' partition found. Exiting 1 ...'
                exit 1
                ;;
            1)
                return 0
                ;;
            *)
                printf -- '%s\n' 'Multiple '"${parttype}"' partitions found. We expect exactly one. Cowardly exiting 1 ...'
                exit 1
                ;;
        esac
        printf -- '%s\n' 'Partition list does not look valid. Cowardly exiting 1 ...'
        exit 1
    fi
}

function no_zpool_exists () {
    declare zpool_list
    zpool_list="$(zpool list -H)"
    [[ "$(wc -l <<<"${zpool_list}")" -eq '0' ]] && return 0
    return 1
}

function zpool_drive_id () {
    declare drive_id
    drive_id="$(find -L /dev/disk/by-id -samefile "${1:?}")"
    if [[ "$(wc -l <<<"${drive_id}")" -eq '1' ]]; then
        printf -- '%s' "${drive_id}"
        return 0
    fi
    printf -- '%s\n' 'More than zpool partition entry in /dev/disk/by-id, exiting 1 ...'
    exit 1
}

function set_zpool_password () {
    printf -- '%s\n' 'password' > '/etc/zfs/'"${zpool_name}"'.key'
    chmod '000' '/etc/zfs/'"${zpool_name}"'.key'
}

function import_pool () {
    zpool import -d '/dev/disk/by-id' -R '/mnt' "${zpool_name}" -N -f
    zfs load-key "${zpool_name}"
}

function create_pool () {
    # Create a temporary pool that is not cached
    zpool create -f                                                     \
                 -o 'ashift=12'                                         \
                 -o 'autotrim=on'                                       \
                 -O 'acltype=posix'                                     \
                 -O 'compression=on'                                    \
                 -O 'relatime=on'                                       \
                 -O 'xattr=sa'                                          \
                 -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 "${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 () {
    declare zpool_drive drive_by_id
    zpool_drive="${1:?}"
    drive_by_id="$(zpool_drive_id "${zpool_drive}")"

    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
    fi
}

function mount_system () {
    zfs mount "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
    zfs mount -a
    
    declare efi_part
    efi_part="$(get_parts 'efi')"
    if we_have_exactly_one_part 'efi' "${efi_part}"; then
        mkdir -p '/mnt/efi'
        mount "${efi_part}" '/mnt/efi'
    fi
}

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

function install_archlinux () {
    sed -ri -e 's'$'\x1''^.*?(ParallelDownloads)[^\r\n\f]*'$'\x1''\1 = 5'$'\x1''g' '/etc/pacman.conf'
    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 () {
    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 () {
    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 () {
    printf -- '%s\n' 'KEYMAP=de-latin1' > '/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 () {
    sed -ri -e 's'$'\x1''(HOOKS=)(.*?[\(| ])(filesystems)([\)| ][^\r\n\f]*)'$'\x1''\1\2zfs \3\4'$'\x1''g' '/mnt/etc/mkinitcpio.conf'
}

function add_zfs_files_to_new_os () {
    for zfs_file in '/etc/hostid' '/etc/zfs/zpool.cache' '/etc/zfs/zpool.key'; do
        rsync -av --itemize-changes {'','/mnt'}"${zfs_file}"
    done
}

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

function get_pkg_info () {
    declare from_where local_pkg_info local_pkg_version version_search
    from_where="${1}"
    version_search='/Version/{print $3}'
    pkg_info="$(paru -$([[ "${from_where}" == 'local' ]] && printf -- '%s' 'Q' || printf -- '%s' 'S')i "${2}" 2>&1)"
    if [[ "${from_where}" == 'local' ]] && grep -Piq -- '^error: package .*? was not found' <<<"${pkg_info}"; then
        return 1
    else
        local_pkg_version="$(awk "${version_search}" <<<"${pkg_info}")"
    fi
    printf -- '%s' "${local_pkg_version}"
    return 0
}

function paru_with_zfs_first () {
    if [[ "${#}" -eq '0' ]]; then
        declare -A local_pkg_info
        /usr/bin/paru -Sy
        if local_pkg_info['zfs-dkms']="$(get_pkg_info 'local' 'zfs-dkms')" && local_pkg_info['zfs-utils']="$(get_pkg_info 'local' 'zfs-utils')"; then
            local_pkg_info['zfs-dkms']="$(get_pkg_info 'local' 'zfs-dkms')"
            local_pkg_info['zfs-utils']="$(get_pkg_info 'local' 'zfs-utils')"

            declare -A remote_pkg_info
            remote_pkg_info['zfs-dkms']="$(get_pkg_info 'remote' 'zfs-dkms')"
            remote_pkg_info['zfs-utils']="$(get_pkg_info 'remote' 'zfs-utils')"

            /usr/bin/paru -S --needed archlinux-keyring

            if [[ "${local_pkg_info['zfs-dkms']}" == "${remote_pkg_info['zfs-dkms']}" ]] && \
                [[ "${local_pkg_info['zfs-utils']}" == "${remote_pkg_info['zfs-utils']}" ]]; then
                /usr/bin/paru -Su
            else
                /usr/bin/paru -Sy 'zfs-dkms' 'zfs-utils' \
                    --assume-installed zfs-dkms="${local_pkg_info['zfs-dkms']}" \
                    --assume-installed zfs-dkms="${remote_pkg_info['zfs-dkms']}" \
                    --assume-installed zfs-utils="${local_pkg_info['zfs-utils']}" \
                    --assume-installed zfs-utils="${remote_pkg_info['zfs-utils']}"
                /usr/bin/paru -Su
            fi
        else
            /usr/bin/paru -S --needed archlinux-keyring
            /usr/bin/paru -Su
        fi
    else
        /usr/bin/paru "${@}"
    fi
}

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 get_aur_helper () {
    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'
    chmod -R 'build:' 'paru'
    pushd 'paru'
    sudo --user 'build' makepkg -si --noconfirm
    popd
    rm -rf 'paru'
    popd
    alias paru='paru_with_zfs_first'
}

function paru_install () {
    sudo --user build paru -S --noconfirm "${@}"
}

function install_os_in_chroot () {
    ### 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

    get_aur_helper
    paru_install 'zfs-dkms' 'zfs-utils'
    hwclock --systohc
    locale-gen
    source /etc/locale.conf
    mkinitcpio -P

    # Install ZFSBootMenu and deps
    git clone --depth=1 https://github.com/zbm-dev/zfsbootmenu/ '/tmp/zfsbootmenu'
    paru_install 'cpanminus' 'kexec-tools' 'fzf' 'util-linux'
    pushd '/tmp/zfsbootmenu'
    make
    make install
    cpanm --notest --installdeps .
    popd
    rm -rf '/tmp/zfsbootmenu'
}

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

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

[Network]
DHCP=ipv4
IPForward=yes

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

function configure_dns () {
    rm '/mnt/etc/resolv.conf'
    ln -s '/run/systemd/resolve/stub-resolv.conf' '/mnt/etc/resolv.conf'
    # sed -i 's/^#DNS=.*/DNS=1.1.1.1/' /mnt/etc/systemd/resolved.conf
    systemctl enable 'systemd-resolved' --root='/mnt'
}

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

function configure_zfs_mount_gen () {
    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 configure_zfsbootmenu () {
    mkdir -p '/mnt/efi/EFI/ZBM'
    curl -s 'https://raw.githubusercontent.com/zbm-dev/zfsbootmenu/master/etc/zfsbootmenu/mkinitcpio.conf' | sed -r -e '/^#/d' -e '/^$/d' > '/mnt/etc/zfsbootmenu/mkinitcpio.conf'
    cat > '/mnt/etc/zfsbootmenu/config.yaml' <<EOF
Global:
  ManageImages: true
  BootMountPoint: /efi
  InitCPIO: true

Components:
  Enabled: false
EFI:
  ImageDir: /efi/EFI/ZBM
  Versions: false
  Enabled: true
Kernel:
  CommandLine: ro loglevel=0 zbm.import_policy=hostid
  Prefix: vmlinuz
EOF
# Up here maybe 'ro quiet' instead of 'ro'
    zfs set org.zfsbootmenu:commandline='rw nowatchdog rd.vconsole.keymap=de-latin1' "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
}

function gen_zfsbootmenu () {
    arch-chroot /mnt /bin/bash -xe <<"EOF"
source /etc/locale.conf
find /efi/EFI/ZBM -type f -delete
generate-zbm
EOF
}

function get_disk_with_efipart () {
    declare disks_with_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_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=="ebd0a0a2-b9e5-4433-87c0-68b6b72699c7")] | length == 1) ) | .path')"
    if [[ "$(wc -l <<<"${disks_with_efipart}")" -eq '1' ]] && [[ "$(wc -c <<<"${disks_with_efipart}")" -gt '1' ]]; then
        printf -- '%s' "${disks_with_efipart}"
        return 0
    fi
    return 1
}

function add_zbm_to_efi () {
    if ! efibootmgr | grep -Pi -- 'ZFSBootMenu'; then
        declare efi_disk
        efi_disk="$(get_disk_with_efipart)"
        efibootmgr --disk "${efi_disk}" \
          --part 1 \
          --create \
          --label "ZFSBootMenu" \
          --loader "\EFI\ZBM\vmlinuz.efi" \
          --verbose
    fi
}

function umount_all () {
    umount '/mnt/efi'
    zfs umount -a
    zpool export "${zpool_name}"
}

function finalize_os_setup () {
    set_root_pw
    configure_networking
    configure_dns
    configure_zfs
    configure_zfs_mount_gen
    configure_zfsbootmenu
    gen_zfsbootmenu
    add_zbm_to_efi
    umount_all
}

function main () {
    if we_are_changerooted; then
        install_os_in_chroot
    else
        set_ntp
        install_pkgs 'jq'
        declare zfs_part
        zfs_part="$(get_parts 'zfs')"
        if we_have_exactly_one_part 'zfs' "${zfs_part}"; then
            printf -- 'Creating zpool on partition '"'"'%s'"'"' ...\n' "${zfs_part}"
            install_zfs
            setup_zpool "${zfs_part}"
            mount_system
            copy_zpool_cache
            systemctl start reflector
            install_archlinux
            gen_fstab
            set_hostname
            set_locale
            add_zfs_hook_to_initramfs
            enter_chroot
            # We're done in chroot
            finalize_os_setup
        fi
    fi
}

main