#compdef zfs

_zfs() {
	local context state line expl
	typeset -A opt_args
	local -a subcmds rw_properties rw_propnames ro_properties

	subcmds=(
		"create" "destroy" "clone" "promote" "rename" "snapshot"
		"rollback" "list" "set" "get" "inherit" "mount" "unmount"
		"share" "unshare" "send" "receive" "allow" "unallow"
		"upgrade"
	)

	ro_properties=(
		"name" "type" "creation" "used" "available" "referenced"
		"compressratio" "mounted" "origin"
	)

	# TODO: Be cleverer about what values can be set.  Is there any way to
	# set the sorting for *size properties to false by default?
	rw_properties=(
		"aclinherit:value:(discard noallow secure passthrough)"
		"aclmode:value:(discard groupmask passthrough)"
		"atime:value:(on off)"
		"canmount:value:(on off)"
		"checksum:value:(on off fletcher2 fletcher4 sha256)"
		"compression:value:(on off lzjb gzip gzip-{1..9})"
		"copies:value:(1 2 3)"
		"devices:value:(on off)"
		"exec:value:(on off)"
		"mountpoint:value:{if [[ -prefix /* ]]; then _path_files -/; else _wanted mountpoints expl 'mountpoint (type \"/\" to start completing paths)' compadd legacy none; fi}"
		"quota:number or 'none':{if [[ -prefix [0-9]## ]]; then _message -e 'number'; elif [[ $PREFIX == quota= ]]; then _wanted none expl 'number or none' compadd none; else _wanted none expl 'quota' compadd none; fi}"
		"readonly:value:(on off)"
		"recordsize:value:(512 1K 2K 4K 8K 16K 32K 64K 128K)"
		"reservation:value:" # <size>, "none"
		"setuid:value:(on off)"
		"shareiscsi:value:(on off)" # or "type=<type>"
		"sharenfs:value:(on off)" # or share(1M) options
		"snapdir:value:(hidden visible)"
		"version:value:(1 2 current)"
		"volblocksize:value:(512 1K 2K 4K 8K 16K 32K 64K 128K)"
		"volsize:value:" # <size>
		"xattr:value:(on off)"
		"zoned:value:(on off)"
	)

	rw_propnames=( ${rw_properties%%:*} )

	if [[ $service == "zfs" ]]; then
		_arguments -C -A "-*" \
			'-\?[Help]' \
			'*::command:->subcmd' && return 0

		if (( CURRENT == 1 )); then
			_wanted commands expl "zfs subcommand" compadd -a subcmds
			return
		fi
		service="$words[1]"
		curcontext="${curcontext%:*}=$service:"
	fi

	case $service in
	("create")
		_arguments -A "-*" \
			'-p[Create parent datasets]' \
			'-o[Set initial properties]:property:_values -s , "property" $rw_properties' \
			- set1 \
			':filesystem:' \
			- set2 \
			'-s[Create sparse volume]' \
			'-b[Set volblocksize]:blocksize:' \
			'-V[Set size]:size:' \
			':volume:'
		;;

	("destroy")
		_arguments -A "-*" \
			'-r[Recursively destroy all children]' \
			'-R[Recursively destroy all dependents]' \
			'-f[Force unmounts]' \
			':filesystem/volume/snapshot:_zfs_dataset'
		;;

	("snapshot")
		_arguments -A "-*" \
			'-r[Recursively snapshot all descendant datasets]' \
			':filesystem/volume:_zfs_dataset -t fs -t vol -S@'
		;;

	("rollback")
		_arguments -A "-*" \
			'-r[Recursively destroy more recent snapshots]' \
			'-R[Recursively destroy more recent snapshots and clones]' \
			'-f[Force unmounts]' \
			':snapshot:_zfs_dataset -t snap'
		;;

	("clone")
		# XXX needs to bail if there are no snapshots
		_arguments -A "-*" \
			'-p[Create parent datasets]' \
			':snapshot:_zfs_dataset -t snap' \
			':filesystem/volume:'
		;;

	("promote")
		# XXX complete only cloned filesystems
		_arguments \
			':filesystem:_zfs_dataset -t fs' \
		;;

	("rename")
		_arguments -A "-*" \
			- set1 \
			'-p[Create parent datasets]' \
			':filesystem/volume/snapshot:_zfs_dataset' \
			':filesystem/volume/snapshot:' \
			- set2 \
			'-r[Recursively rename snapshots of all descendent datasets]' \
			':snapshot:_zfs_dataset -t snap' \
			':snapshot:'
		;;

	("list")
		_arguments -A "-*" \
			'-r[Recursively display children]' \
			'-H[Scripting mode]' \
			'-o[Properties to list]:property:_values -s , "property" $ro_properties $rw_propnames' \
			'*-s[Sort key (ascending)]:property:_values "property" $ro_properties $rw_propnames' \
			'*-S[Sort key (descending)]:property:_values "property" $ro_properties $rw_propnames' \
			'-t[Dataset types to list]:dataset type:_values -s , "dataset type" filesystem snapshot volume' \
			'*:filesystem/volume/snapshot:_zfs_dataset'
		;;

	("set")
		_arguments \
			':property:_values -s , "property" $rw_properties' \
			'*:filesystem/volume:_zfs_dataset -t fs -t vol'
		;;

	("get")
		_arguments -A "-*" \
			"-r[Recursively display children's properties]" \
			'-H[Scripting mode]' \
			'-p[Display numbers exactly]' \
			'-s[Specify sources]:source:_values -s , "source" local default inherited temporary none' \
			'-o[Specify fields]:field:_values -s , "field" name property value source' \
			':property:_values -s , "property" $ro_properties $rw_propnames all' \
			'*:filesystem/volume/snapshot:_zfs_dataset'
		;;

	("inherit")
		_arguments -A "-*" \
			'-r[Recursively inherit property for all children]' \
			':property:_values -s , "property" $ro_properties $rw_properties' \
			'*:filesystem/volume:_zfs_dataset -t fs -t vol'
		;;

	("mount")
		_arguments -A "-*" \
			'-o[Mount options]:mount options:_values -s , "option" {,no}{devices,exec,setuid} ro rw' \
			'-O[Overlay mount]' \
			'-v[Report mount progress]' \
			- set1 \
			':filesystem:_zfs_dataset -t fs' \
			- set2 \
			'-a[Mount all available ZFS filesystems]'
		;;

	("unmount")
		_arguments -A "-*" \
			- set1 \
			'-f[Force unmount]' \
			':filesystem:_zfs_dataset -t fs -t mtpt' \
			- set2 \
			'-a[Unmount all ZFS filesystems]'
		;;

	("share")
		_arguments -A "-*" \
			- set1 \
			'-a[Share all available ZFS filesystems]' \
			- set2 \
			':filesystem:_zfs_dataset -t fs'
		;;

	("unshare")
		_arguments -A "-*" \
			- set1 \
			'-a[Unshare all shared ZFS filesystems]' \
			- set2 \
			'-F[Force unshare]' \
			':filesystem:_zfs_dataset -t fs -t mtpt'
		;;

	("send")
		_arguments -A "-*" \
			'-i[Generate an incremental stream]:snapshot:_zfs_dataset -t snap' \
			':snapshot:_zfs_dataset -t snap'
		;;

	("receive")
		_arguments -A "-*" \
			'-v[Verbose]' \
			'-n[Do not receive the stream]' \
			'-F[Force a rollback if necessary]' \
			- set1 \
			':filesystem/volume/snapshot:_zfs_dataset' \
			- set2 \
			'-d[Set path prefix]:filesystem:_zfs_dataset -t fs'
		;;

	("allow"|"unallow")
		_message "unimplemented zfs subcommand: $service"
		;;

	("upgrade")
		_arguments -A "-*" \
			- set1 \
			'-v[Verbose]' \
			- set2 \
			'-a[Upgrade all filesystems on all pools]' \
			'-r[Upgrade descendent filesystems, too]' \
			'-V[Upgrade to specified version]:version:(1 2)' \
			- set3 \
			'-r[Upgrade descendent filesystems, too]' \
			'-V[Upgrade to specified version]:version:(1 2)' \
			':filesystem:_zfs_dataset -t fs' \
		;;

	(*)
		_message "unknown zfs subcommand: $service"
		;;
	esac
}

_zfs "$@"