#!/bin/zsh
setopt no_unset warn_create_global extended_glob

load(){
	typeset -gA dev_child dev_name dev_alias dev_type dev_fs dev_mtp dev_size dev_label dev_discard
	local line kname pkname
	local -a record
	lsblk -rnpo kname,pkname,name,type,fstype,mountpoint,size,label,disc-gran | while IFS= read line; do
		record=( "${(@s. .)line}" )  # split manually so empty fields get assigned properly
		kname=$record[1]
		pkname=$record[2]

		dev_alias[$record[3]]=$kname
		dev_name[$kname]=$record[3]
		dev_type[$kname]=$record[4]
		dev_fs[$kname]=$record[5]
		dev_mtp[$kname]=$record[6]
		dev_size[$kname]=$record[7]
		dev_label[$kname]=$record[8]
		dev_discard[$kname]=$record[9]

		if (( $+dev_child[$pkname] )); then
			dev_child[$pkname]+=$'\0'$kname
		else
			dev_child[$pkname]=$kname
		fi
	done
}

order_aux(){
	local dev child
	for dev in "$@"; do
		if ! (( $+dev_child[$dev] )); then
			erase_tsort_list+=( "$dev -" )
			continue
		fi
		order_aux "${(0@)dev_child[$dev]}"
		for child in "${(0@)dev_child[$dev]}"; do
			erase_tsort_list+=( "$child $dev" )

			# echo "${dev_name[$child]#/dev/}  ->  ${dev_name[$dev]#/dev/}"
		done
	done
}

order(){
	local dev
	local -a devices
	# Resolve alias names to canonical kernel device names
	for dev in "$@"; do
		if (( $+dev_alias[$dev] )); then
			devices+=( "$dev_alias[$dev]" )
		elif (( $+dev_name[$dev] )); then
			devices+=( "$dev" )
		else
			dev=$dev:A  # resolve symlinks
			if (( $+dev_alias[$dev] )); then
				devices+=( "$dev_alias[$dev]" )
			elif (( $+dev_name[$dev] )); then
				devices+=( "$dev" )
			else
				printf >&2 "Unrecognised device: %s\n" "${(qqq)dev}"
				exit 2
			fi
		fi
	done
	typeset -ga erase_tsort_list
	order_aux "$devices[@]"

	typeset -ga erase_order
	erase_order=( "${(@f)"$(tsort <<<"${(F)erase_tsort_list}")"}" ) || exit $?
	unset erase_tsort_list

	# remove dummy "-" used when ordering
	while (( $erase_order[(I)-] )); do
		unset "erase_order[${erase_order[(I)-]}]"
	done
}

print-dev(){
	printf "%s: %s\n" "$1" "$2 (${dev_name[$2]#/dev/})  ${dev_size[$2]:---}  ${dev_label[$2]:---}  $dev_type[$2]  ${dev_fs[$2]:---}  ${dev_mtp[$2]:---}"
}

print-order() {
	local dev
	for dev in $erase_order; do
		if (( $pretend )); then
			print-dev "Would erase" "$dev"
		else
			print-dev "Will erase" "$dev"
		fi
		# blkid $dev
	done
}

erase_all() {
	local dev
	for dev in $erase_order; do
		print-dev "Erasing" "$dev"
		erase "$dev"
	done
}

typeset -g pretend
if [[ $1 == "-f" ]]; then
	shift
	+canfail() {
		echo -E '#>' "${(q)@}"
		"$@"
	}
	pretend=0
else
	+canfail() {
		echo -E '#>' "${(q)@}"
	}
	pretend=1
fi
+() { +canfail "$@" || exit $? }

try-discard() {
	[[ $dev_discard[$1] == 0B ]] && return
	local -a discard_args
	discard_args=( -l $((16*1024)) )  # default for devices larger than 2GiB
	case $dev_size[$1] in
		(*[BkK]) discard_args=( );;  # discard whole without progress
		(*M)
			if [[ ${dev_size[$1]%M} -le 2048 ]]; then
				discard_args=( -v )  # discard whole with progress
			fi
			;;
		(*G)
			if [[ ${dev_size[$1]%G} -le 2 ]]; then
				discard_args=( -v )  # discard whole with progress
			fi
			;;
	esac
	+canfail blkdiscard $discard_args "$1"
}

cleanup-vg-symlinks() {
	[[ -n $1 ]] || exit 99  # sanity check
	[[ -d /dev/$1 ]] || return
	local -a broken_symlinks empty_dirs

	broken_symlinks=( /dev/$1/*(N-@) )
	if (( $#broken_symlinks )); then
		+ rm $broken_symlinks
	fi

	empty_dirs=( /dev/$1(N/^F) )
	if (( $#empty_dirs )); then
		+ rmdir $empty_dirs
	fi
}

erase() {
	if [[ -n $dev_mtp[$1] ]]; then
		+ umount -R "$dev_mtp[$1]"
	fi

	case $dev_fs[$1] in
		(swap)
			+canfail swapoff "$1"
			;;
		(LVM2_member)
			local vg
			local -a vg_pvs
			pvs --unbuffered --noheadings -o vg_name | while read vg; do
				[[ -z $vg ]] && continue
				vg_pvs=( "${(@f)"$(vgs --unbuffered --noheadings -o pv_name "$vg")"}" )
				# strip leading and trailing whitespace from each line
				vg_pvs=( "${(@)vg_pvs//# ##/}" )
				vg_pvs=( "${(@)vg_pvs//% ##}" )
				if [[ 1 == $#vg_pvs && "$vg_pvs[1]" == "$dev" ]]; then
					# Last physical volume in volume group
					+ vgchange -a n "$vg"
					+ vgremove "$vg"
					cleanup-vg-symlinks "$vg"
				else
					+ vgreduce "$vg" "$1"
				fi
			done
			+ pvremove "$1"
			;;
	esac

	+ wipefs -af "$1"
	try-discard "$1"

	case $dev_type[$1] in
		(lvm)
			+ lvchange -a n "$dev_name[$1]"
			+ lvremove "$dev_name[$1]"
			;;
		(raid*)
			+ mdadm --stop "$dev_name[$1]"
			;;
	esac

}
#typeset -ft erase

() {
	load
	order "$@"
	print-order

	if (( $pretend )); then
		echo -n "Proceed with simulation? "
	else
		echo -n "Erase above devices? "
	fi
	local REPLY
	read -q || exit 1
	unset REPLY

	printf "\n\n3... " || exit $?
	sleep 1 || exit $?

	printf "2... " || exit $?
	sleep 1 || exit $?

	printf "1..." || exit $?
	sleep 1 || exit $?

	printf "\n\n" || exit $?

	erase_all

	if (( $pretend )); then
		echo "Simulation finished"
	else
		blkid "$@"
		true
	fi
} "$@"