#!/usr/bin/env bash

set -x

TMPDIR=/tmp

createZdevice() {
	# Check Zram Class created
	if ! [[ -d "/sys/class/zram-control" ]]; then
		modprobe --verbose zram >> "$ZLOG" 2>&1 || return 1
		RAM_DEV='0'
	else
		RAM_DEV="$(cat /sys/class/zram-control/hot_add)"
	fi

	if [[ -n $MEM_SIZE ]]; then
		echo "$MEM_SIZE" > "/sys/block/zram${RAM_DEV}/mem_limit"
	else
		echo "zram${RAM_DEV} no mem_limit" >> "$ZLOG"
	fi

	if [[ -n $DISK_SIZE ]]; then
		echo "$DISK_SIZE" > "/sys/block/zram${RAM_DEV}/disksize"
	else
		echo "zram${RAM_DEV} no disksize" >> "$ZLOG"
		return 1
	fi

	if [[ -n $ALG ]]; then
		echo "$ALG" > "/sys/block/zram${RAM_DEV}/comp_algorithm"
	else
		echo "zram${RAM_DEV} no comp_algorithm" >> "$ZLOG"
	fi

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

createZswap() {
	createZdevice || return 1
	mkswap --label "zram-config${RAM_DEV}" "/dev/zram${RAM_DEV}" >> "$ZLOG" 2>&1 || return 1

	if [[ -n $PRIORITY ]]; then
		swapon -v -p "$PRIORITY" "/dev/zram${RAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	else
		echo "zram${RAM_DEV} no swap_priority" >> "$ZLOG"
		return 1
	fi

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

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

	echo "swap		/zram${RAM_DEV}		zram-config${RAM_DEV}" >> "${TMPDIR}"/zram-device-list
}

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

	if [[ -n $BIND_DIR ]]; then
		mkdir -p "${ZDIR}${BIND_DIR}" >> "$ZLOG" 2>&1 || return 1

		if [[ -n $TARGET_DIR ]]; then
			dirPerm="$(stat -c "%a" "$TARGET_DIR")"
			dirUser="$(stat -c "%u" "$TARGET_DIR")"
			dirGroup="$(stat -c "%g" "$TARGET_DIR")"

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

			mount --verbose --bind "${TARGET_DIR}/" "${ZDIR}${BIND_DIR}/" >> "$ZLOG" 2>&1 || return 1
			mount --verbose --make-private "${ZDIR}${BIND_DIR}/" >> "$ZLOG" 2>&1 || 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 "dirMountOpt: ${dirMountOpt}; dirFsType: ${dirFSType}" >> "$ZLOG"
			createZdevice || return 1

			mke2fs -v -t "$dirFSType" "/dev/zram${RAM_DEV}" >> "$ZLOG" 2>&1 || return 1
			mkdir -p "${ZDIR}/zram${RAM_DEV}" >> "$ZLOG" 2>&1 || return 1
			mount --verbose --types "$dirFSType" -o "$dirMountOpt" "/dev/zram${RAM_DEV}" "${ZDIR}/zram${RAM_DEV}/" >> "$ZLOG" 2>&1 || return 1
			mkdir -p "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" >> "$ZLOG" 2>&1 || 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" 2>&1 || return 1

			chown "${dirUser}:${dirGroup}" "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" >> "$ZLOG" 2>&1 || return 1
			chmod "$dirPerm" "${ZDIR}/zram${RAM_DEV}/upper" "${ZDIR}/zram${RAM_DEV}/workdir" "$TARGET_DIR" >> "$ZLOG" 2>&1 || return 1

			echo "${ZTYPE}		/zram${RAM_DEV}		${TARGET_DIR}		${BIND_DIR}" >> "${TMPDIR}"/zram-device-list
		else
			echo "No mount directory provided in ztab" >> "$ZLOG"
			return 1
		fi
	else
		echo "No bind directory provided in ztab" >> "$ZLOG"
		return 1
	fi
}

createZlog() {
	if dpkg -s 'nginx' &> /dev/null && [[ $(systemctl is-active nginx.service) == "active" ]]; then
		systemctl stop nginx.service >> "$ZLOG" 2>&1 || return 1
	fi
	if [[ $(systemctl is-active rsyslog.service) == "active" ]]; then
		systemctl stop rsyslog.service >> "$ZLOG" 2>&1 || return 1
	fi

	createZdir

	systemctl restart rsyslog.service >> "$ZLOG" 2>&1 || return 1
	if dpkg -s 'nginx' &> /dev/null; then
		systemctl restart nginx.service >> "$ZLOG" 2>&1 || return 1
	fi

	journalctl --flush >> "$ZLOG" 2>&1 || return 1

	if [[ -n $OLDLOG_DIR ]]; then
		echo -e "olddir ${OLDLOG_DIR}\\ncreateolddir 755 root root\\nrenamecopy" > /etc/logrotate.d/00_oldlog
	else
		echo "createZlog: no oldlog directory provided in ztab" >> "$ZLOG"
	fi
}

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

removeZlog() {
	if [[ -n $OLDLOG_DIR ]]; then
		rm -f /etc/logrotate.d/00_oldlog
	fi

	echo "$ZRAM_DEV" >> "$ZLOG"
	local DEV_NUM
	DEV_NUM="$(echo "$ZRAM_DEV" | tr -dc '0-9')"

	if [[ -n $TARGET_DIR ]]; then
		if dpkg -s 'nginx' &> /dev/null && [[ $(systemctl is-active nginx.service) == "active" ]]; then
			systemctl stop nginx.service >> "$ZLOG" 2>&1 || return 1
		fi
		if [[ $(systemctl is-active rsyslog.service) == "active" ]]; then
			systemctl stop rsyslog.service >> "$ZLOG" 2>&1 || return 1
		fi

		if ! (umount --verbose "${TARGET_DIR}/" >> "$ZLOG" 2>&1); then
			umount --verbose --lazy "${TARGET_DIR}/" >> "$ZLOG" 2>&1 || return 1
		fi
	else
		return 1
	fi

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

	if [[ -n $ZRAM_DEV ]]; then
		umount --verbose "${ZDIR}${ZRAM_DEV}/" >> "$ZLOG" 2>&1 || return 1
		umount "overlay${DEV_NUM}" >> "$ZLOG" 2>&1 || return 1
		rm -rfv "${ZDIR}${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	else
		return 1
	fi

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

	if dpkg -s 'nginx' &> /dev/null && [[ $(systemctl is-active nginx.service) == "active" ]]; then
		systemctl stop nginx.service >> "$ZLOG" 2>&1 || return 1
	fi
	if [[ $(systemctl is-active rsyslog.service) == "active" ]]; then
		systemctl stop rsyslog.service >> "$ZLOG" 2>&1 || return 1
	fi

	echo "$DEV_NUM" > /sys/class/zram-control/hot_remove
	zramctl -r "/dev${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1

	systemctl restart rsyslog.service >> "$ZLOG" 2>&1 || return 1
	if dpkg -s 'nginx' &> /dev/null; then
		systemctl restart nginx.service >> "$ZLOG" 2>&1 || return 1
	fi

	journalctl --flush >> "$ZLOG" 2>&1 || return 1
	echo "/dev$ZRAM_DEV removed" >> "$ZLOG"
}

removeZdir() {
	echo "$ZRAM_DEV" >> "$ZLOG"
	local DEV_NUM
	DEV_NUM="$(echo "$ZRAM_DEV" | tr -dc '0-9')"

	if [[ -n "$TARGET_DIR" ]]; then
		if ! (umount --verbose "${TARGET_DIR}/" >> "$ZLOG" 2>&1); then
			umount --verbose --lazy "${TARGET_DIR}/" >> "$ZLOG" 2>&1 || return 1
		fi
	else
		return 1
	fi

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

	if [[ -n $ZRAM_DEV ]]; then
		umount --verbose "${ZDIR}${ZRAM_DEV}/" >> "$ZLOG" 2>&1 || return 1
		umount "overlay${DEV_NUM}" >> "$ZLOG" 2>&1 || return 1
		rm -rfv "${ZDIR}${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	else
		return 1
	fi

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

	echo "$DEV_NUM" > /sys/class/zram-control/hot_remove
	zramctl -r "/dev${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	echo "/dev$ZRAM_DEV removed" >> "$ZLOG"
}

removeZswap() {
	local DEV_NUM
	DEV_NUM="$(echo "$ZRAM_DEV" | tr -dc '0-9')"

	swapoff "/dev${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	echo "$DEV_NUM" > /sys/class/zram-control/hot_remove
	zramctl -r "/dev${ZRAM_DEV}" >> "$ZLOG" 2>&1 || return 1
	echo "/dev$ZRAM_DEV removed" >> "$ZLOG"
}

enableZephemeral() {
	if [[ $ZEPHEMERAL == "false" ]]; then
		mkinitramfs -o /boot/initrd
		sudo sed -i '/root=PARTUUID/ s|$| init=/bin/ro-root.sh|' /boot/cmdline.txt
		echo "initramfs initrd followkernel" >> /boot/config.txt
		echo "ramfsfile=initrd" >> /boot/config.txt
		echo "ramfsaddr=-1" >> /boot/config.txt
		cp "${ZSHARE}/ro-root.sh" /bin/ro-root.sh
		chown root:root /bin/ro-root.sh
		chmod a+x /bin/ro-root.sh
		echo "Zephemeral enabled, after reboot all writes & changes will not be persistent"
	else
		echo "Zephemeral is already enabled"
	fi
}

disableZephemeral() {
	if [[ $ZEPHEMERAL == "true" ]]; then
		sed -i -e 's| init=/bin/ro-root.sh||g'  /boot/cmdline.txt
		sed -i -e '/initramfs initrd followkernel/d' /boot/config.txt
		sed -i -e '/ramfsfile=initrd/d' /boot/config.txt
		sed -i -e '/ramfsaddr=-1/d' /boot/config.txt
		rm -f /bin/ro-root.sh
		echo "Zephemeral disabled, after reboot system will be write and change enabled"
	else
		echo "Zephemeral is already disabled"
	fi
}

ZSHARE="/usr/local/share/zram-config"
ZLOG="${ZSHARE}/log/zram-config.log"
ZDIR="/opt/zram"
if [[ -f /boot/cmdline.txt ]] && grep -qs "init=/bin/ro-root.sh" /boot/cmdline.txt; then
	ZEPHEMERAL="true"
else
	ZEPHEMERAL="false"
fi

case "$1" in
	start)
		echo "zram-config start $(date +%Y-%m-%d-%H:%M:%S)" >> "$ZLOG"
		rm -fv "${TMPDIR}"/zram-device-list.new >> "$ZLOG"
		file="/etc/ztab"
		ZTAB_EMPTY="true"
		while read -r line; do
			case "$line" in
				"#"*)
					# Skip comment line
					continue
					;;

				"")
					# Skip empty line
					continue
					;;

				*)
					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 [[ -f "${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 "entry ${entry} already exists as zram."
							echo "$entry" >> "${TMPDIR}"/zram-device-list.new
							continue
						fi
					fi
					test -f "${TMPDIR}"/zram-device-list.new && mv "${TMPDIR}"/zram-device-list.new "${TMPDIR}"/zram-device-list

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

						dir)
							TARGET_DIR="$5"
							BIND_DIR="$6"
							if [[ $ZEPHEMERAL == "false" ]]; then
								createZdir
							fi
							;;

						log)
							TARGET_DIR="$5"
							BIND_DIR="$6"
							OLDLOG_DIR="$7"
							if [[ $ZEPHEMERAL == "false" ]]; then
								createZlog
							fi
							;;
					esac
					;;
			esac
		done < "$file"
		if [[ $ZTAB_EMPTY == "true" ]]; then
			echo "'/etc/ztab' is empty and needs to be configured" >> "$ZLOG"
			exit 1
		fi
		;;

	stop)
		echo "zram-config stop $(date +%Y-%m-%d-%H:%M:%S)" >> "$ZLOG"
		tac "${TMPDIR}"/zram-device-list > "${TMPDIR}"/zram-device-list.rev
		file="${TMPDIR}/zram-device-list.rev"
		while read -r line; do
			case "$line" in
				"#"*)
					# Skip comment line
					continue
					;;

				"")
					# Skip empty line
					continue
					;;

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

						dir)
							ZTYPE="$1"
							ZRAM_DEV="$2"
							TARGET_DIR="$3"
							BIND_DIR="$4"
							if [[ $ZEPHEMERAL == "false" ]]; then
								removeZdir
							fi
							;;

						log)
							ZTYPE="$1"
							ZRAM_DEV="$2"
							TARGET_DIR="$3"
							BIND_DIR="$4"
							if [[ $ZEPHEMERAL == "false" ]]; then
								removeZlog
							fi
							;;
					esac
					;;
			esac
		done < "$file"
		rm -fv "${TMPDIR}"/zram-device-list.rev "${TMPDIR}"/zram-device-list >> "$ZLOG"
		;;

	enable-ephemeral)
		enableZephemeral
		;;

	disable-ephemeral)
		disableZephemeral
		;;

	*)
		echo "Usage: zram-config {start|stop|enable-ephemeral|disable-ephemeral}" > 2
		exit 1
		;;
esac