#!/bin/bash

DATE=$(date +%Y%m%d%H%M%S)

function Usage(){
    echo "Usage: $0 <option> [option]"
    echo "       $0 -commit snapshot1"
    echo "       $0 -rollback last"
    echo "	     $0 -deepfreeze"
	echo "       $0 -l"
	echo "       $0 -c"
	echo "       $0 -d"
}

function Exit() {
	if [ -d /media/root-rw ]; then
		umount /media/root-rw
	fi
	exit 0
}

if [[ $EUID -ne 0 ]]; then
    echo "错误：当前用户不是 root，请使用 sudo 或切换到 root 用户运行。"
    exit 1
fi

mkdir -p "/home/.config" ||
	exit 1
while read spec file vfstype opts pass freq; do
	if [ "$file" = "/data" ]; then
		case "$spec" in
			LABEL=* | \#LABEL=*)
				kylinupoverlay="/dev/disk/by-label/${spec#*=}"
				;;
			PARTLABEL=* | \#PARTLABEL=*)
				kylinupoverlay="/dev/disk/by-label/${spec#*=}"
				;;
			UUID=* | \#UUID=*)
				kylinupoverlay="/dev/disk/by-uuid/${spec#*=}"
				;;
			/dev/* | \#/dev/*)
				kylinupoverlay="${spec#\#}"
				;;
			*)
				kylinupoverlay="/dev/${spec}"
				;;
		esac
	fi
done < /etc/fstab
if [ -d /media/root-rw ]; then
	mount $kylinupoverlay /media/root-rw
else
	mkdir -p /media/root-rw
	mount $kylinupoverlay /media/root-rw
fi
option=${1#*-}
cmdline=$(cat /proc/cmdline)
cmdline=${cmdline#*/}
path=/boot/efi/EFI/BOOT/grub.cfg
case $option in
	y)
		if [ -f "/boot/extlinux/extlinux.conf" ]; then
			line=$(grep "append" /boot/extlinux/extlinux.conf)
			if [[ ! $line =~ "overlay" ]]; then
				sed -i "s/${line##* }/& overlay/g" /boot/extlinux/extlinux.conf
				echo "now reboot to open overlay"
			fi
		else
			if [ -f "/etc/default/grub" ]; then
				if [ ! "$(cat /etc/default/grub | grep overlay)" ]; then
					sed -i "s/GRUB_CMDLINE_LINUX=\"/&overlay /g" /etc/default/grub
					if [ $? -ne 0 ];then
								echo "fail"
						Exit
					fi
				fi
				update-grub >/dev/null 2>&1
				if [ ! "$(cat /proc/cmdline | grep overlay)" ]; then
					echo "now reboot to open overlay"
				fi
			else
				line=$(grep "linux" $path)
				if [[ ! $line =~ "overlay" ]]; then
					sed -i "s/${line##* }/& overlay/g" $path
					echo "now reboot to open overlay"
				fi
			fi
		fi
		;;
	n)
		if [ -f "/boot/extlinux/extlinux.conf" ]; then
			sed -i "s/ overlay//g" /boot/extlinux/extlinux.conf
			echo "now reboot to close overlay"
		else
			if [ -f "/etc/default/grub" ]; then
				if [ "$(cat /etc/default/grub | grep overlay)" ]; then
							sed -i "s/overlay //g" /etc/default/grub
					if [ $? -ne 0 ];then
						echo "fail"
						Exit
						fi
				fi
				update-grub >/dev/null 2>&1
				if [ "$(cat /proc/cmdline | grep overlay)" ]; then
					echo "now reboot to close overlay"
				fi
			else
				sed -i "s/ overlay//g" $path
				echo "now reboot to close overlay"
			fi
		fi
		;;
	commit)
		if [ ! "$(echo "$cmdline" | grep overlay)" ]; then
			echo "please open overlay first"
			Exit
		fi
		if [ -f /media/root-rw/backup/rollback_deepfreeze ]; then
            		echo "ERROR: A deepfreeze flag is there, please close it first."
            		Exit
        	fi
		if ls /home/.config/mk_snapshot_* 1> /dev/null 2>&1 ; then
                        echo "ERROR: A snapshot flag is there, please reboot first."
                        Exit
                fi
		if ls /home/.config/mk_rollback_* 1> /dev/null 2>&1 ; then
                        echo "ERROR: A rollback flag is there, please reboot first."
                        Exit
                fi
		.  /etc/kylinoverlay.conf
		if [ -f /media/root-rw/backup/index ];then
		    if [ $(cat /media/root-rw/backup/index) -ge $max_update_num ]; then
			echo "the number of snapshot is match or bigger than max_update_num,please check the max_update_num in /etc/kylinoverlay.conf"
		    	Exit 
		    fi
		fi
		echo $2 > "/home/.config/mk_snapshot_$DATE" ||
			Exit
		echo "now reboot to make snapshot"
		;;
    rollback)
		if [ ! "$(echo "$cmdline" | grep overlay)" ]; then
                        echo "please open overlay first"
                        Exit
                fi
		if ls /home/.config/mk_snapshot_* 1> /dev/null 2>&1 ; then
            		echo "ERROR: A snapshot flag is there, please reboot first."
            		Exit
        	fi
		if [ -f /media/root-rw/backup/rollback_deepfreeze ]; then
            		echo "ERROR: A deepfreeze flag is there, please close it first."
            		Exit
        	fi
		if ls /home/.config/mk_rollback_* 1> /dev/null 2>&1 ; then
                        echo "ERROR: A rollback flag is there, please reboot first."
                        Exit
                fi
        	echo $2 > "/home/.config/mk_rollback_$DATE" ||
			Exit
        	echo "now reboot to rollback."
        ;;
    deepfreeze)
		if [ ! "$(echo "$cmdline" | grep overlay)" ]; then
                        echo "please open overlay first"
                        Exit
                fi
		if ls /home/.config/mk_snapshot_* 1> /dev/null 2>&1 ; then
            		echo "ERROR: A snapshot flag is there, please reboot first."
            		Exit
        	fi
		if ls /home/.config/mk_rollback_* 1> /dev/null 2>&1 ; then
                        echo "ERROR: A rollback flag is there, please reboot first."
                        Exit
                fi
        	if [ -f "/media/root-rw/backup/rollback_deepfreeze" ]; then
                	sudo rm -rf /media/root-rw/backup/rollback_deepfreeze
                	echo "overlay deepfreeze close"
        	else
                	sudo touch "/media/root-rw/backup/rollback_deepfreeze"
                	echo "overlay deepfreeze open"
        	fi
		;;
	l)
		for file in $(ls -a "/media/root-rw/backup/"); do
			tmp=${file##*/}
			suf=${tmp:0:7}
			if [ $suf = "update_" ]; then
				echo ${tmp#*_}
			fi
		done
		;;
	c)
		if [ "$(ls /home/.config)" ]; then
			rm -rf /home/.config/*
			if [ $? -ne 0 ];then
				echo "cancle fail"
				Exit
			else
				echo "overlay action cancle successed"
			fi
		fi
		;;
	d)
		updatenum=$(cat "/media/root-rw/backup/index")
		if [ -d "/media/root-rw/backup/update_${updatenum}"* ]; then
			rm -rf /media/root-rw/backup/update_${updatenum}*
			if [ $? -ne 0 ];then
	                        echo "delete snapshot fail"
                        	Exit
        	        else
                	        echo "delete snapshot successed, please reboot"
				updatenum=$((updatenum - 1))
	                        echo $updatenum > "/media/root-rw/backup/index"
                	fi
		else
			echo "can't find snapshot"
		fi
		;;
    	*)
		echo "no such options"
        Usage
		;;
esac
Exit
