#!/bin/bash

# ZSC-Sync - zfs snapshot cron sync
# Usage: <script> <remote> <remote_dataset>

ZSC_CONFIG="$HOME/.zsc"

REMOTE=$1
if test -z "$REMOTE"; then
	echo "No remote specified"
	exit -1
fi

REMOTE_ROOT_DATASET=$2
if test -z "$REMOTE_ROOT_DATASET"; then
	echo "No remote root dataset specified"
	exit -1
fi


### LOCK ###
echo "Locking $0"
if ! lockfile-create -p -r 2 "$0"; then
	echo "Another backup is running, aborting"
	exit 49
fi


update_remote_state() {
	echo "Fetching remote snapshots..."
	remote_snapshots=$(ssh $REMOTE zfs list -H -t snapshot -o name -s creation)

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

	echo "Fetching subdatasets..."
	remote_subdatasets=$(ssh $REMOTE zfs list -H -o name -s creation -r "$REMOTE_ROOT_DATASET")

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

	echo "Fetching resume tokens..."
	resume_tokens=$(ssh $REMOTE zfs get -H -p receive_resume_token -o name,value -t filesystem)

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

	return 0
}


### Resume existing transfer ###
# Returns 0 if no action was performed or transfer is successful, 1 on error
resume_transfer() {
	remote_dataset=$1

	RESUME_TOKEN=$(echo -e "$resume_tokens" | grep -P "^$remote_dataset\t.*\$" | cut -d$'\t' -f2)
	if [ "$RESUME_TOKEN" != "" ] && [ "$RESUME_TOKEN" != "-" ]; then
		echo "Resuming interrupted transfer..."
		zfs send -v -t "$RESUME_TOKEN" | ssh $REMOTE zfs receive -Fu -s "$remote_dataset" || return 10
		echo "Resumed transfer finished."
		update_remote_state || return 20
		return 0
	fi

	return 0
}


handle_dataset() {
	DATASET=$1
	RECURSIVE=$2

	PREFIX=$3
	if test -z "$PREFIX"; then
		PREFIX="zsc-"
	fi

	if test "!no_prefix!" = "$PREFIX"; then
		PREFIX=""
	fi

	TARGET_REMOTE_DATASET=$4
	if test -z "$TARGET_REMOTE_DATASET"; then
		TARGET_REMOTE_DATASET="$REMOTE_ROOT_DATASET/"$(echo -e "$DATASET" | sed -r "s/\//_/g")
	fi

	echo "[$DATASET -> $TARGET_REMOTE_DATASET]"
	echo


	resume_transfer "$TARGET_REMOTE_DATASET" || return 11


	### Sync existing snapshots ###
	REMOTE_SNAPSHOTS=($(echo -e "$remote_snapshots" | grep "$TARGET_REMOTE_DATASET@" | cut -d"@" -f 2))
	LOCAL_SNAPSHOTS=($(zfs list -H -t snapshot -o name -s creation "$DATASET" | cut -d"@" -f 2))

	echo "Remote snapshots: ${#REMOTE_SNAPSHOTS[@]}"
	echo "Local snapshots: ${#LOCAL_SNAPSHOTS[@]}"
	echo

	for local_snapshot in "${LOCAL_SNAPSHOTS[@]}"; do
		found=false
		for remote_snapshot in "${REMOTE_SNAPSHOTS[@]}"; do
			if [ "$remote_snapshot" = "$local_snapshot" ] ; then
				found=true
				break
			fi
		done

		if [ $found = true ] ; then
			echo "> $local_snapshot found remotely"
			FROM_SNAPSHOT=$local_snapshot
			TO_SNAPSHOT=""
		else
			echo "> $local_snapshot needs sync"
			TO_SNAPSHOT=$local_snapshot
		fi
	done

	echo "Snapshots are synchronized up to $FROM_SNAPSHOT"
	echo

	if test ! -z "$TO_SNAPSHOT"; then
		echo "Synchronizing snapshots from $FROM_SNAPSHOT to $TO_SNAPSHOT ..."

		if test ! -z "$FROM_SNAPSHOT"; then
			incr="-I $FROM_SNAPSHOT"
		fi
		if [ $RECURSIVE = true ] ; then
			recu="-R"
		fi
		zfs send -v -p $recu $incr "$DATASET@$TO_SNAPSHOT" | ssh $REMOTE zfs receive -Fu -s "$TARGET_REMOTE_DATASET" || return 12

		echo
		echo "Done"
		echo
	else
		echo "No snapshot to sync"
	fi

	return 0
}

init_subdataset() {
	DATASET=$1
	PARENT_DATASET=$2
	LATEST_SYNCHRONIZED_SNAPSHOT=$3

	TARGET_REMOTE_DATASET="$REMOTE_ROOT_DATASET/"$(echo -e $PARENT_DATASET | sed -r "s/\//_/g")"/"${DATASET:$(echo -e "$PARENT_DATASET" | wc -c)}


	resume_transfer "$TARGET_REMOTE_DATASET" || return 13

	if [ $(echo -e "$remote_subdatasets" | grep -e "^$TARGET_REMOTE_DATASET\$" | wc -l) -eq 0 ]; then
		echo "Subdataset not found on $TARGET_REMOTE_DATASET, initializing..."

		LOCAL_SNAPSHOT=($(zfs list -H -t snapshot -o name -s creation "$DATASET" | cut -d"@" -f 2 | tail -n 1))

		if test -z "$LOCAL_SNAPSHOT"; then
			echo "No local snapshot"
			return 14
		fi

		echo "[$DATASET@$LOCAL_SNAPSHOT -> $TARGET_REMOTE_DATASET]"
		zfs send -v -p "$DATASET@$LOCAL_SNAPSHOT" | ssh $REMOTE zfs receive -Fu -s "$TARGET_REMOTE_DATASET" || return 15

		echo
		echo "--- Done"
		echo
	elif [ $(echo -e "$remote_snapshots" | grep -e "^$TARGET_REMOTE_DATASET@$LATEST_SYNCHRONIZED_SNAPSHOT\$" | wc -l) -eq 0 ]; then
		echo
		echo "--- Subdataset $TARGET_REMOTE_DATASET is out of date compared to parent dataset. Syncing first..."

		handle_dataset "$DATASET" false "" "$TARGET_REMOTE_DATASET" || return 16

		echo
		echo "--- Done"
		echo
	fi

	return 0
}


update_remote_state || exit 21

HAD_ERROR=0

for line in $(cat "$ZSC_CONFIG") ; do
	IFS=: read dataset recursive d w m y prefix <<< "$line"
	echo
	echo
	echo "##### Dataset: $dataset Recursive: $recursive Prefix: $prefix d: $d w: $w m: $m y: $y #####"
	echo

	if [ $recursive = "true" ]; then
		SUBDATASETS=($(zfs list -o name | grep "$dataset/"))


		# Find latest synchronized parente snapshot
		REMOTE_DATASET="$REMOTE_ROOT_DATASET/"$(echo -e "$dataset" | sed -r "s/\//_/g")
		REMOTE_SNAPSHOTS=($(echo -e "$remote_snapshots" | grep "$REMOTE_DATASET@" | cut -d"@" -f 2))
		LOCAL_SNAPSHOTS=($(zfs list -H -t snapshot -o name -s creation "$dataset" | cut -d"@" -f 2))

		for local_snapshot in "${LOCAL_SNAPSHOTS[@]}"; do
			found=false
			for remote_snapshot in "${REMOTE_SNAPSHOTS[@]}"; do
				if [ "$remote_snapshot" = "$local_snapshot" ] ; then
					found=true
					break
				fi
			done

			if [ $found = true ] ; then
				latest_synchronized_snapshot=$local_snapshot
			fi
		done
		echo "Latest synchronized snapshot locally: $latest_synchronized_snapshot"
		echo


		# Init subdatasets
		echo "Subdatasets: ${SUBDATASETS[@]}"
		echo
		for p in "${SUBDATASETS[@]}"; do
			init_subdataset $p $dataset "$latest_synchronized_snapshot"
			r=$?
			if [ $r -ne 0 ]; then
				echo "Error while initing subdataset $p, continuing"
				if [ $HAD_ERROR ]; then
					HAD_ERROR=$r
				fi
			fi
		done
		echo
		echo "----------"
		echo

		update_remote_state || exit 22
	fi

	handle_dataset $dataset $recursive $prefix
	r=$?
	if [ $r -ne 0 ]; then
		echo "Error while initing subdataset $p, continuing"
		if [ $HAD_ERROR ]; then
			HAD_ERROR=$r
		fi
	fi

	echo
	echo "----------------------------------------------------------------------"
	echo
done


### UNLOCK ###
echo "Unlocking $0"
lockfile-remove "$0"

exit $HAD_ERROR