SCRIPTPATH="$(dirname "$0")"

errormsg() {
	echo "$1" >&2
	exit 1
}

# Check deps
depcheck() {
	for x in \
		"dd" \
		"fallocate" \
		"sfdisk" \
		"mkfs.vfat" \
		"mkfs.ext4" \
		"pvcreate" \
		"vgcreate" \
		"lvcreate" \
		"vgchange" \
		"mkfs.btrfs" \
		"cryptsetup" \
		"mktemp" \
		"mount" \
		"umount" \
		"mkdir" \
		"cat" \
		"install" \
		"apk" \
		"bwrap" \
		"patch" \
	; do
		if ! which "$x" >/dev/null 2>&1; then
			errormsg "Missing dependency $x, exiting."
		fi
	done
}

splitargs() {
	echo "$1" | sed 's/ /\n/g'
}

# BWRAP "chrooting"
wrapped() {
	bwrap \
		--new-session \
		--die-with-parent \
		--unshare-pid \
		--unshare-uts \
		--unshare-ipc \
		--bind "$tmpdir" / \
		--dev-bind /dev /dev \
		--ro-bind /sys /sys \
		--proc /proc \
		--tmpfs /tmp \
		--tmpfs /run \
		-- $@
}

setup() {
	: "${image:=nnd-rootfs}"
	: "${arch:=x86_64}"
	: "${kflav:=full}"
	: "${svcs:=bundle.core bundle.net}"
	: "${size:=1G}"

	VGN="nnd"
	LVN="rootlv"
	CRYPT_DM="nnd-decrypt"
	EFI_LABEL="nnd-efi"
	BOOT_LABEL="nnd-boot"
	CRYPT_LABEL="nnd-root-encrypted"

	# Always clean up
	trap cleanup_err INT TERM EXIT HUP

	# Set up file as loopback
	fallocate -v -l "$size" "$image"
	if ! [ -r "$image" ]; then
		errormsg "Something went wrong while preallocating sparse image"
	fi
	loopback="$(losetup --show -P -f "$image")"
	if ! [ -r "$loopback" ]; then
		errormsg "Something went wrong while setting up loopback image"
	fi

	# Create partition table
	sfdisk "$loopback" <<-EOT
	label: gpt
	size=64MiB,name=$EFI_LABEL,type=U
	size=512MiB,name=$BOOT_LABEL,type=L
	type=L
	EOT

	# Format partitions
	mkfs.vfat  -n "$EFI_LABEL"  "${loopback}p1"
	mkfs.ext4  -L "$BOOT_LABEL" "${loopback}p2"

	vgcreate "$VGN" "${loopback}p3"
	lvcreate "$VGN" -n "$LVN" -l '100%FREE'
	vgchange -a y "$VGN"

	# Generate encryption keyfile
	keyfile="$(mktemp "nnd-keyfile.XXXXXXXXX")"
	touch "$keyfile"
	chmod 400 "$keyfile"
	chown root:root "$keyfile"
	dd bs=512 count=4 if=/dev/urandom of="$keyfile" iflag=fullblock

	# Create encrypted root
	cryptsetup -q luksFormat "/dev/$VGN/$LVN" "$keyfile"
	cryptsetup -q open -d "$keyfile" "/dev/$VGN/$LVN" "$CRYPT_DM"

	mkfs.btrfs -L "$ROOT_LABEL" "/dev/mapper/$CRYPT_DM"

	# Set up chroot
	tmpdir="$(mktemp -d "nnd-image.XXXXXXXXX")"
	if [ -z "$tmpdir" ]; then
		errormsg "Something went wrong during working directory preparation, bailing out"
	fi
	mount -v -t btrfs "/dev/mapper/$CRYPT_DM" "$tmpdir"

	mkdir -p "$tmpdir/boot"
	mount -v -t ext4 "${loopback}p2" "$tmpdir/boot"

	mkdir -p "$tmpdir/boot/efi"
	mount -v -t vfat "${loopback}p1" "$tmpdir/boot/efi"

	# Copy decryption key to well-known path
	cp -v "$keyfile" "$tmpdir"/crypto_keyfile.bin

	install -Dm444 "$SCRIPTPATH"/repositories "$tmpdir"/etc/apk/repositories 
	# Set up bootstrap rootfs
	apk add \
		-v \
		--initdb \
		--allow-untrusted \
		--root "$tmpdir" \
		--repositories-file "$tmpdir/etc/apk/repositories" \
		--arch "$arch" \
		lfsbase \
		busybox \
		nnd-signkey \
		alpine-keys \
		apk-tools

	## Add btrfs and lvm to required initrd modules
	# TODO: Only append and don't replace
	install -Dm444 "$SCRIPTPATH"/mkinitfs.conf "$tmpdir"/etc/mkinitfs/mkinitfs.conf

	## Provide temporary override for grub, TODO: provide our own grub cfgs eventually
	install -Dm444 <(
	cat <<-EOF
	GRUB_DISTRIBUTOR="nnd"
	GRUB_TIMEOUT=2
	GRUB_DISABLE_SUBMENU=y
	GRUB_DISABLE_RECOVERY=true
	GRUB_CMDLINE_LINUX_DEFAULT="loglevel=3"
	GRUB_CMDLINE_LINUX="cryptroot=/dev/mapper/$VGN-$LVN cryptdm=$CRYPT_DM cryptdiscards=yes cryptkey=yes rootfstype=btrfs root=/dev/mapper/$CRYPT_DM $kextparam"
	EOF
	) "$tmpdir"/etc/default/grub

	# Set up default resolv.conf (in case no DNS server is configured via DHCP or otherwise)
	## TODO: Maybe distribute this via packages?
	install -Dm644 "$SCRIPTPATH"/resolv.conf "$tmpdir"/etc/resolv.conf

	# Set up rest of rootfs
	wrapped apk add -v \
		nnd \
		linux-"$kflav" \
		nnd-s6-linux-init-default \
		nnd-s6-services \
		grub-efi \
		iproute2 \
		bcnm@edge \
		mdevd \
		mdev-conf \
		utmps \
		findmnt \
		dosfstools \
		e2fsprogs \
		btrfs-progs \
		cryptsetup \
		net-predictable \
		lvm2

	# Add net-predictable to mdev.conf
	(cd "$tmpdir" && patch -p0 < "$SCRIPTPATH"/mdev.patch)

	# Install grub to ESP
	wrapped grub-install --removable --efi-directory /boot/efi --boot-directory /boot

	# Unlock root login
	wrapped passwd -ud root
}

addpkg() {
	wrapped apk add -v $@
}

addtty() {
	for tty in $@; do
		[ -z "$tty" ] && continue
		cp -rv "$tmpdir"/usr/share/nnd/s6/dist/rc/getty.tty1 "$tmpdir"/etc/s6/rc/getty."$tty"
		sed -i 's/tty1/'"$tty"'/g' "$tmpdir"/etc/s6/rc/getty."$tty"/run

		wrapped nnd-s6 default bundle.tty
		wrapped nnd-s6 custom bundle.tty
		touch "$tmpdir"/etc/s6/rc/bundle.tty/contents.d/getty."$tty"
	done
}

addsvc() {
	# Set-up services
	for bundle in $@; do
		[ -z "$bundle" ] && continue
		touch "$tmpdir"/etc/s6/rc/default/contents.d/"$bundle"
	done
	wrapped nnd-s6 dist
	wrapped nnd-s6 generate
}

pack() {
	# Set service set as default
	wrapped nnd-s6 swap
	cleanup
}

cleanup_err() {
	cleanup
	rm -v "$image"
}

cleanup() {
	# Disable traps and cleanup
	trap '' INT TERM EXIT HUP
	umount -Rv "$tmpdir"
	rmdir -v "$tmpdir"
	rm -v "$keyfile"
	cryptsetup -q close "$CRYPT_DM"
	vgchange -a n "$VGN"
	losetup -d "$loopback"
}

addsshkey() {
	mkdir -vp "$tmpdir"/root/.ssh
	cat "$1" >> "$tmpdir"/root/.ssh/authorized_keys
}

custom() {
	echo "You are now entering your image's shell."
	echo "Any customizations made here will remain in the image." 
	echo "Upon exit, the image will be packed."
	wrapped sh || true # Handle graceful exit
}