#!/usr/bin/env bash

createZdevice() {
	local retval=0

	echo "createZdevice: Beginning creation of zDevice."
	if ! [[ -d "/sys/class/zram-control" ]]; then
		modprobe --verbose zram num_devices=0 &>> "$ZLOG" || return 1
	fi

	RAM_DEV="$(cat /sys/class/zram-control/hot_add)"

	if [[ -z $RAM_DEV ]]; then
		echo "createZdevice: Failed to find an open zram device, trying one more time!" >> "$ZLOG"
		RAM_DEV="$(cat /sys/class/zram-control/hot_add)"
		if [[ -z $RAM_DEV ]]; then
			echo "createZdevice: Failed to find an open zram device. Exiting!" | tee -a "$ZLOG"
			return 1
		fi
	fi

	if ! echo "$ALG" > "/sys/block/zram${RAM_DEV}/comp_algorithm"; then
		echo "createZdevice: zram${RAM_DEV}, Failed to set compression algorithm! Exiting!" | tee -a "$ZLOG"
		retval=1
	fi
	if ! echo "$DISK_SIZE" > "/sys/block/zram${RAM_DEV}/disksize"; then
		echo "createZdevice: zram${RAM_DEV}, Failed to set disk size! Exiting!" | tee -a "$ZLOG"
		retval=1
	fi
	if ! echo "$MEM_SIZE" > "/sys/block/zram${RAM_DEV}/mem_limit"; then
		echo "createZdevice: zram${RAM_DEV}, Failed to set memory limit! Exiting!" | tee -a "$ZLOG"
		retval=1
	fi
	if [[ $retval -ne 0 ]]; then
		echo "$RAM_DEV" > /sys/class/zram-control/hot_remove
		return 1
	fi

	if [[ $MEM_SIZE == 0 ]]; then
		echo "createZdevice: zram${RAM_DEV} no mem_limit" >> "$ZLOG"
	fi

	echo "createZdevice: zram${RAM_DEV} created comp_algorithm=${ALG} mem_limit=${MEM_SIZE} disksize=${DISK_SIZE}." >> "$ZLOG"
}

createZswap() {
	echo "createZswap: Beginning creation of zswap device." >> "$ZLOG"
	createZdevice || return 1
	mkswap --label "zram-config${RAM_DEV}" "/dev/zram${RAM_DEV}" &>> "$ZLOG" || return 1

	if [[ -n $PRIORITY ]]; then
		swapon -v -p "$PRIORITY" "/dev/zram${RAM_DEV}" &>> "$ZLOG" || return 1
	else
		echo "createZswap: No swap priority provided for zram${RAM_DEV}. Exiting!" | tee -a "$ZLOG"
		return 1
	fi

	if [[ -n $PAGE_CLUSTER ]]; then
		sysctl vm.page-cluster="$PAGE_CLUSTER" &>> "$ZLOG" || return 1
	else
		echo "createZswap: zram${RAM_DEV} no page_cluster" >> "$ZLOG"
	fi

	if [[ -n $SWAPPINESS ]]; then
		sysctl vm.swappiness="$SWAPPINESS" &>> "$ZLOG" || return 1
	else
		echo "createZswap: zram${RAM_DEV} no swappiness" >> "$ZLOG"
	fi

	echo "swap		/zram${RAM_DEV}		zram-config${RAM_DEV}" >> "$TMPDIR"/zram-device-list
	echo "createZswap: Completed zswap device creation." >> "$ZLOG"
}

createZdir() {
	local dirPerm
	local dirUser
	local dirGroup
	local dirMountOpt
	local dirFSType

	if [[ -z $BIND_DIR ]]; then
		echo "createZdir: No bind directory provided in '/etc/ztab'. Exiting!" | tee -a "$ZLOG"
		return 1
	elif [[ -z $TARGET_DIR ]]; then
		echo "createZdir: No mount directory provided in '/etc/ztab'. Exiting!" | tee -a "$ZLOG"
		return 1
	fi

	echo "createZdir: Beginning creation of ${ZDIR}/zram${RAM_DEV}." >> "$ZLOG"
	mkdir -p "${ZDIR}${BIND_DIR}" &>> "$ZLOG" || return 1

	dirPerm="$(stat -c "%a" "$TARGET_DIR")"
	dirUser="$(stat -c "%u" "$TARGET_DIR")"
	dirGroup="$(stat -c "%g" "$TARGET_DIR")"

	echo "createZdir: dirPerm - ${TARGET_DIR} ${dirPerm} ${dirUser}:${dirGroup}" >> "$ZLOG"

	mount --verbose --bind "${TARGET_DIR}/" "${ZDIR}${BIND_DIR}/" &>> "$ZLOG" || return 1
	mount --verbose --make-private "${ZDIR}${BIND_DIR}/" &>> "$ZLOG" || return 1

	dirMountOpt="$(awk -v a="${ZDIR}${BIND_DIR}" '$2 == a {print $4}' /proc/mounts | head -1)"
	dirFSType="$(awk -v a="${ZDIR}${BIND_DIR}" '$2 == a {print $3}' /proc/mounts | head -1)"

	echo "createZdir: dirMountOpt - ${dirMountOpt}; dirFsType: ${dirFSType}" >> "$ZLOG"
	createZdevice || return 1

	mke2fs -v -t "$dirFSType" "/dev/zram${RAM_DEV}" &>> "$ZLOG" || return 1
	mkdir -p "${ZDIR}/zram${RAM_DEV}" &>> "$ZLOG" || return 1
	mount --verbose --types "$dirFSType" -o "$dirMountOpt" "/dev/zram${RAM_DEV}" "${ZDIR}/zram${RAM_DEV}/" &>> "$ZLOG" || return 1
	mkdir -p "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" &>> "$ZLOG" || return 1
	mount --verbose --types overlay -o "redirect_dir=on,lowerdir=${ZDIR}${BIND_DIR},upperdir=${ZDIR}/zram${RAM_DEV}/upper,workdir=${ZDIR}/zram${RAM_DEV}/workdir" "overlay${RAM_DEV}" "$TARGET_DIR" &>> "$ZLOG" || return 1
	chown "${dirUser}:${dirGroup}" "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" &>> "$ZLOG" || return 1
	chmod "$dirPerm" "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" &>> "$ZLOG" || return 1

	echo "${ZTYPE}		/zram${RAM_DEV}		${TARGET_DIR}		${BIND_DIR}" >> "$TMPDIR"/zram-device-list

	if [[ $ZTYPE == "log" ]] && [[ -n $OLDLOG_DIR ]]; then
		echo -e "olddir ${OLDLOG_DIR}\\ncreateolddir 755 root root\\nrenamecopy" > /etc/logrotate.d/00_oldlog
	elif [[ $ZTYPE == "log" ]]; then
		echo "createZdir: No oldlog directory provided in '/etc/ztab', skipping oldlog configuration." >> "$ZLOG"
	fi
	echo "createZdir: Creation of ${ZDIR}/zram${RAM_DEV} complete." >> "$ZLOG"
}

mergeOverlay() {
	echo "mergeOverlay: Beginning merge of ${ZDIR}${BIND_DIR}." >> "$ZLOG"
	ls -la "$ZDIR" "${ZDIR}${BIND_DIR}" "${ZDIR}${ZRAM_DEV}/upper" >> "$ZLOG"
	cd /usr/local/lib/zram-config/ || return 1
	(./overlay merge --yes --lowerdir="${ZDIR}${BIND_DIR}" --upperdir="${ZDIR}${ZRAM_DEV}/upper") &>> "$ZLOG" || return 1
	bash -x -- *.sh &>> "$ZLOG"
	rm -fv -- *.sh &>> "$ZLOG" || return 1
	echo "mergeOverlay: Merge of ${ZDIR}${BIND_DIR} complete." >> "$ZLOG"
}

removeZdir() {
	local count=0

	[[ -n $OLDLOG_DIR ]] && rm -f /etc/logrotate.d/00_oldlog

	echo "removeZdir: Beginning removal of device /dev${ZRAM_DEV}." >> "$ZLOG"

	[[ -z $TARGET_DIR ]] && return 1
	if ! umount --verbose "${TARGET_DIR}/" &>> "$ZLOG"; then
		[[ -x $(command -v lsof) ]] && lsof "${TARGET_DIR}/" &>> "$ZLOG"
		umount --verbose --lazy "${TARGET_DIR}/" &>> "$ZLOG" || return 1
	fi

	mergeOverlay &>> "$ZLOG" || return 1

	[[ -z $ZRAM_DEV ]] && return 1
	if ! umount --verbose "${ZDIR}${ZRAM_DEV}/" &>> "$ZLOG"; then
		umount --verbose --lazy "${ZDIR}${ZRAM_DEV}/" &>> "$ZLOG" || return 1
	fi
	rm -rfv "${ZDIR}${ZRAM_DEV}" &>> "$ZLOG" || return 1

	[[ -z $BIND_DIR ]] && return 1
	if ! umount --verbose "${ZDIR}${BIND_DIR}/" &>> "$ZLOG"; then
		umount --verbose --lazy "${ZDIR}${BIND_DIR}/" &>> "$ZLOG" || return 1
	fi
	rm -rfv "${ZDIR}${BIND_DIR}" &>> "$ZLOG" || return 1

	if [[ -z $SERVICE ]]; then  # We don't care about device reset when shutting down system
		until echo "${ZRAM_DEV//[!0-9]/}" > /sys/class/zram-control/hot_remove || [[ count -ge 5 ]]; do
			count=$(( count + 1 ))
			sleep 5
		done
	fi

	echo "removeZdir: Device /dev$ZRAM_DEV removed." >> "$ZLOG"
}

removeZswap() {
	echo "removeZswap: Beginning removal of device /dev${ZRAM_DEV}." >> "$ZLOG"

	swapoff --verbose "/dev${ZRAM_DEV}" &>> "$ZLOG" || return 1
	echo "${ZRAM_DEV//[!0-9]/}" > /sys/class/zram-control/hot_remove || return 1

	echo "removeZswap: Device /dev$ZRAM_DEV removed." >> "$ZLOG"
}

syncZdir() {
	echo "syncZdir: Beginning sync of device /dev${ZRAM_DEV}." >> "$ZLOG"

	[[ -z $TARGET_DIR ]] && return 1
	if ! umount --verbose "${TARGET_DIR}/" &>> "$ZLOG"; then
		[[ -x $(command -v lsof) ]] && lsof "${TARGET_DIR}/" &>> "$ZLOG"
		umount --verbose --lazy "${TARGET_DIR}/" &>> "$ZLOG" || return 1
	fi

	mergeOverlay &>> "$ZLOG" || return 1

	mkdir -p "${ZDIR}${ZRAM_DEV}/upper" "${ZDIR}${ZRAM_DEV}/workdir" "$TARGET_DIR" &>> "$ZLOG" || return 1
	mount --verbose --types overlay -o "redirect_dir=on,lowerdir=${ZDIR}${BIND_DIR},upperdir=${ZDIR}${ZRAM_DEV}/upper,workdir=${ZDIR}${ZRAM_DEV}/workdir" "overlay${ZRAM_DEV//[!0-9]/}" "$TARGET_DIR" &>> "$ZLOG" || return 1

	echo "syncZdir: Device /dev$ZRAM_DEV synced." >> "$ZLOG"
}

serviceConfiguration() {
	if [[ $1 == "stop" ]]; then
		echo "serviceConfiguration: Stopping services that interfere with zram device configuration." >> "$ZLOG"
		if [[ $OS == "alpine" ]]; then
			if rc-service syslog status &> /dev/null; then
				export syslogActiveAlpine="true"
				rc-service syslog stop &>> "$ZLOG" || return 1
			fi
		else
			if [[ $(systemctl is-active rsyslog.service) == "active" ]]; then
				export rsyslogActive="true"
				systemctl --no-block stop syslog.socket &>> "$ZLOG" || return 1
			fi
			if [[ $(systemctl is-active systemd-journald.service) == "active" ]]; then
				export journaldActive="true"
				journalctl --flush &>> "$ZLOG" || return 1
				systemctl --no-block stop systemd-journald.socket systemd-journald-audit.socket systemd-journald-dev-log.socket &>> "$ZLOG" || return 1
			fi
		fi
	elif [[ $1 == "start" ]]; then
		echo "serviceConfiguration: Restarting services that interfere with zram device configuration." >> "$ZLOG"
		if [[ -n $journaldActive ]]; then
			systemctl --no-block restart systemd-journald.socket systemd-journald-audit.socket systemd-journald-dev-log.socket &>> "$ZLOG" || return 1
		fi
		if [[ -n $rsyslogActive ]]; then
			systemctl --no-block restart syslog.socket &>> "$ZLOG" || return 1
		fi
		if [[ -n $syslogActiveAlpine ]]; then
			rc-service syslog start &>> "$ZLOG" || return 1
		fi
	fi
}

TMPDIR="/tmp"
ZDIR="/opt/zram"
ZLOG="/usr/local/share/zram-config/log/zram-config.log"
OS="$(grep -o '^ID=.*$' /etc/os-release | cut -d'=' -f2)"
if [[ -s /run/systemd/shutdown/scheduled ]]; then
	SERVICE="1"
fi

case "$1" in
	start)
		echo "zram-config start $(printf "%(%F-%T-%Z)T\\n" "-1")" | tee -a "$ZLOG"
		ZTAB_EMPTY="true"
		while read -r line; do
			case "$line" in
				"#"*)
					# Skip comment line
					continue
					;;

				"")
					# Skip empty line
					continue
					;;

				*)
					# shellcheck disable=SC2086
					set -- $line
					echo "ztab create ${1} ${2} ${3} ${4} ${5} ${6} ${7} ${8} ${9}" >> "$ZLOG"
					ZTAB_EMPTY="false"
					ZTYPE="$1"
					ALG="$2"
					MEM_SIZE="$3"
					DISK_SIZE="$4"
					if [[ -s "$TMPDIR"/zram-device-list ]]; then
						if [[ $1 == "swap" ]]; then
							entry="$(grep "^swap" "$TMPDIR"/zram-device-list)"
						else
							entry="$(grep "${1}.*${5}" "$TMPDIR"/zram-device-list)"
						fi
						if [[ -n $entry ]]; then
							echo "Start: Entry ${entry} already exists as a zram device, skipping recreation of device." >> "$ZLOG"
							continue
						fi
					fi

					case "$1" in
						swap)
							PRIORITY="$5"
							PAGE_CLUSTER="$6"
							SWAPPINESS="$7"
							createZswap
							;;

						dir|log)
							TARGET_DIR="$5"
							BIND_DIR="$6"
							OLDLOG_DIR="$7"
							serviceConfiguration "stop"
							createZdir
							;;
					esac
					;;
			esac
		done < /etc/ztab
		if [[ $ZTAB_EMPTY == "true" ]]; then
			echo "Start: Configuration file '/etc/ztab' is empty and needs to be configured. Exiting!" | tee -a "$ZLOG"
			exit 1
		fi
		;;

	stop)
		echo "zram-config stop $(printf "%(%F-%T-%Z)T\\n" "-1")" | tee -a "$ZLOG"
		if ! [[ -s "$TMPDIR"/zram-device-list ]]; then
			echo "Stop: zram-config not running. Exiting!" | tee -a "$ZLOG"
			exit 0
		fi
		tac "$TMPDIR"/zram-device-list > "$TMPDIR"/zram-device-list.rev
		while read -r line; do
			case "$line" in
				"#"*)
					# Skip comment line
					continue
					;;

				"")
					# Skip empty line
					continue
					;;

				*)
					# shellcheck disable=SC2086
					set -- $line
					echo "ztab remove ${1} ${2} ${3} ${4}" >> "$ZLOG"
					case "$1" in
						swap)
							ZRAM_DEV="$2"
							removeZswap
							;;

						dir|log)
							ZRAM_DEV="$2"
							TARGET_DIR="$3"
							BIND_DIR="$4"
							[[ -z $SERVICE ]] && serviceConfiguration "stop"
							removeZdir
							;;
					esac
					;;
			esac
		done < "$TMPDIR"/zram-device-list.rev
		rm -fv "$TMPDIR"/zram-device-list.rev "$TMPDIR"/zram-device-list >> "$ZLOG"
		;;

	sync)
		echo "zram-config sync $(printf "%(%F-%T-%Z)T\\n" "-1")" | tee -a "$ZLOG"
		if ! [[ -s "$TMPDIR"/zram-device-list ]]; then
			echo "Sync: zram-config not running. Exiting!" | tee -a "$ZLOG"
			exit 0
		fi
		tac "$TMPDIR"/zram-device-list > "$TMPDIR"/zram-device-list.rev
		while read -r line; do
			case "$line" in
				"#"*)
					# Skip comment line
					continue
					;;

				"")
					# Skip empty line
					continue
					;;

				*)
					# shellcheck disable=SC2086
					set -- $line
					echo "ztab sync ${1} ${2} ${3} ${4}" >> "$ZLOG"
					case "$1" in
						dir|log)
							ZRAM_DEV="$2"
							TARGET_DIR="$3"
							BIND_DIR="$4"
							[[ -z $SERVICE ]] && serviceConfiguration "stop"
							syncZdir
							;;
					esac
					;;
			esac
		done < "$TMPDIR"/zram-device-list.rev
		rm -fv "$TMPDIR"/zram-device-list.rev >> "$ZLOG"
		;;

	*)
		echo "Usage: zram-config {start|stop|sync}"
		exit 0
		;;
esac
[[ -z $SERVICE ]] && serviceConfiguration "start"
