#!/usr/bin/env bash
# Copyright (C) 2018  Luke Shumaker
# Copyright (C) 2023  Umorpha Systems
# SPDX-License-Identifier: AGPL-3.0-or-later
declare -r NAME=osi-mk
declare -r VERSION=20231023

# Why is this different than mkosi[1]?
#
# - mkosi claims to be "legacy-free"--but they call everything that's
#   not systemd "legacy"; that clearly won't do for creating OpenRC
#   images.
#
# - mkosi claims to be "legacy-free", only supporting GPT
#   disk-labels--but btrfs can be booted directly, without a separate
#   disk-label, or a separate ESP partition.  To a btrfs disk, GPT/ESP
#   is legacy.
#
# - Using a raw btrfs disk means that it can easily be mounted without
#   first dissecting the disk-label.
#
# [1]: https://github.com/systemd/mkosi

set -euE -o pipefail
source "$(dirname -- "${BASH_SOURCE[0]}")/lib/osi.sh"

loaded_modules=()
load_module() {
	local module
	if ! [[ -f $1 ]]; then
		error $EXIT_INVALIDARGUMENT 'Module does not exist: %s' "$1"
	fi
	module="$(realpath -- "$1")"
	if in_array "$module" "${loaded_modules[@]}"; then
		return 0
	fi
	loaded_modules+=("$module")
	# shellcheck source=/dev/null
	source "$1"
}

osi-mk:genfstab() {
	local arg_mountpoint=$1
	{
		# This header mimics the stock fstat from 'filesystem'.
		echo '# Static information about the filesystems.'
		echo '# See fstab(5) for details.'
		echo
		{
			# This doesn't quite mimic the header from 'filesystem'
			#
			# - I swapped "file system" for "device", for clarity,
			#   and so that it's one word for `column -t`.
			# - I swapped "pass" for "fsck", for clarity.
			# - I removed the space after the "#", so it works with
			#   `column -t`.
			echo '#<device> <mountpoint> <type> <options> <dump> <fsck>'
			case "${arg_conf[format]}" in
				erofs)
					echo "UUID=${arg_conf[fsuuid]}" / erofs defaults 0 1
					;;
				*)
					genfstab -t uuid "$arg_mountpoint" | grep '^[#]' | awk '$3 != "swap"'
					;;
			esac
		} | column -t
	} >"${arg_mountpoint}/etc/fstab"
}

osi-mk:directories() {
	local arg_mountpoint=$1
	local spec outside inside
	for spec in "${arg_directories[@]}"; do
		outside="${spec%:*}"
		inside="${spec#"${outside}:"}"
		# TODO: maybe rsync would be better?
		mkdir -p -- "$(dirname -- "${arg_mountpoint}/${inside}")"
		print 'Copying %q to %q:%q' "$outside" "$arg_file" "$inside"
		cp -aT --no-preserve=ownership -- "$outside" "${arg_mountpoint}/${inside}"
	done
}

osi-mk:hostname() {
	local arg_mountpoint=$1

	echo "${arg_conf[hostname]}" >"$arg_mountpoint/etc/hostname"
}

osi-mk:mkinitcpio() {
	local arg_mountpoint=$1
	print 'Configuring mkinitcpio to include all drivers'
	mkdir -p -- "$arg_mountpoint/etc/mkinitcpio.conf.d"
	cat >"$arg_mountpoint/etc/mkinitcpio.conf.d/drivers.conf" <<'EOF'
#!/hint/bash
# Copyright (C) 2023  Umorpha Systems
# SPDX-License-Identifier: AGPL-3.0-or-later

# Remove 'autodetect' from HOOKS; include all drivers.
for ((i=0; i<${#HOOKS[@]}; i++)); do
	if [[ ${HOOKS[i]} == autodetect ]]; then
		HOOKS=("${HOOKS[@]:0:i}" "${HOOKS[@]:i+1}" );
	fi
done
EOF
}

osi-mk:grub-install() {
	local arg_mountpoint=$1
	cat <<-'EOT' >> "$arg_mountpoint/etc/default/grub"
			GRUB_TIMEOUT=0
			EOT
	# shellcheck disable=SC2016
	arch-chroot -- "$arg_mountpoint" sh -c \
		    'grub-install --target=i386-pc "$(awk '\''$2 == "/" { print $1 }'\'' </proc/mounts)"'
}

osi-mk:grub-mkconfig() {
	local arg_mountpoint=$1
	arch-chroot -- "$arg_mountpoint" grub-mkconfig -o /boot/grub/grub.cfg
}

main() {
	local arg_orig=("$@")
	local arg_mode=outside
	local arg_mountpoint=
	local arg_edit=false
	local arg_edit_base=

	local arg_directories=("$(dirname -- "${BASH_SOURCE[0]}")/osi-mk.d/grub-hack":usr/bin)
	local arg_modules=()
	local arg_packages=()
	local arg_package_files=()
	declare -A arg_conf=(
		[format]=raw_btrfs
		[bootloader]=grub
	)

	local args
	if ! args="$(getopt -n "${0##*/}" -o "s:e::d:m:p:P:C:hV" -l "inside::,edit::,directory:,module:,package:,package-file:,conf:,help,version" -- "$@")"; then
		arg_mode=error
	else
		eval "set -- $args"
		while true; do
			case "$1" in
				--inside) shift; arg_mode=inside; arg_mountpoint=$1; shift;;
				-e|--edit) shift; arg_edit=true; arg_edit_base=$1; shift;;

				-d|--directory) shift; arg_directories+=("$1"); shift;;
				-m|--module) shift; arg_modules+=("$1"); shift;;
				-p|--package) shift; arg_packages+=("$1"); shift;;
				-P|--package-file) shift; arg_package_files+=("$1"); shift;;
				-C|--conf) shift; arg_conf["${1%%=*}"]="${1#*=}"; shift;;

				-V|--version) shift; arg_mode=version;;
				-h|--help) shift; arg_mode=usage;;
				--) shift; break;;
				*) error $EXIT_FAILURE 'Internal error.  The programmer writing this tool screwed up.';;
			esac
		done
		case "$arg_mode" in
			outside|inside)
				if (( $# != 1 )); then
					if (( $# == 0 )); then
						error 0 "Expected 1 positional argument, got none"
					else
						error 0 "Expected 1 positional argument, got %d: %s" "$#" "${*@Q}"
					fi
					arg_mode=error
				else
					arg_file=$1
				fi
				for module in "${arg_modules[@]}"; do
					if ! [[ -f $module ]]; then
						error 0 'Module does not exist: %s' "$module"
						arg_mode=error
					fi
				done
				for dirspec in "${arg_directories[@]}"; do
					if ! [[ -d "${dirspec%:*}" ]]; then
						error 0 'Directory does not exist: %s' "${dirspec%:*}"
						arg_mode=error
					fi
				done
				for package_file in "${arg_package_files[@]}"; do
					if ! [[ -f "$package_file" ]]; then
						error 0 'Package file does not exist: %s' "$package_file"
						arg_mode=error
					fi
				done
				if [[ $arg_mode = outside ]]; then
					if [[ ( $arg_edit = false || -n $arg_edit_base ) && -e $arg_file ]]; then
						error $EXIT_INVALIDARGUMENT 'Image file already exists, refusing to overwrite: %s' "$arg_file"
					fi
					if $arg_edit; then
						if ! [[ -f ${arg_edit_base:-$arg_file} ]]; then
							error $EXIT_INVALIDARGUMENT 'Image must already exist to --edit: %s' "${arg_edit_base:-$arg_file}"
						fi
					fi
					case "${arg_conf[format]}" in
						raw_btrfs)
							if [[ -z ${arg_conf[size]:-} ]]; then
								error $EXIT_INVALIDARGUMENT 'Must specify --conf=size= when creating a new image'
							fi
							;;
						erofs)
							if $arg_edit; then
								error $EXIT_INVALIDARGUMENT '--conf=format=erofs does not support --edit'
							fi
							if [[ ${arg_conf[bootloader]} != none ]]; then
								error $EXIT_INVALIDARGUMENT '--conf=format=erofs requires --conf=bootloader=none'
							fi
							;;
						*)
							error $EXIT_INVALIDARGUMENT 'Unrecognized --conf=format= value: %q' "${arg_conf[format]}"
							;;
					esac
					case "${arg_conf[bootloader]}" in
						grub) :;;
						none) :;;
						*)
							error $EXIT_INVALIDARGUMENT 'Unrecognized --conf=bootloader= value: %q' "${arg_conf[bootloader]}"
							;;
					esac
				fi
				;;
		esac
	fi
	case "$arg_mode" in
		error)
			print "Try '%q --help' for more information" "${0##*/}" >&2
			return $EXIT_INVALIDARGUMENT
			;;
		version)
			print "%s (osi-tools) %s" "$NAME" "$VERSION"
			return $EXIT_SUCCESS
			;;
		usage)
			print 'Usage: %s [OPTIONS] FILENAME.img' "${0##*/}"
			print 'Operating System Image: Make'
			echo
			print 'Create a mountable, bootable OS image.'
			echo
			print 'OPTIONS:'
			#     000000000011111111112222222222333333333344444444445555555555666666666677777777778
			#     012345678901234567890123456789012345678901234567890123456789012345678901234567890
			print '  -s SIZE,           --size=SIZE                    set the size of the image'
			print '  -e[BASE.img],      --edit[=BASE.img]              edit an existing image'
			# --inside is internal-only; undocumented
			echo
			print '  -d OUTSIDE:INSIDE, --directory=OUTSIDE:INSIDE     include the given directory'
			print '  -m MOD.sh,         --module=MOD.sh                include the given module'
			print '  -p PKGNAME,        --package=PKGNAME              include the given package (or group)'
			print '  -P PKG.pkg.tar.xz, --package-file=PKG.pkg.tar.xz  include the given package file'
			print '  -C KEY=VAL,        --conf=KEY=VAL                 set set a config option'
			echo
			print '  -h,                --help                         display this help'
			print '  -V,                --version                      output version information'
			return $EXIT_SUCCESS
			;;

		# main code starts here
		outside) # the part that runs as a normal user without the image mounted
			mount_dev="$arg_file"
			mount_opt=''
			gprintf -v prefix '%s [format]' "$NAME"
			{
				case "${arg_conf[format]}" in
					raw_btrfs)
						if $arg_edit; then
							if [[ -n $arg_edit_base ]]; then
								cp -T --reflink -- "$arg_edit_base" "$arg_file"
							fi
							if [[ -n ${arg_conf[fsuuid]:-} ]]; then
								# Explicit new FSUUID.
								btrfstune -f -U "${arg_conf[fsuuid]:-}" "$arg_file"
							elif [[ -n $arg_edit_base ]]; then
								# Random new FSUUID.
								btrfstune -fu "$arg_file"
							fi
							if [[ -n ${arg_conf[size]:-} ]]; then
								# Calculate sizes in exact bytes now, to avoid any discrepancies
								# between truncate(1) and btrfs-filesystem(8).
								# Use truncate(1) to do this.  It's a little gross,
								# but because TMPDIR is likely on tmpfs, it's time-cheap,
								# and because of sparse files, it's space-cheap.
								tmpfile=$(mktemp -t -- "${0##*/}.XXXXXXXXXX")
								trap "rm -- ${tmpfile@Q}" EXIT
								case "${arg_conf[size]}" in
									'+'*|'-'*|'<'*|'>'*|'/'*|'%'*)
										truncate --reference="$arg_file" --size="${arg_conf[size]}" -- "$tmpfile";;
									*)
										truncate --size="${arg_conf[size]}" -- "$tmpfile";;
								esac
								old_size=$(stat --format='%s' -- "$arg_file")
								new_size=$(stat --format='%s' -- "$tmpfile")
								rm -- "$tmpfile"
								trap - EXIT

								# Do the resize
								arg_mountpoint=$(mktemp -dt -- "${0##*/}.XXXXXXXXXX")
								trap "rmdir -- ${arg_mountpoint@Q}" EXIT
								if (( new_size > old_size )); then
									truncate --size="$new_size" -- "$arg_file"
									sudo -- "$(dirname -- "${BASH_SOURCE[0]}")/osi-mount" --root -- "$arg_file" "$arg_mountpoint" btrfs filesystem resize max "$arg_mountpoint"
								elif (( new_size < old_size )); then
									sudo -- "$(dirname -- "${BASH_SOURCE[0]}")/osi-mount" --root -- "$arg_file" "$arg_mountpoint" btrfs filesystem resize "$new_size" "$arg_mountpoint"
									truncate --size="$new_size" -- "$arg_file"
								fi
								rmdir -- "$arg_mountpoint"
								trap - EXIT
							fi
						else
							truncate --size="${arg_conf[size]}" -- "$arg_file"
							mkfs.btrfs ${arg_conf[fsuuid]:+"--uuid=${arg_conf[fsuuid]}"} -- "$arg_file"
						fi
						;;
					erofs)
						if $arg_edit; then
							error 2 '--conf=format=erofs does not support --edit'
						fi
						if [[ -z "${arg_conf[fsuuid]:-}" ]]; then
							arg_conf[fsuuid]=$(uuidgen)
							arg_orig=(--conf=fsuuid="${arg_conf[fsuuid]}" "${arg_orig[@]}")
						fi

						mount_dev=$(mktemp -dt -- "${0##*/}.XXXXXXXXXX")
						# Make our own tmpfs because $TMPDIR might be mounted
						# with 'noexec' or something.  Also, so we can clean it
						# up without a `sudo rm -rf`.
						sudo mount -t tmpfs tmpfs:osi-mk "$mount_dev"
						trap "sudo umount ${mount_dev@Q}; rmdir -- ${mount_dev@Q}" EXIT
						mount_opt=bind
						;;
				esac
			} >& >(sed "s|^|${prefix} |")

			arg_mountpoint=$(mktemp -dt -- "${0##*/}.XXXXXXXXXX")
			r=0
			sudo -- "$(dirname -- "${BASH_SOURCE[0]}")/osi-mount" --root \
			     --rwdir=/var/cache/pacman/pkg \
			     --options="$mount_opt" \
			     -- \
			     "$mount_dev" "$arg_mountpoint" \
			     "${BASH_SOURCE[0]}" --inside="$arg_mountpoint" \
			     "${arg_orig[@]}" || r=$?
			rmdir -- "$arg_mountpoint"
			if (( r != 0 )); then
				return $r
			fi

			case "${arg_conf[format]}" in
				erofs)
					tmpdir=$(mktemp -dt -- "${0##*/}.XXXXXXXXXX")
					trap "sudo umount ${mount_dev@Q}; rmdir -- ${mount_dev@Q}; rm -f -- ${tmpdir@Q}/erofs.img; rmdir -- ${tmpdir@Q}" EXIT
					sudo mkfs.erofs \
					     -U"${arg_conf[fsuuid]}" \
					     ${arg_conf[erofs_compression]:+"-z${arg_conf[erofs_compression]}"} \
					     "$tmpdir/erofs.img" "$mount_dev"
					cp -T "$tmpdir/erofs.img" "$arg_file"
					;;
			esac
			;;
		inside) # the part that runs as root with the image mounted
			needs_sudo

			### Load modules ###
			packages=("${arg_packages[@]}")
			pre_install=()
			post_install=(
				50:osi-mk:directories
			)

			if ! $arg_edit || [[ -n $arg_edit_base ]] || [[ -n "${arg_conf[fsuuid]:-}" ]]; then
				post_install+=(00:osi-mk:genfstab)
			fi
			case "${arg_conf[format]}" in
				raw_btrfs)
					packages+=(btrfs-progs)
					;;
				erofs)
					packages+=(erofs-utils)
					;;
			esac
			case "${arg_conf[bootloader]}" in
				grub)
					packages+=(grub)
					post_install+=(99:osi-mk:grub-mkconfig)
					if ! $arg_edit; then
						post_install+=(98:osi-mk:grub-install) # before '99:osi-mk:grub-mkconfig'
					fi
					;;
			esac
			case "${arg_conf[initramfs]:-}" in
				mkinitcpio)
					pre_install+=(00:osi-mk:mkinitcpio)
					;;
			esac
			if [[ -n "${arg_conf[hostname]:-}" ]]; then
				post_install+=(00:osi-mk:hostname)
			fi

			for module in "${arg_modules[@]}"; do
				load_module "$module"
			done
			cache_packages+=("${packages[@]}")

			### Download ###
			gprintf -v prefix '%s [download:repos]' "$NAME"
			{
				# Download syncdbs to the image
				mkdir -p -- "$arg_mountpoint"/var/{lib/pacman,log}
					pacman -r "$arg_mountpoint" --config=/usr/share/pacman/defaults/pacman.conf.x86_64 \
					       -Sy --noconfirm
			} |& sed "s|^|${prefix} |"
			gprintf -v prefix '%s [download:check-config]' "$NAME"
			{
				# Validate that the user specified an answer to
				# every question that pacman asks.
				opt_fail=false
				while read -r -a options; do
					name="${options[0]%:}"
					options=("${options[@]:1}")

					base=${name%%[<>=]*}
					if in_array "${arg_conf["$base"]:-}" "${options[@]}" "${options[@]#*/}"; then
						print 'inserting package %s=%s' "$name" "${arg_conf["$base"]}"
						packages+=("${arg_conf["$base"]}")
						cache_packages+=("${arg_conf["$base"]}")
					else
						error 0 "must set option '%q' to one of [%s]" "$base" "${options[*]}"
						opt_fail=true
					fi
				done < <("$(dirname -- "${BASH_SOURCE[0]}")/pacman-choices" \
						 -r "$arg_mountpoint" --config=/usr/share/pacman/defaults/pacman.conf.x86_64 \
						 --oneline -- "${cache_packages[@]}")
				if [[ $opt_fail == true ]]; then
					exit $EXIT_NOTCONFIGURED
				fi
			} >& >(sed "s|^|${prefix} |")
			gprintf -v prefix '%s [download:packages]' "$NAME"
			{
				if (( ${#cache_packages[@]} > 0 )); then # this check is important for --edit
					# Download needed packages to the host cache
					pacman -r "$arg_mountpoint" --config=/usr/share/pacman/defaults/pacman.conf.x86_64 \
					       -Syw --noconfirm -- "${cache_packages[@]}"

					# Copy needed packages from host cache to image cache
					mkdir -p -- "$arg_mountpoint"/var/cache/pacman/pkg
					pacman -r "$arg_mountpoint" --config=/usr/share/pacman/defaults/pacman.conf.x86_64 \
					       -Sp --print-format='%l' -- "${cache_packages[@]}" \
						| sed -En 's,^file://(.*),\1\n\1.sig,p' \
						| xargs -d $'\n' -r cp -t "$arg_mountpoint/var/cache/pacman/pkg" --
				fi
			} |& sed "s|^|${prefix} |"

			### pre_install ###
			while IFS=: read -r n fn; do
				printf -v prefix '%s [pre_install:%s:%s]' "$NAME" "$n" "$fn"
				{
					print Begin
					"$fn" "$arg_mountpoint"
					print End
				} |& sed "s|^|${prefix} |"
			done < <(printf '%s\n' "${pre_install[@]}" | sort)

			### Install ###
			gprintf -v prefix '%s [install]' "$NAME"
			{
				if (( ${#packages[@]} > 0 )); then # this check is important for --edit
					# The --hookdir bit is to hack around https://bugs.archlinux.org/task/49347
					pacstrap -M -C /usr/share/pacman/defaults/pacman.conf.x86_64 -- "$arg_mountpoint" \
						 --hookdir="$arg_mountpoint/etc/pacman.d/hooks" \
						 --needed \
						 "${packages[@]}"
				fi
				if (( ${#arg_package_files[@]} > 0 )); then
					dir="$(mktemp -d -- "$arg_mountpoint/tmp/package-files.XXXXXXXXXX")"
					trap "rm -rf -- ${dir@Q}" EXIT
					cp -t "$dir" -- "${arg_package_files[@]}"
					arch-chroot -- "$arg_mountpoint" pacman -U --noconfirm -- "${arg_package_files[@]/#*\//"/tmp/${dir##*/tmp/}"}"
					rm -rf -- "$dir"
					trap - EXIT
				fi
			} |& sed "s|^|${prefix} |"

			### post_install ###
			while IFS=: read -r n fn; do
				printf -v prefix '%s [post_install:%s:%s]' "$NAME" "$n" "$fn"
				{
					print Begin
					"$fn" "$arg_mountpoint"
					print End
				} |& sed "s|^|${prefix} |"
			done < <(printf '%s\n' "${post_install[@]}" | sort)

			### End ###
			print '%s Done' "$NAME"
			;;

		*) error $EXIT_FAILURE 'Internal error.  The programmer writing this tool screwed up.';;
	esac
}

main "$@"