#!/bin/sh
# shellcheck disable=SC2034,SC2154

. /lib/dracut-zfs-lib.sh

decode_root_args || return 0

GENERATOR_FILE=/run/systemd/generator/sysroot.mount
GENERATOR_EXTENSION=/run/systemd/generator/sysroot.mount.d/zfs-enhancement.conf

if [ -e "$GENERATOR_FILE" ] && [ -e "$GENERATOR_EXTENSION" ]; then
	# We're under systemd and dracut-zfs-generator ran to completion.
	info "ZFS: Delegating root mount to sysroot.mount at al."

	# We now prevent Dracut from running this thing again.
	rm -f "$hookdir"/mount/*zfs*
	return
fi

info "ZFS: No sysroot.mount exists or zfs-generator did not extend it."
info "ZFS: Mounting root with the traditional mount-zfs.sh instead."

# Delay until all required block devices are present.
modprobe zfs 2>/dev/null
udevadm settle

ZFS_DATASET=
ZFS_POOL=

if [ "${root}" = "zfs:AUTO" ] ; then
	if ! ZFS_DATASET="$(find_bootfs)" ; then
		# shellcheck disable=SC2086
		zpool import -N -a ${ZPOOL_IMPORT_OPTS}
		if ! ZFS_DATASET="$(find_bootfs)" ; then
			warn "ZFS: No bootfs attribute found in importable pools."
			zpool export -aF

			rootok=0
			return 1
		fi
	fi
	info "ZFS: Using ${ZFS_DATASET} as root."
fi

ZFS_DATASET="${ZFS_DATASET:-${root}}"
ZFS_POOL="${ZFS_DATASET%%/*}"

if import_pool "${ZFS_POOL}" ; then
	# Load keys if we can or if we need to
	if [ "$(zpool list -H -o feature@encryption "${ZFS_POOL}")" = 'active' ]; then
		# if the root dataset has encryption enabled
		ENCRYPTIONROOT="$(zfs get -H -o value encryptionroot "${ZFS_DATASET}")"
		if ! [ "${ENCRYPTIONROOT}" = "-" ]; then
			KEYSTATUS="$(zfs get -H -o value keystatus "${ENCRYPTIONROOT}")"
			# if the key needs to be loaded
			if [ "$KEYSTATUS" = "unavailable" ]; then
				# decrypt them
				ask_for_password \
					--tries 5 \
					--prompt "Encrypted ZFS password for ${ENCRYPTIONROOT}: " \
					--cmd "zfs load-key '${ENCRYPTIONROOT}'"
			fi
		fi
	fi
	# Let us tell the initrd to run on shutdown.
	# We have a shutdown hook to run
	# because we imported the pool.
	info "ZFS: Mounting dataset ${ZFS_DATASET}..."
	if mount_dataset "${ZFS_DATASET}" ; then
		ROOTFS_MOUNTED=yes
		return 0
	fi
fi

rootok=0