#!/bin/bash
do {
	#////////////////////////////////////
	# 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 ---------------------------------------------------------------
	source /DietPi/dietpi/func/dietpi-globals
	G_CHECK_ROOT_USER
	export G_PROGRAM_NAME='DietPi-Set_userdata'
	#Import DietPi-Globals ---------------------------------------------------------------

	setvar EXIT_CODE = '0'

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

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

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

	setvar FREESPACE_AVAILABLE_TARGET = '0'
	setvar FREESPACE_REQUIRED_SOURCE = '0'

	setvar RUN_MOVE_DATA = '1'
	proc 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 test -L $TARGET_DIRECTORY {

			rm -R $TARGET_DIRECTORY

		}

		mkdir -p $TARGET_DIRECTORY

		# - Copy source to target, if it contains any files/folders
		if test -z $(find $SOURCE_DIRECTORY -maxdepth 0 -empty) {

			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") )) {

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

			}

			rm $fp_test_target &> /dev/null


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

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

				rm -R "$SOURCE_DIRECTORY"/*

			} else {

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

			}

		}

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

		# - Remove source base folder if its a symlink
		if test -L $SOURCE_DIRECTORY {

			rm -R $SOURCE_DIRECTORY

		}

		# - Create symlink to G_FP_DIETPI_USERDATA if required
		if test $TARGET_DIRECTORY != $G_FP_DIETPI_USERDATA {

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

		}

		# - 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 test ! -d $G_FP_DIETPI_USERDATA && test ! -L $G_FP_DIETPI_USERDATA {

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

	}

	#-------------------------------------------------------------------------------------
	#Run
	setvar 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...
	{

		# Sanity checks
		# - Check for both inputs
		if test -z $SOURCE_DIRECTORY || test -z $TARGET_DIRECTORY {

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

		# - Check if symlink is already pointing to target directory.
		} elif test $(readlink -f $G_FP_DIETPI_USERDATA) = $TARGET_DIRECTORY {

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

		# - Check if source directory exists
		} elif test ! -d $SOURCE_DIRECTORY {

			setvar LOGFILE_OUTPUT_TEXT = ""ERROR: source directory $SOURCE_DIRECTORY does not exist.""
			setvar 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") )) {

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

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

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

		}

		#Ensure we can create and write to target directory
		mkdir -p $TARGET_DIRECTORY &> /dev/null
		if test ! -d $TARGET_DIRECTORY {

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

		# - Create a test file inside target directory
		} else {

			setvar test_file_name = "".testfile_dietpi_userdata""

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

			if test -f "$TARGET_DIRECTORY"/"$test_file_name" {
				rm "$TARGET_DIRECTORY"/"$test_file_name"
			} else {
				setvar LOGFILE_OUTPUT_TEXT = ""ERROR: Unable to create test file in target directory $TARGET_DIRECTORY. Check permissions.""
				setvar EXIT_CODE = '1'
				break
			}

		}

		#Ensure enough freespace in target
		setvar 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..."
		setvar 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 )) {

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

		}

		#Run, attempt to move data.
		Run_Move_Data

		#Done
		setvar RUN_MOVE_DATA = '0'

	}

	#-----------------------------------------------------------------------------------
	#Print any errors and send to logfile
	rm $FP_LOGFILE &> /dev/null
	if test -n $LOGFILE_OUTPUT_TEXT {


		# - Info
		if (( $EXIT_CODE == 0 )) {

			G_DIETPI-NOTIFY 2 $LOGFILE_OUTPUT_TEXT

		# - Error
		} else {

			G_DIETPI-NOTIFY 1 $LOGFILE_OUTPUT_TEXT

		}

		echo -e ""

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

	} else {

		G_DIETPI-NOTIFY 2 User data location setup completed.

		echo -e ""

	}

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