#!/usr/bin/env bash
# install arch linux on zfs root, stage one
# Craig Jennings <c@cjennings.net>
#
# https://openzfs.github.io/openzfs-docs/Getting%20Started/Arch%20Linux/Root%20on%20ZFS.html

# set -e # halt on any error

#  -------------------------- Prerequisites --------------------------

sed -i "s/^#ParallelDownloads = 5$/ParallelDownloads = 15/" /etc/pacman.conf
loadkeys us
timedatectl set-ntp true

###  --------------------------- Choose Disk ---------------------------

all_disk_ids=( $(ls /dev/disk/by-id/) )
echo ""; echo "Select the disk id to use. All data will be erased."
select disk_id in "${all_disk_ids[@]}"; do
    # ensure valid selection
    if [[ -n $disk_id ]]; then
        selection=$disk_id
        break
    else
        echo "Invalid. Try again."
    fi
done

# Confirm the selected disk
read -p "Confirm: '$selection' [y/n]? " choice
if [[ "$choice" != "y" ]]; then
    echo "Exiting..."
    exit 1
fi

DISK="/dev/disk/by-id/$selection"
MNT=/mnt   # Set a mount point
SWAPSIZE=4 # Set swap size in GB
RESERVE=1  # Set how much space to leave at the end of disk, minimum 1GB

###  --------------------------- Erase Disk -------------------------

echo ""; echo "### Erasing Disk"
blkdiscard -f "${DISK}" || true    # discard all sectors on flash-based storage
sgdisk --zap-all "${DISK}"         # clear the disk

###  ------------------------- Partition Disk ------------------------

echo ""; echo "### Partitioning Disk"
parted --script --align=optimal  "${DISK}" -- \
 mklabel gpt \
 mkpart EFI 2MiB 1GiB \
 mkpart bpool 1GiB 5GiB \
 mkpart rpool 5GiB -$((SWAPSIZE + RESERVE))GiB \
 mkpart swap  -$((SWAPSIZE + RESERVE))GiB -"${RESERVE}"GiB \
 mkpart BIOS 1MiB 2MiB \
 set 1 esp on \
 set 5 bios_grub on \
 set 5 legacy_boot on

###  ---------------------- Setup Encrypted Swap ---------------------

echo ""; echo "### Encrypted Swap"
for i in ${DISK}; do
   cryptsetup open --type plain --key-file /dev/random "${i}"-part4 "${i##*/}"-part4
   mkswap /dev/mapper/"${i##*/}"-part4
   swapon /dev/mapper/"${i##*/}"-part4
done

#  ------------------- Create Boot And Root Pools ------------------

# This step creates a separate boot pool for /boot with the features limited to
# only those that GRUB supports, allowing the root pool to use any/all features.

echo ""; echo "### Checking ZFS Module"
modprobe zfs # ensure zfs module is loaded

echo ""; echo "### Creating Boot Pool"
# shellcheck disable=SC2046
zpool create -d \
    -o feature@async_destroy=enabled \
    -o feature@bookmarks=enabled \
    -o feature@embedded_data=enabled \
    -o feature@empty_bpobj=enabled \
    -o feature@enabled_txg=enabled \
    -o feature@extensible_dataset=enabled \
    -o feature@filesystem_limits=enabled \
    -o feature@hole_birth=enabled \
    -o feature@large_blocks=enabled \
    -o feature@lz4_compress=enabled \
    -o feature@spacemap_histogram=enabled \
    -o ashift=12 \
    -o autotrim=on \
    -O acltype=posixacl \
    -O canmount=off \
    -O compression=lz4 \
    -O devices=off \
    -O normalization=formD \
    -O relatime=on \
    -O xattr=sa \
    -O mountpoint=/boot \
    -R "${MNT}" \
    bpool \
    $(for i in ${DISK}; do
       printf '%s ' "${i}-part2";
      done)

echo ""; echo "### Creating Root Pool"
# shellcheck disable=SC2046
zpool create \
    -o ashift=12 \
    -o autotrim=on \
    -R "${MNT}" \
    -O acltype=posixacl \
    -O canmount=off \
    -O compression=zstd \
    -O dnodesize=auto \
    -O normalization=formD \
    -O relatime=on \
    -O xattr=sa \
    -O mountpoint=/ \
    rpool \
   $(for i in ${DISK}; do
      printf '%s ' "${i}-part3";
     done)

echo ""; echo "### Creating Unencrypted Root System Container"
# create UNENCRYPTED root system container
zfs create \
 -o canmount=off \
 -o mountpoint=none \
rpool/archlinux

#  --------------------- Create System Datasets --------------------

echo ""; echo "### Creating System Datasets"
zfs create -o canmount=noauto -o mountpoint=/  rpool/archlinux/root
zfs mount rpool/archlinux/root
zfs create -o mountpoint=legacy rpool/archlinux/home
mkdir "${MNT}"/home
mount -t zfs rpool/archlinux/home "${MNT}"/home
zfs create -o mountpoint=legacy  rpool/archlinux/var
zfs create -o mountpoint=legacy rpool/archlinux/var/lib
zfs create -o mountpoint=legacy rpool/archlinux/var/log
zfs create -o mountpoint=none bpool/archlinux
zfs create -o mountpoint=legacy bpool/archlinux/root
mkdir "${MNT}"/boot
mount -t zfs bpool/archlinux/root "${MNT}"/boot
mkdir -p "${MNT}"/var/log
mkdir -p "${MNT}"/var/lib
mount -t zfs rpool/archlinux/var/lib "${MNT}"/var/lib
mount -t zfs rpool/archlinux/var/log "${MNT}"/var/log

#  ---------------------- Format And Mount ESP ---------------------

echo ""; echo "### Format And Mount ESP"
for i in ${DISK}; do
 mkfs.vfat -n EFI "${i}"-part1
 mkdir -p "${MNT}"/boot/efis/"${i##*/}"-part1
 mount -t vfat -o iocharset=iso8859-1 "${i}"-part1 "${MNT}"/boot/efis/"${i##*/}"-part1
done

mkdir -p "${MNT}"/boot/efi
mount -t vfat -o iocharset=iso8859-1 "$(echo "${DISK}" | sed "s|^ *||"  | cut -f1 -d' '|| true)"-part1 "${MNT}"/boot/efi

###  -------------------------- Install Base -------------------------

echo ""; echo "### Installing Base"
# install packages with pacstrap
pacstrap "${MNT}" \
         base \
         base-devel \
         dkms \
         efibootmgr \
         git \
         grub \
         intel-ucode \
         linux-firmware \
         linux-lts \
         linux-lts-headers \
         man-db \
         man-pages \
         vi

# create fstab and remove all zroot entries
echo ""; echo "### Creating fstab"
genfstab -U -p "${MNT}" >> "${MNT}"/etc/fstab
sed -i '/^# zroot/d' "${MNT}"/etc/fstab
sed -i '/^zroot/d' "${MNT}"/etc/fstab
sed -i '/^$/d' "${MNT}"/etc/fstab
echo "" >> "${MNT}"/etc/fstab # one blank line at the end

# copy over dns settings to the new system
echo ""; echo "### Copying DNS Settings"
cp -v /etc/resolv.conf "${MNT}"/etc

#  ----------------------------- Chroot ----------------------------
# copy second part of this script to the new system

echo ""; echo "### Copying script then chrooting"
sed '1,/^#part2$/d' `basename $0` >  "${MNT}"/root
chroot "${MNT}" /usr/bin/env DISK="${DISK}" /root/stage2.sh

#  -------------------------- After Chroot -------------------------

umount -Rl "${MNT}"
zfs snapshot -r rpool@initial-installation
zfs snapshot -r bpool@initial-installation

zpool export -a

echo ""; echo "### Done with Stage Two"

#part2
#!/usr/bin/env bash
# install arch linux on zfs root, stage two
# Craig Jennings <c@cjennings.net>

printf '\033c'

password="welcome" # root password; will force change after login
hostname="velox"
disk=/dev/nvme0n1

yay_repo="https://aur.archlinux.org/yay.git"
source_dir="/usr/src"
logfile=/root/zfsarch_stage2.log

cd # go home

#  --------------------- Add ArchZFS Repository --------------------

pacman-key --init
pacman-key --refresh-keys
pacman-key --populate

curl --fail-early --fail -L https://archzfs.com/archzfs.gpg \
|  pacman-key -a - --gpgdir /etc/pacman.d/gnupg

pacman-key \
--lsign-key \
--gpgdir /etc/pacman.d/gnupg \
DDF7DB817396A49B2A2723F7403BD972F75D9D76

tee -a /etc/pacman.d/mirrorlist-archzfs <<- 'EOF'
## See https://github.com/archzfs/archzfs/wiki
## France
#,Server = https://archzfs.com/$repo/$arch

## Germany
#,Server = https://mirror.sum7.eu/archlinux/archzfs/$repo/$arch
#,Server = https://mirror.biocrafting.net/archlinux/archzfs/$repo/$arch

## India
#,Server = https://mirror.in.themindsmaze.com/archzfs/$repo/$arch

## United States
#,Server = https://zxcvfdsa.com/archzfs/$repo/$arch
EOF

tee -a /etc/pacman.conf <<- 'EOF'

#[archzfs-testing]
#Include = /etc/pacman.d/mirrorlist-archzfs

#,[archzfs]
#,Include = /etc/pacman.d/mirrorlist-archzfs
EOF

# this #, prefix is a workaround for ci/cd tests
# remove them
sed -i 's|#,||' /etc/pacman.d/mirrorlist-archzfs
sed -i 's|#,||' /etc/pacman.conf
sed -i 's|^#||' /etc/pacman.d/mirrorlist

#  ---------- Install Packages And Linux Compatible Kernel ---------
pacman -Sy

kernel_compatible_with_zfs="$(pacman -Si zfs-linux \
| grep 'Depends On' \
| sed "s|.*linux=||" \
| awk '{ print $1 }')"
pacman -U --noconfirm https://america.archive.pkgbuild.com/packages/l/linux/linux-"${kernel_compatible_with_zfs}"-x86_64.pkg.tar.zst

pacman -S --noconfirm zfs-linux zfs-utils

#  ---------------------- Configure Mkinitcpio ---------------------

sed -i 's|filesystems|zfs filesystems|' /etc/mkinitcpio.conf
mkinitcpio -P

#  ------------------------- General Hostid ------------------------

zgenhostid -f -o /etc/hostid

#  ---------------------- Apply Grub Workaround ----------------------
# Note: This workaround needs to be applied for every GRUB update, as the update will overwrite the changes.

echo 'export ZPOOL_VDEV_NAME_PATH=YES' >> /etc/profile.d/zpool_vdev_name_path.sh
# shellcheck disable=SC1091
. /etc/profile.d/zpool_vdev_name_path.sh

# GRUB fails to detect rpool name, hard code as "rpool"
sed -i "s|rpool=.*|rpool=rpool|"  /etc/grub.d/10_linux

#  -------------------------- Install Grub -------------------------

mkdir -p /boot/efi/archlinux/grub-bootdir/i386-pc/
mkdir -p /boot/efi/archlinux/grub-bootdir/x86_64-efi/
for i in ${DISK}; do
 grub-install --target=i386-pc --boot-directory \
     /boot/efi/archlinux/grub-bootdir/i386-pc/  "${i}"
done
grub-install --target x86_64-efi --boot-directory \
 /boot/efi/archlinux/grub-bootdir/x86_64-efi/ --efi-directory \
 /boot/efi --bootloader-id archlinux --removable
if test -d /sys/firmware/efi/efivars/; then
   grub-install --target x86_64-efi --boot-directory \
    /boot/efi/archlinux/grub-bootdir/x86_64-efi/ --efi-directory \
    /boot/efi --bootloader-id archlinux
fi

# Import both bpool and rpool at boot:
echo 'GRUB_CMDLINE_LINUX="zfs_import_dir=/dev/"' >> /etc/default/grub

# Generate GRUB menu:
mkdir -p /boot/grub
grub-mkconfig -o /boot/grub/grub.cfg
cp /boot/grub/grub.cfg \
 /boot/efi/archlinux/grub-bootdir/x86_64-efi/grub/grub.cfg
cp /boot/grub/grub.cfg \
 /boot/efi/archlinux/grub-bootdir/i386-pc/grub/grub.cfg

# For both legacy and EFI booting: mirror ESP content:
espdir=$(mktemp -d)
find /boot/efi/ -maxdepth 1 -mindepth 1 -type d -print0 \
| xargs -t -0I '{}' cp -r '{}' "${espdir}"
find "${espdir}" -maxdepth 1 -mindepth 1 -type d -print0 \
    | xargs -t -0I '{}' sh -vxc "find /boot/efis/ -maxdepth 1 -mindepth 1 -type d -print0 | xargs -t -0I '[]' cp -r '{}' '[]'"

###  ----------------------- Config Environment ---------------------

# required software
pacman -S --noconfirm sed networkmanager reflector
systemctl enable NetworkManager
systemctl enable reflector.timer

# parallel downloads
sed -i "s/^#ParallelDownloads = 5$/ParallelDownloads = 15/" /etc/pacman.conf

# aur config
sed -i '/^CFLAGS=/s/-march=x86-64 -mtune=generic/-march=native/' /etc/makepkg.conf
sed -i 's/^#RUSTFLAGS="-C opt-level=2"/RUSTFLAGS="-C opt-level=2 -C target-cpu=native"/g'  /etc/makepkg.conf
sed -i 's/^#MAKEFLAGS="-j2"/MAKEFLAGS="-j$(nproc)"/g' /etc/makepkg.conf
sed -i 's/^COMPRESSXZ=(xz -c -z -)/COMPRESSXZ=(xz -c -z --threads=0 -)/g' /etc/makepkg.conf
sed -i 's/^COMPRESSZST=(zstd -c -z -q -)/COMPRESSZST=(zstd -c -z -q --threads=0 -)/g' /etc/makepkg.conf

# set up local time
sudo ln -sf /usr/share/zoneinfo/US/Central /etc/localtime
hwclock --systohc

# set up locale
echo "en_US.UTF-8 UTF-8" >> /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf

# set up hostname
echo "$hostname" > /etc/hostname
echo "127.0.0.1       localhost" >> /etc/hosts
echo "::1             localhost" >> /etc/hosts
echo "127.0.1.1       $hostname.localdomain $hostname" >> /etc/hosts

# console settings
echo "KEYMAP=us" > /etc/vconsole.conf

#  ------------------------ Enable AUR Helper ------------------------

pacman -Syy
pacman -S --needed --noconfirm base-devel git openssh wget
build_dir="$source_dir/yay"
mkdir -p "$build_dir"
chown "$username:$username" -R "/home/$username"
sudo -u "$username" git clone --depth 1 "$yay_repo" "$build_dir"
cd "$build_dir" && sudo -u "$username" makepkg --noconfirm -si

###  ------------------------- Enable ZFS Services ----------------------
# add kernel modules

pacman -S --noconfirm zfs-dkms

# enable zfs services
systemctl enable zfs-import-cache
systemctl enable zfs-import.target
systemctl enable zfs-mount.service
systemctl enable zfs-share
systemctl enable zfs-zed
systemctl enable zfs.target

###  ----------------------------- Wrap Up --------------------------

# set root password; must change first login
echo "root:$password" | chpasswd
chage -d 0 root

#  --------------------------- Exit Chroot ---------------------------
exit