#!/bin/sh

set -e -u
#set -x

PATH=/bin:/sbin:/usr/bin:/usr/sbin

zfs_get_root()
{
	local _res

	_res=$(
		mount -p | while read _dev _fs _type _rest ; do
			[ "/" = "$_fs" ] || continue
			if [ "zfs" != "$_type" ] ; then
				echo 'non-zfs'
				return 1
			fi
			echo "$_dev"
			return 0
		done
	)

	if [ $? -ne 0 ] ; then
		return $?
	fi

	echo $_res
	return 0
}

get_mp()
{
	local _wanted_dev
	local _res

	_wanted_dev="$1"
	_res=$(
		mount -p | while read _dev _mp _type _rest ; do
			[ "$_dev" = "$_wanted_dev" ] || continue
			echo "$_mp"
			return 0
		done
	)

	if [ $? -ne 0 ] ; then
		return $?
	fi

	echo $_res
	return 0

}

zfs_get_prop()
{
	local _prop _dataset

	_prop=$2
	_dataset=$1
	zfs get -H -p -o value $_prop $_dataset
}

zfs_get_local_props()
{
	local _dataset

	_dataset=$1
	zfs get -H -p -s local,received -o property,value all ${_dataset} | while read property value ; do
		printf '%s' "-o $property=$value "
	done
}

zpool_get_prop()
{
	local _prop _pool

	_prop=$2
	_pool=$1
	zpool get -H -p -o value $_prop $_pool
}

is_mounted()
{
	[ $(zfs_get_prop $1 mounted) = "yes" ]
}

echo_loud()
{
	printf '\a\a\a\a'
	echo ""
	echo ""
	echo "$@"
	echo ""
	echo ""
	printf '\a\a\a\a'
}

setup()
{
	if ! root_filesystem=$(zfs_get_root) ; then
		echo "Non-ZFS root filesystem" 1>&2
		return 1
	fi

	root_home=${root_filesystem%/*}
	root_pool=${root_home%/*}
	root_stamp=${root_filesystem##*/}

	#stamp=$(date +%Y%m%d-%H%M)
	stamp=$(date +%Y%m%d)
}

cmd_destroy()
{
	set -e -u

	setup

	local _be _bootfs _mp _root _origin
	local _snap

	_be=${1}

	if [ "${_be#@}" != "${_be}" ] ; then
		_snap="${root_filesystem}${_be}"
		zfs destroy -r ${_snap}
		return
	fi

	_root="${root_home}/${_be}"
	_origin=$(zfs_get_prop ${_root} origin)

	_bootfs=$(zpool_get_prop ${root_pool} bootfs)
	if [ "${_root}" = "${_bootfs}" ] ; then
		echo "Can not destroy boot BE" 1>&2
		return 1
	fi
	if is_mounted ${_root} ; then
		_mp=$(get_mp ${_root})
		if [ "${_mp}" = "/" ] ; then
			echo "Can not destroy active BE" 1>&2
			return 1
		fi
	fi
	if [ "${_origin}" = "-" ] ; then
		echo "Can not destroy trunk BE" 1>&2
		return 1
	fi

	echo_loud "Destroying $_root${_origin:+ with origin ${_origin}}"
	zfs destroy -r ${_root}
	zfs destroy -r ${_origin}
}

cmd_clone()
{
	set -e -u

	setup

	local _be _from _fromsnap new_root _new
	local _cloned_props

	_be=${1}
	_from=${2%%@*}
	_fromsnap=${2##*@}
	_from="${root_home}/${_from}"
	new_root="${root_home}/${_be}"

	echo_loud "Creating new root $new_root from ${_from}@${_fromsnap}"

	zfs list -rH -o name -s name -t filesystem ${_from} | while read _subord ; do
		_canmount=$(zfs_get_prop ${_subord} canmount)
		if [ "$_canmount" = "on" ] ; then
			# Fixup!
			zfs set canmount=noauto ${_subord}
			_canmount="noauto"
		fi

		# Get locally set properties to reproduce them in the clone.
		# XXX canmount is expected to be locally set.
		_cloned_props=$(zfs_get_local_props ${_subord})
		_new="${new_root}${_subord#${_from}}"
		zfs clone ${_cloned_props} ${_subord}@${_fromsnap} ${_new}
	done
}

cmd_create()
{
	set -e -u

	setup

	local _be _from new_root _new
	local _cloned_props
	_be=${1:-${stamp}}
	_from=${2:+${root_home}/$2}
	_from=${_from:-${root_filesystem}}
	new_root="${root_home}/${_be}"

	echo_loud "Creating new root $new_root from ${_from}@${_be}"

	zfs snapshot -r ${_from}@${_be}

	zfs list -rH -o name -s name -t filesystem ${_from} | while read _subord ; do
		_canmount=$(zfs_get_prop ${_subord} canmount)
		if [ "$_canmount" = "on" ] ; then
			# Fixup!
			zfs set canmount=noauto ${_subord}
			_canmount="noauto"
		fi

		# Get locally set properties to reproduce them in the clone.
		# XXX canmount is expected to be locally set.
		_cloned_props=$(zfs_get_local_props ${_subord})
		_new="${new_root}${_subord#${_from}}"
		zfs clone ${_cloned_props} ${_subord}@${_be} ${_new}
	done
}

cmd_activate()
{
	set -e -u

	setup

	if [ -z "$1" ] ; then
		echo "no BE is specified" 1>&2
		return 1
	fi

	_be=${1}
	new_root="${root_home}/${_be}"
#	is_mounted $new_root || zfs mount $new_root
	new_root_mp=$(get_mp $new_root)

	echo_loud "setting bootfs=$new_root on pool ${root_pool}"
	zpool set bootfs=${new_root} ${root_pool}

#	echo_loud "installing new boot blocks on ada0"	# XXX fixme using vfs.zfs.boot.primary_vdev
#	gpart bootcode -b ${new_root_mp}/boot/pmbr -p ${new_root_mp}/boot/gptzfsboot -i 1 ada0
}

cmd_promote()
{
	set -e -u

	setup

	local _root

	_root=${1:+${root_home}/$1}
	_root=${_root:-${root_filesystem}}

	zfs list -rH -o name -s name -t filesystem ${_root} | while read _subord ; do
		if ! out=$(zfs promote ${_subord} 2>&1) ; then
			case "$out" in
			*'not a cloned filesystem'*) ;; # OK
			*)
				echo "$out" 1>&2
				return 1
			;;
			esac
		fi
	done
}

#TODO make beadm compliant, show subordinates
cmd_list()
{
	set -e -u

	setup

	local _bootfs

	_bootfs=$(zpool_get_prop ${root_pool} bootfs)

	{
	printf "BE ACTIVE MOUNTPOINT MOUNTED USED CREATED\n"
	zfs list -rH -d1 -o name,mountpoint,creation -s name -t filesystem ${root_home} | while read _fs _mp _creation; do
		local _used _origin

		_be=${_fs#${root_home}}
		_mounted='N'
		_active='N'

		[ -z "${_be}" ] && continue
		_be=${_be#/}

		if is_mounted ${_fs} ; then
			_mounted='Y'
			_mp=$(get_mp ${_fs})
			if [ "${_mp}" = "/" ] ; then
				_active='Y'
			fi
		fi
		if [ "${_fs}" = "${_bootfs}" ] ; then
			_active="${_active}R"
		fi
		_origin=$(zfs_get_prop ${_fs} origin)
		if [ "${_origin}" != '-' ] ; then
			_used=$(zfs destroy -n -v -R ${_origin} | sed -n 's/would reclaim //p')
		else
			_used='-'
		fi
		_creation=$(echo ${_creation} | sed 's/ /^/g')
		printf "%s %s %s %s %s %s\n" ${_be} ${_active} ${_mp} ${_mounted} ${_used} ${_creation}
	done
	} | column -t | sed 's/\^/ /g'
}

cmd_etcupgrade()
{
	set -e -u

	setup

	if [ -z "$1" ] ; then
		echo "no BE is specified" 1>&2
		return 1
	fi

	_be=${1}
	new_root="${root_home}/${_be}"

	if [ ! -d sys ] || [ ! -f Makefile ] ; then
		echo "Current directory doesn't look like source tree root" 1>&2
		return 1
	fi

	new_root_mp="/${new_root}"

	echo_loud "running etcupdate for $new_root_mp"
	_etcupdate_destdir_opts="-D ${new_root_mp} -d ${new_root_mp}/usr/local/var/db/etcupdate"
	etcupdate -F ${_etcupdate_destdir_opts}
	echo etcupdate status ${_etcupdate_destdir_opts}
	etcupdate status ${_etcupdate_destdir_opts}
}

cmd_srcupgrade()
{
	set -e -u

	setup

	if [ -z "$1" ] ; then
		echo "no BE is specified" 1>&2
		return 1
	fi

	_be=${1}
	new_root="${root_home}/${_be}"

	if [ ! -d sys ] || [ ! -f Makefile ] ; then
		echo "Current directory doesn't look like source tree root" 1>&2
		return 1
	fi

	new_root_mp="/${new_root}"

	rm -rf ${new_root_mp}/boot/kernel*
	rm -rf ${new_root_mp}/*.core

	echo_loud "installworld => $new_root_mp"
	make -s installworld DESTDIR=${new_root_mp}
	echo_loud "installkernel => $new_root_mp"
	make -s installkernel DESTDIR=${new_root_mp}
	echo_loud "install boot => $new_root_mp"
	( cd sys/boot && make install DESTDIR=${new_root_mp} )
}

mount_fstab()
{
	set -e -u

	new_root_mp=$1
	mount_late=$2

	cat $new_root_mp/etc/fstab | while read _dev _mp _type _opts _etc ; do
		case "$_dev" in '#'*|'' ) continue ;; esac
		[ "$_type" = "swap" ] && continue
		[ "$_type" = "nfs" ] && continue
		case "$_opts" in *noauto* ) continue ;; esac

		case "$_opts" in
		*,late,*|late,*|*,late )
			[ "$mount_late" != "yes" ] && continue
			;;
		*)
			[ "$mount_late" = "yes" ] && continue
			;;
		esac

		case "$_type" in
		procfs | linprocfs | fdescfs | tmpfs)
			mount -t "$_type" "$_dev" $new_root_mp$_mp
			;;
		*)
			if is_mounted "$_dev" ; then
				mount -t nullfs $(get_mp $_dev) $new_root_mp$_mp
			else
				mount -t "$_type" "$_dev" $new_root_mp$_mp
			fi
			;;
		esac
	done
}

cmd_mount()
{
	set -e -u

	setup

	if [ -z "$1" ] ; then
		echo "no BE is specified" 1>&2
		return 1
	fi

	local _be new_root _where
	_be=${1}
	new_root="${root_home}/${_be}"

	# mount BE root
	is_mounted $new_root || zfs mount $new_root

	# mount devfs
	mount -t devfs devfs /$new_root/dev

	# mount filesystems from BE's fstab
	mount_fstab /$new_root "no"

	# mount BE subordinates
	zfs list -rH -o mountpoint,name,canmount -s mountpoint -t filesystem $new_root | \
	while read _mp _name _canmount ; do
		# skip filesystems that must not be mounted
		[ "$_canmount" = "off" ] && continue

		case "$_mp" in
		"" | "legacy" | "/" | "/$new_root")
			# do nothing for filesystems with unset or legacy mountpoint
			# or those that would be mounted over /
			;;
		"/$new_root/"*)
			# filesystems with mountpoint relative to BE
			is_mounted $_name || zfs mount $_name
			;;
		*)
			# filesystems with mountpoint elsewhere
			if is_mounted $_name ; then
				_where=$(get_mp $_name)
				if [ "${_where#/$new_root}" = "$_where" ] ; then
					# mounted outside of BE hierarchy
					[ ! -d /$new_root$_mp ] && mkdir -p /$new_root$_mp
					mount -t nullfs $(get_mp $_name) /$new_root$_mp
				fi
			else
				# not mounted
				[ ! -d /$new_root$_mp ] && mkdir -p /$new_root$_mp
				 mount -t zfs $_name /$new_root$_mp
			fi
			;;
		esac
	done

	# mount shared (non-BE) filesystems
	zfs list -H -o mountpoint,name,canmount -s mountpoint -t filesystem | \
	while read _mp _name _canmount ; do
		# skip filesystems that are not auto mounted
		[ "$_canmount" = "on" ] || continue

		# skip all filesystems belonging to BEs
		case $_name in $root_home/*) continue ;; esac

		case "$_mp" in
		"" | "legacy" | "/" | "/$new_root")
			# do nothing for filesystems with unset or legacy mountpoint
			# or those that would be mounted over /
			;;
		*)
			if is_mounted $_name ; then
				_where=$(get_mp $_name)
				if [ "${_where#/$new_root}" = "$_where" ] ; then
					# mounted outside of BE hierarchy
					[ ! -d /$new_root$_mp ] && mkdir -p /$new_root$_mp
					mount -t nullfs $(get_mp $_name) /$new_root$_mp
				fi
			else
				# not mounted
				[ ! -d /$new_root$_mp ] && mkdir -p /$new_root$_mp
				 mount -t zfs $_name /$new_root$_mp
			fi
			;;
		esac
	done

	# mount late filesystems from BE's fstab
	mount_fstab /$new_root "yes"
}

cmd_umount()
{
	set -e -u

	setup

	if [ -z "$1" ] ; then
		echo "no BE is specified" 1>&2
		return 1
	fi

	local _be _force new_root _where
	_be=${1}
	_force=${2:+'-f'}
	new_root="${root_home}/${_be}"

	# if BE is not mounted, then we have nothing to do
	is_mounted $new_root || return 0
	_where=$(get_mp $new_root)

	if [ -z "$_where" ] ; then
		echo "failed to determined where BE $_be is mounted" 1>&2
		return 1
	fi
	if [ "$_where" = "/" ] ; then
		echo "can not unmount current root BE" 1>&2
		return 1
	fi
	mount -p | tail -r | while read _dev _mp _type _rest ; do
		case $_mp in
		$_where | $_where/*)
			umount $_force $_mp
			;;
		esac
	done
}

cmd_chroot()
{
	set -e -u

	setup

	local _be _root _where

	_be=${1}
	shift

	_root="${root_home}/${_be}"

	# mount BE root
	if ! is_mounted ${_root} ; then
		echo "${_root} is not mounted" 1>&2
		return 1
	fi

	_where=$(get_mp ${_root})
	if [ -z "$_where" ] ; then
		echo "failed to determined where BE $_be is mounted" 1>&2
		return 1
	fi
	if [ "$_where" = "/" ] ; then
		echo "refusing to chroot to /" 1>&2
		return 1
	fi

	chroot ${_where} "$@"
}

cmd_snapshot()
{
	set -e -u

	setup

	local _snap _be

	_snap=${1}
	_be=${2:+${root_home}/$2}
	_be=${_be:-${root_filesystem}}

	zfs snapshot -r ${_be}@${_snap}
}

cmd_rollback()
{
	set -e -u

	setup

	local _snap _be

	_snap=${1}
	_be=${2:+${root_home}/$2}
	_be=${_be:-${root_filesystem}}

	zfs list -rH -o name -s name -t filesystem ${_be} | while read _subord ; do
		zfs rollback ${_subord}@${_snap}
	done
}

cmd_Xupgrade()
{
	set -e -u

	local _newbe

	_newbe=$(date +%Y%m%d)
	cmd_create $_newbe
	cmd_mount $_newbe
	cmd_srcupgrade $_newbe
}

cmd=$1
shift
eval cmd_$cmd "$@"