#!/usr/bin/env bash
# Copyright (C) 2023-2024  Umorpha Systems
# SPDX-License-Identifier: AGPL-3.0-or-later

# Dependencies:
#  - /usr/bin/env       (coreutils, busybox)
#  - cat                (coreutils, busybox)
#  - dd                 (coreutils, busybox--only needed fo typ=uboot-sunxi)
#  - install            (coreutils, busybox)
#  - mkdir              (coreutils, busybox)
#  - mktemp             (coreutils, busybox)
#  - rm                 (coreutils, busybox)
#  - rmdir              (coreutils, busybox)
#  - losetup (optional) (util-linux; the busybox one doesn't work)
#  - umount             (util-linux, busybox)
#  - blkid              (util-linux)
#  - mount              (util-linux)
#  - bash               (bash)
#  - sgdisk             (gptfdisk--only needed fo typ=uboot-sunxi)
#  - systemd-repart     (systemd)
#  - lvcreate           (lvm2)
#  - lvresize           (lvm2)
#  - pvcreate           (lvm2)
#  - pvs                (lvm2)
#  - vgchange           (lvm2)
#  - vgcreate           (lvm2)
#  - vgremove           (lvm2)
#  - vgs                (lvm2)
#  - pv                 (pv)
#  - mkfs.ext4          (e2fsprogs)
#  - mkfs.vfat          (dosfstools)
#  - arch-chroot        (arch-install-scripts)
#  - umorpha-mount
#
# Dependencies of arch-chroot (as of v28):
#  - cat        (coreutils, busybox)
#  - install    (coreutils, busybox)
#  - ln         (coreutils, busybox)
#  - touch      (coreutils, busybox)
#  - chroot     (coreutils; the busybox one doesn't work)
#  - readlink   (coreutils; the busybox one doesn't work)
#  - mountpoint (util-linux, busybox)
#  - umount     (util-linux, busybox)
#  - unshare    (util-linux, busybox)
#  - mount      (util-linux)
#  - grep       (grep, busybox)
#  - bash       (bash)
#
# Dependencies of umorpha-mount:
#  - sh     (bash, busybox)
#  - getopt (util-linux, busybox)
#  - mkdir  (coreutils, busybox)
#  - mount  (util-linux)
#
# Of those, these are the ones that aren't already in an initcpio:
#  - arch-chroot    (arch-install-scripts)
#  - bash           (bash)
#  - chroot         (coreutils)
#  - readlink       (coreutils)
#  - pv             (pv)
#  - systemd-repart (systemd)
#  - mkfs.ext4      (e2fsprogs)
#  - mkfs.vfat      (dosfstools)
#  - lvcreate       (lvm2)
#  - lvresize       (lvm2)
#  - pvcreate       (lvm2)
#  - pvs            (lvm2)
#  - vgchange       (lvm2)
#  - vgcreate       (lvm2)
#  - vgremove       (lvm2)
#  - vgs            (lvm2)

set -euE -o pipefail

v() {
	printf >&2 '  $ %s\n' "${*@Q}"
	"$@"
}

in_array() {
	local needle haystack straw
	needle=$1
	haystack=("${@:2}")
	for straw in "${haystack[@]}"; do
		if [[ "$needle" == "$straw" ]]; then
			return 0
		fi
	done
	return 1
}

main() {
	local typ device rootfs slot
	typ=$1
	rootfs=$2
	device=$3
	slot='a'

	echo ":: Installing '${rootfs}' to '${device}' for '${typ}'"
	if [[ "$rootfs" != *rootfs* ]]; then
		echo >&2 "!! Filename '${rootfs}' does not look like a rootfs filename"
		return 2
	fi
	case "$typ" in
		uefi|bios|uboot-sunxi) :;;
		*)
			echo >&2 "!! Install type '${typ}' is not valid"
			return 2
			;;
	esac

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

	cleanup=()
	_do_cleanup() {
		if [[ $? == 0 ]]; then
			echo ":: Cleaning up..."
		else
			echo >&2 "!! ERROR! Cleaning up..."
		fi
		for (( i=${#cleanup[@]}-1; i >= 0; i-- )); do
			eval "${cleanup[$i]}"
		done
	}
	trap '_do_cleanup' EXIT

	local tmpdir
	tmpdir="$(mktemp -dt "${0##*/}.XXXXXXXXXX")"
	cleanup+="rmdir -- ${tmpdir@Q}"

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

	local is_loop=false
	if [[ -f "$device" ]]; then
		echo "   file '${device}' is not a block device, setting up a loopback device..."
		device=$(v losetup --find --show --partscan -- "$device")
		echo "   ... '${device}'"
		cleanup+=("v losetup --detach ${device@Q}")
		is_loop=true
	fi
	local devicep="$device"
	if [[ "$devicep" == *[0-9] ]]; then
		devicep+='p'
	fi

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

	echo ":: Setting up GPT"

	n=0
	cat >"$tmpdir/00-esp.conf" <<-EOF
		[Partition]
		# EFI System Partition, specified by UEFI
		Type=esp
		Format=vfat
		Label=Umorpha EFI System Partition
		SizeMinBytes=512M
		SizeMaxBytes=512M
	EOF
	n_esp=$((++n))
	sector_size=''
	case "$typ" in
		bios)
			cat >"$tmpdir/05-bios.conf" <<-EOF
				[Partition]
				# GRUB BIOS Boot partition, specified by GRUB
				Type=21686148-6449-6E6F-744E-656564454649
				Label=Umorpha BIOS Boot Partition
				SizeMinBytes=1M
				SizeMaxBytes=1M
			EOF
			n_bios=$((++n))
			;;
		uboot-sunxi)
			sector_size=512
			# Commented out because `BegOffset=` and `EndOffset=`
			# aren't real systemd-repart options.
			#
			#cat >"$tmpdir/05-uboot.conf" <<-EOF
			#	[Partition]
			#	# U-Boot Environment partition, specified by U-Boot
			#	Type=3DE21764-95BD-54BD-A5C3-4ABE786F38A8
			#	Label=Umorpha U-Boot Partition
			#	BegOffset=8192b
			#	EndOffset=1M
			#EOF
			#n_uboot=$((++n))
			;;
	esac
	cat >"$tmpdir/10-lvm.conf" <<-EOF
		[Partition]
		# LVM2 partition, specified by ???
		# https://en.wikipedia.org/w/index.php?title=User_talk:Claunia&diff=prev&oldid=1182832243
		Type=E6D6D379-F507-44C2-A23C-238F2A3DF928
		Label=Umorpha LVM Partition
	EOF
	n_lvm=$((++n))

	v systemd-repart \
		--empty=allow --dry-run=no \
		--no-pager \
		--definitions="$tmpdir" \
		${sector_size:+--sector-size="$sector_size"} \
		"$device"

	rm -- "$tmpdir"/*.conf

	if [[ $typ == uboot-sunxi ]]; then
		# Shrink the number of entries in the partition table
		# to make room for U-Boot; the location that sunxi
		# chips need us to put U-Boot at is by default
		# occupied by the partition table.
		#
		# def+min table location is LBA=2, def+min entry size
		# is 128b, def number of entries is 128, which means
		# that it ranges from 1KiB to 17KiB.
		#
		# U-Boot needs to go at 8KiB, meaning that if we want
		# leave the partition table at the default location it
		# can be a maximum of 7KiB.
		#
		# 7KiB / (128b/entry) = 56entries
		v sgdisk --resize-table=56 "$device"

		# TODO: I'd like to set up a partition for U-Boot, to
		# make it clear that it's there (instead of putting it
		# in unpartitioned space).  But the default
		# FirstUsableLBA is set to 1MiB, and so sgdisk will
		# silently bump the partition's start LBA to 1MiB.
		# And I don't see a way to have sgdisk change the
		# FirstUsableLBA.
		#
		#v sgdisk \
		#	--new=3:$((8*1024/512)):$((1024*1024/512)) \
		#	--typecode=3DE21764-95BD-54BD-A5C3-4ABE786F38A8 \
		#	--change-name=3:'Umorpha U-Boot Partition' \
		#	--sort
		#	"$device"
	fi

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

	echo ":: Setting up LVM"

	if ! in_array "${devicep}${n_lvm}" $(pvs --noheadings --options=pv_name); then
		v pvcreate "${devicep}${n_lvm}"
	fi
	if ! in_array vg_umorpha $(vgs --noheadings --options=vg_name); then
		v vgcreate vg_umorpha "${devicep}${n_lvm}"
	fi
	if $is_loop; then
		cleanup+=("v vgchange --activate=n vg_umorpha")
	fi

	# Use short `-c` instead of `--format` so this works with busybox.
	size="$(stat -c '%s' -- "$rootfs")"

	if [[ -e /dev/mapper/vg_umorpha-lv_root_$slot ]]; then
		# Ugg, lvresize exits with non-zero for no-ops, which we can't
		# even really do a pre-flight check for since it might round a
		# size up.
		v lvresize --size="${size}b" "/dev/mapper/vg_umorpha-lv_root_$slot" || true
	else
		v lvcreate --zero=y --wipesignatures=y --yes --name="lv_root_$slot" \
			 --size="${size}b" vg_umorpha
	fi

	v pv -- "$rootfs" >"/dev/mapper/vg_umorpha-lv_root_$slot"

	if ! [[ -e /dev/mapper/vg_umorpha-lv_root_overlay ]]; then
		v lvcreate --zero=y --wipesignatures=y --yes \
		        --name=lv_root_overlay --size=1G vg_umorpha
	fi
	if [[ -z "$(blkid --output=export /dev/mapper/vg_umorpha-lv_root_overlay | sed -n s/^TYPE=//p)" ]]; then
		v mkfs.ext4 /dev/mapper/vg_umorpha-lv_root_overlay
	fi

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

	echo ":: Configuring boot"

	local selector_boot selector_root selector_overlay
	selector_boot="$(   blkid --output=export "${devicep}${n_esp}"                   | grep ^UUID=)"
	selector_root="$(   blkid --output=export "/dev/mapper/vg_umorpha-lv_root_$slot" | grep ^UUID=)"
	selector_overlay="$(blkid --output=export /dev/mapper/vg_umorpha-lv_root_overlay | grep ^UUID=)"
	v declare -p selector_boot selector_root selector_overlay

	v umorpha-mount \
		--root="/dev/mapper/vg_umorpha-lv_root_$slot:auto:ro" \
		--overlay='/dev/mapper/vg_umorpha-lv_root_overlay:auto:rw' \
		--boot="${devicep}${n_esp}:auto:rw" \
		--basedir="$tmpdir"
	cleanup+=("rmdir -- ${tmpdir@Q}/root")
	cleanup+=("v umount ${tmpdir@Q}/root")
	cleanup+=("rmdir -- ${tmpdir@Q}/overlay")
	cleanup+=("v umount ${tmpdir@Q}/overlay")
	cleanup+=("rmdir -- ${tmpdir@Q}/mnt")
	cleanup+=("v umount ${tmpdir@Q}/mnt")
	cleanup+=("v umount ${tmpdir@Q}/mnt/boot")

	case "$typ" in
		uefi|uboot-sunxi)
			if [[ $typ == uboot-sunxi ]]; then
				v dd if="$(echo "$tmpdir/mnt"/usr/lib/u-boot/*/u-boot-sunxi-with-spl.bin)" of="$device" conv=notrunc bs=512 seek=16
			fi
			mkdir -p -- "$tmpdir/mnt/etc/kernel"
			echo "root=$selector_root overlay=$selector_overlay boot=$selector_boot rw" >"$tmpdir/mnt/etc/kernel/cmdline"
			mkdir -p -- "$tmpdir/mnt/boot/EFI/Linux"
			bootloader_src=$(echo "$tmpdir"/mnt/usr/lib/systemd/boot/efi/systemd-boot*.efi)
			bootloader_dst=${bootloader_src##*-}
			bootloader_dst="$tmpdir/mnt/boot/EFI/BOOT/${bootloader_dst^^}"
			v install -D "$bootloader_src" "$bootloader_dst"
			;;
		bios)
			v arch-chroot -- "$tmpdir/mnt" grub-install --target=i386-pc "$device"
			install -Dm600 /dev/stdin "$tmpdir/mnt/boot/grub/grub.cfg" <<-EOF
				insmod part_gpt
				insmod part_msdos
				if [ -s \$prefix/grubenv ]; then
				  load_env
				fi
				if [ -s \$prefix/grub_a.cfg ]; then
				  configfile \$prefix/grub_a.cfg
				fi
				if [ -s \$prefix/grub_b.cfg ]; then
				  configfile \$prefix/grub_b.cfg
				fi
			EOF
			{
				for krnl in "$tmpdir/mnt"/usr/lib/modules/*/vmlinuz; do
					pkgbase=$(cat -- ${krnl%vmlinuz}pkgbase)
					cat <<-EOF
						menuentry 'Parabola GNU/Linux-libre, slot ${slot}, ${pkgbase} kernel'
						  set search --no-floppy --fs-uuid --set=root ${selector_boot#UUID=}
						  echo 'Loading ${pkgbase} kernel...'
						  linux /vmlinuz-${pkgbase} root=$selector_root overlay=$selector_overlay boot=$selector_boot rw
						  echo 'Loading initial ramdisk...'
						  initrd /initramfs-${pkgbase}.img
						  echo 'Booting...'
						  boot
					EOF
				done
			} | install -Dm600 /dev/stdin "$tmpdir/mnt/boot/grub/grub_${slot}.cfg"
			;;
	esac

	if ! v arch-chroot -- "$tmpdir/mnt" sh -c "printf '%s\n' usr/lib/modules/*/vmlinuz | /usr/share/libalpm/scripts/mkinitcpio install"; then
		v arch-chroot -- "$tmpdir/mnt" sh -c "MKINITCPIO_PROCESS_PRESET=linux-libre-fallback bash -x /bin/mkinitcpio -k /usr/lib/modules/6.7.4-gnu-1/vmlinuz -U /boot/EFI/Linux/parabola-linux-libre-fallback.efi -S autodetect --microcode '/boot/*-ucode.img'"
	fi
	v arch-chroot -- "$tmpdir/mnt" ldconfig
	v arch-chroot -- "$tmpdir/mnt" systemd-machine-id-setup
	v arch-chroot -- "$tmpdir/mnt" update-ca-trust
}

main "$@"