#!/bin/bash
{
	#////////////////////////////////////
	# DietPi Function:
	#
	#////////////////////////////////////
	# Created by Daniel Knight / daniel.knight@dietpi.com / dietpi.com
	#
	#////////////////////////////////////
	#
	# Info:
	# - Sets up user data directory.
	# - Allows automated moving of user data in DietPi. Automatically generates a symlink from /mnt/dietpi_userdata to target directory if needed.
	# - Also moves the Dphys swapfile to $TARGET_DIRECTORY. Required to prevent locked .swapfile from allow us to delete $SOURCE_DIRECTORY.
	#
	# Usage:
	# - /DietPi/dietpi/func/dietpi-set_userdata SOURCE_DIRECTORY TARGET_DIRECTORY	| Setup user data directory, move data if needed. if TARGET_DIRECTORY='auto' , auto target location.Returns 1 if failed.
	#////////////////////////////////////

	#Import DietPi-Globals ---------------------------------------------------------------
	. /DietPi/dietpi/func/dietpi-globals
	G_CHECK_ROOT_USER
	export G_PROGRAM_NAME='DietPi-Set_userdata'
	#Import DietPi-Globals ---------------------------------------------------------------

	EXIT_CODE=0

	SOURCE_DIRECTORY="$1"
	TARGET_DIRECTORY="$2"

	LOGFILE_OUTPUT_TEXT=''
	FP_LOGFILE='/var/log/dietpi-move_userdata.log'

	SWAPFILE_SIZE=$(/DietPi/dietpi/func/dietpi-set_dphys-swapfile | awk '{print $1}')

	FREESPACE_AVAILABLE_TARGET=0
	FREESPACE_REQUIRED_SOURCE=0

	RUN_MOVE_DATA=1
	Run_Move_Data(){

		# - stop all running services.
		/DietPi/dietpi/dietpi-services stop

		# - move swap out the way
		/DietPi/dietpi/func/dietpi-set_dphys-swapfile 0 /mnt/.swapfile

		# - Remove directory if its currently a symlink.
		if [ -L "$TARGET_DIRECTORY" ]; then

			rm -R "$TARGET_DIRECTORY"

		fi

		mkdir -p "$TARGET_DIRECTORY"

		# - Copy source to target, if it contains any files/folders
		if [ -z "$(find $SOURCE_DIRECTORY -maxdepth 0 -empty)" ]; then

			G_DIETPI-NOTIFY 0 "Moving your existing data from $SOURCE_DIRECTORY to $TARGET_DIRECTORY"
			G_DIETPI-NOTIFY 2 "Please wait...\n"

			sleep 1

			# - Check if user permissions are supported on target filesystem
			local fp_test_target="$TARGET_DIRECTORY/.permissions_test"

			rm "$fp_test_target" &> /dev/null

			echo 0 > "$fp_test_target"
			chown www-data:www-data "$fp_test_target" &> /dev/null

			local cp_options='-vR'

			if (( $(ls -lha "$fp_test_target" | grep -ci -m1 "www-data") )); then

				cp_options+='p'
				G_DIETPI-NOTIFY 2 "Target filesystem supports user permissions:"
				G_DIETPI-NOTIFY 2 " - Attempting to preserve permissions during transfer."
				sleep 1

			fi

			rm "$fp_test_target" &> /dev/null


			# - Begin transfer
			cp "$cp_options" "$SOURCE_DIRECTORY"/* "$TARGET_DIRECTORY"/

			# - Remove all files in source
			if (( $? == 0 )); then

				rm -R "$SOURCE_DIRECTORY"/*

			else

				LOGFILE_OUTPUT_TEXT="ERROR: Failed to copy $SOURCE_DIRECTORY/* to $TARGET_DIRECTORY."
				EXIT_CODE=1
				break

			fi

		fi

		# - move swapfile
		/DietPi/dietpi/func/dietpi-set_dphys-swapfile "$SWAPFILE_SIZE" "$TARGET_DIRECTORY"/.swapfile

		# - Remove source base folder if its a symlink
		if [ -L "$SOURCE_DIRECTORY" ]; then

			rm -R "$SOURCE_DIRECTORY"

		fi

		# - Create symlink to G_FP_DIETPI_USERDATA if required
		if [ "$TARGET_DIRECTORY" != "$G_FP_DIETPI_USERDATA" ]; then

			rm -R "$G_FP_DIETPI_USERDATA" &> /dev/null
			ln -sf "$TARGET_DIRECTORY" "$G_FP_DIETPI_USERDATA"

		fi

		# - Set permissions for userdata dirs:
		/DietPi/dietpi/dietpi-software setpermissions

		# - Start services back up again
		/DietPi/dietpi/dietpi-services start

	}

	#/////////////////////////////////////////////////////////////////////////////////////
	# Main Loop
	#/////////////////////////////////////////////////////////////////////////////////////
	#init

	# - If there is no directory or symlink for $G_FP_DIETPI_USERDATA, always create a directory.
	if [ ! -d "$G_FP_DIETPI_USERDATA" ] && [ ! -L "$G_FP_DIETPI_USERDATA" ]; then

		#Run full core_env set
		/DietPi/dietpi/func/dietpi-set_core_environment

	fi

	#-------------------------------------------------------------------------------------
	#Run
	target_to_lowercase=$(echo -e "$TARGET_DIRECTORY" | tr '[:upper:]' '[:lower:]')

	G_DIETPI-NOTIFY 3 DietPi Updating user data location
	G_DIETPI-NOTIFY 2 " - From : $SOURCE_DIRECTORY"
	G_DIETPI-NOTIFY 2 " - To   : $TARGET_DIRECTORY"
	G_DIETPI-NOTIFY 2 "Please wait..."

	while (( $RUN_MOVE_DATA ));
	do

		# Sanity checks
		# - Check for both inputs
		if [ -z "$SOURCE_DIRECTORY" ] || [ -z "$TARGET_DIRECTORY" ]; then

			LOGFILE_OUTPUT_TEXT="ERROR: Please provide a source ($SOURCE_DIRECTORY) and target ($TARGET_DIRECTORY) directory for input."
			EXIT_CODE=1
			break

		# - Check if symlink is already pointing to target directory.
		elif [ "$(readlink -f $G_FP_DIETPI_USERDATA)" = "$TARGET_DIRECTORY" ]; then

			LOGFILE_OUTPUT_TEXT="$G_FP_DIETPI_USERDATA is already symlinked to target directory."
			EXIT_CODE=0 #return ok
			break

		# - Check if source directory exists
		elif [ ! -d "$SOURCE_DIRECTORY" ]; then

			LOGFILE_OUTPUT_TEXT="ERROR: source directory $SOURCE_DIRECTORY does not exist."
			EXIT_CODE=1
			break

		# - Check for disallowed directory match
		elif (( $(echo -e "$SOURCE_DIRECTORY" | grep -ci -m1 "$TARGET_DIRECTORY") ||
			$(echo -e "$TARGET_DIRECTORY" | grep -ci -m1 "$SOURCE_DIRECTORY") )); then

			LOGFILE_OUTPUT_TEXT="ERROR: $SOURCE_DIRECTORY and $TARGET_DIRECTORY cannot be within each other. Disallowed directory match."
			EXIT_CODE=1
			break

		# - Only allow full filepaths
		elif [ "${SOURCE_DIRECTORY:0:1}" != "/" ] || [ "${TARGET_DIRECTORY:0:1}" != "/" ]; then

			LOGFILE_OUTPUT_TEXT="ERROR: Both source ($SOURCE_DIRECTORY) and target directories ($TARGET_DIRECTORY) must contain the full filepath (eg: /mnt/drive1)"
			EXIT_CODE=1
			break

		fi

		#Ensure we can create and write to target directory
		mkdir -p "$TARGET_DIRECTORY" &> /dev/null
		if [ ! -d "$TARGET_DIRECTORY" ]; then

			LOGFILE_OUTPUT_TEXT="ERROR: Unable to create target directory $TARGET_DIRECTORY."
			EXIT_CODE=1
			break

		# - Create a test file inside target directory
		else

			test_file_name=".testfile_dietpi_userdata"

			echo 0 > "$TARGET_DIRECTORY"/"$test_file_name"

			if [ -f "$TARGET_DIRECTORY"/"$test_file_name" ]; then
				rm "$TARGET_DIRECTORY"/"$test_file_name"
			else
				LOGFILE_OUTPUT_TEXT="ERROR: Unable to create test file in target directory $TARGET_DIRECTORY. Check permissions."
				EXIT_CODE=1
				break
			fi

		fi

		#Ensure enough freespace in target
		FREESPACE_AVAILABLE_TARGET=$(( $(df -Pk "$TARGET_DIRECTORY" | awk '{print $4}' | sed -n 2p) * 1024 )) #bytes

		G_DIETPI-NOTIFY 2 "Calculating space required for moving data, please wait..."
		FREESPACE_REQUIRED_SOURCE=$(du -cbs "$SOURCE_DIRECTORY" | awk '{print $1}' | sed -n 1p) #bytes

		echo -e " - Available $FREESPACE_AVAILABLE_TARGET bytes"
		echo -e " - Required  $FREESPACE_REQUIRED_SOURCE bytes"

		if (( $FREESPACE_AVAILABLE_TARGET < $FREESPACE_REQUIRED_SOURCE )); then

			LOGFILE_OUTPUT_TEXT="ERROR: Not enough free space in target directory $TARGET_DIRECTORY.\n - Available $FREESPACE_AVAILABLE_TARGET\n - Required $FREESPACE_REQUIRED_SOURCE"
			EXIT_CODE=1
			break

		fi

		#Run, attempt to move data.
		Run_Move_Data

		#Done
		RUN_MOVE_DATA=0

	done

	#-----------------------------------------------------------------------------------
	#Print any errors and send to logfile
	rm "$FP_LOGFILE" &> /dev/null
	if [ -n "$LOGFILE_OUTPUT_TEXT" ]; then


		# - Info
		if (( $EXIT_CODE == 0 )); then

			G_DIETPI-NOTIFY 2 "$LOGFILE_OUTPUT_TEXT"

		# - Error
		else

			G_DIETPI-NOTIFY 1 "$LOGFILE_OUTPUT_TEXT"

		fi

		echo -e ""

		# + send to logfile
		echo -e "$LOGFILE_OUTPUT_TEXT" > "$FP_LOGFILE"

	else

		G_DIETPI-NOTIFY 2 User data location setup completed.

		echo -e ""

	fi

	#-----------------------------------------------------------------------------------
	G_DIETPI-NOTIFY -1 ${EXIT_CODE:=0}
	#-----------------------------------------------------------------------------------
	exit $EXIT_CODE
	#-----------------------------------------------------------------------------------
}