#!/usr/bin/env bash
###################################################################
#
#	File    : auto_snapshots
#	Author  : adb <alex@alexdball.dev>
#	Created : 26-02-22
#
###################################################################
#
#
# Define variables
DEBUG=0
#
# Teminal Variables
red='\033[0;31m'
green='\033[1;32m'
blue='\033[1;34m'
yellow='\033[1;33m'
nc='\033[0m'

#
# Message variables
NOT_ROOT="\n${red}[x]${nc} This program requires ${red}root${nc} privileves\n"
SEL_SUBVOL="${red}[x]${nc} Cannot Proceed. Please enter the ${red}source${nc} subvolume\n"
SEL_SNAPVOL="${red}[x]${nc} Cannot Proceed. Please enter the ${red}'snapshots'${nc} subvolume\n"

# Application Variables
KEEP_SNAPSHOTS=50
RUNTIME=$(date +%Y-%m-%d_%H_%M_%S)
SNAP_MOUNT_DIR="/tmp/auto_snapshots"
BTRFS="$(which btrfs)"

function GetFilesystemInfo(){
    ROOT_INFO="$(df -T | grep "/$")"

    if [ -z "$ROOT_INFO" ]; then
        printf "${red}[x]${nc} Cannot proceed. Failed to determine system partition layout\n"
        exit 0
    fi
    ROOT_DEVICE=$(echo "$ROOT_INFO" | awk '{print $1}')
    ROOT_DEVICE=${ROOT_DEVICE::-1}
    ROOT_PARTITION=$(echo "$ROOT_INFO" | awk '{print $1}')
    ROOT_PARTITION_FS=$(echo "$ROOT_INFO" | awk '{print $2}')
    ROOT_MOUNT_POINT="$(echo "$ROOT_INFO" | awk '{print $NF}')"

    if [ ! "$ROOT_PARTITION_FS" == "btrfs" ]; then
        echo -e "${red}[x]${nc} The root partition is NOT btrfs"
        exit 0
    fi
}

# Check for root privileges
if [ "$(id -u)" != "0" ]; then echo -e "$NOT_ROOT"; exit 1; fi

# Functions
function Usage(){
    printf "\n[i] Usage : %s <parameters>\n\n" "$0"
    printf "Where Parameters :\n"
    printf "\t-v <subvolume>\t\t  - The subvolume to create a snapshot for\n"
    printf "\t-s <snapshots_volume>\t  - The subvolume where snaphots are saved\n"
    printf "\t-d <snapshot_description> - The description of the snapshot.\n"
    printf "\t   Defaults to : <subvolume>_YYYY_MM_DD_HH_mm_ss\n"
    printf "\t-h Displays this message and exits\n"
	exit 0;
}

# Get the given parameters
while getopts v:s:d:h flag
do
    case "${flag}" in
        v) SELECTED_SUB_VOL=${OPTARG};;
        s) SELECTED_SNAP_VOL=${OPTARG};;
        d) SELECTED_SNAP_DESC=${OPTARG};;
        h) Usage;
    esac
done

# Get root partition filesystem.
GetFilesystemInfo

# Check if parameters are valid
if [ "$SELECTED_SUB_VOL" == "" ]; then 
    echo -e "$SEL_SUBVOL" 
    Usage
fi
if [ "$SELECTED_SNAP_VOL" == "" ]; then
    echo -e "$SEL_SNAPVOL"
    Usage
fi

# Mount the partition to the temp directory
# Check if the mount directory exists
if [ ! -d ${SNAP_MOUNT_DIR} ]; then
    echo -e "${yellow}[!]${nc} Mount directory $SNAP_MOUNT_DIR does not exist. Creating..."
    mkdir -p "$SNAP_MOUNT_DIR"
else
    echo -e "${green}[i]${nc} Mount directory $SNAP_MOUNT_DIR exists."
fi

function mount_root_partition() {
    # Check if already mounted
    IS_MOUNTED="$(mount | grep $SNAP_MOUNT_DIR)"

    if [ -z "$IS_MOUNTED" ]; then
        echo
        echo -e "${green}==>${nc} Mounting $ROOT_PARTITION to $SNAP_MOUNT_DIR"
        mount "$ROOT_PARTITION" "$SNAP_MOUNT_DIR" || exit 1
    else
        echo -e "${red}[x]${nc} Partition already mounted. Aborting"
        exit 0
    fi
}

function unmount_root_partition() {
    # Check if already mounted
    IS_MOUNTED="$(mount | grep $SNAP_MOUNT_DIR)"
    if [ -z "$IS_MOUNTED" ]; then
        echo -e "${red}[x]${nc} Partition not mounted. Aborting"
        exit 0
    else
        echo
        echo -e "${green}==>${nc} Unmounting $ROOT_PARTITION from $SNAP_MOUNT_DIR"
        umount "$SNAP_MOUNT_DIR" || exit 1
    fi
}

function CheckSubVolExists(){
    ROOT_SUBVOL=$(btrfs subvol list . | awk '{print $NF}' | grep "$SELECTED_SUB_VOL$")
}

function CheckSnapshotVolExists(){
    SNAP_SUBVOL=$(btrfs subvol list . | awk '{print $NF}' | grep "$SELECTED_SNAP_VOL$")
}

function GetSnapshotCount(){
    # Show existing volume snapshots
    SNAPSHOT_COUNT="$($BTRFS subvol list . | grep "$SELECTED_SNAP_VOL"/ | wc -l)"
}

function GetOldestSnapshot(){
    OLDEST_SNAPSHOT="$($BTRFS subvol list . | grep "$SELECTED_SNAP_VOL"/ | awk '{print $NF}' | head -n1 )"
}

function CreateSnapshot(){
    # Create a snapshot
    echo
    echo -e "${yellow}==>${nc} Creating snapshot of $SELECTED_SUB_VOL in $SELECTED_SNAP_VOL/$SELECTED_SUB_VOL""$SNAP_DESC""$RUNTIME"
    $BTRFS subvolume snapshot "$SELECTED_SUB_VOL" "$SELECTED_SNAP_VOL"/"$SELECTED_SUB_VOL""$SNAP_DESC""$RUNTIME" >/dev/null
}

function DeleteSnapshot() {
    # Check if there are snapshots to be deleted
    if [ "$SNAPSHOT_COUNT" -ge "$KEEP_SNAPSHOTS" ]; then
    	echo -e "${red}==>${nc} Deleting snapshot : ${OLDEST_SNAPSHOT}\n"
    	if [ "$OLDESDT_SNAPSHOT" != "@snapshots" ]; then
            $BTRFS subvolume delete "$OLDEST_SNAPSHOT" > /dev/null
        fi
    fi

}


# check for custom snapshot description
if [ "$SELECTED_SNAP_DESC" != "" ]; then SNAP_DESC=$SELECTED_SNAP_DESC"_"; fi
function debug(){
    echo -e "${green}==>${nc} Found btrfs partition for /\n"
    echo -e "${blue}Device     ${nc}                : $ROOT_DEVICE"
    echo -e "${blue}Partition  ${nc}                : $ROOT_PARTITION"
    echo -e "${blue}FileSystem ${nc}                : $ROOT_PARTITION_FS"
    echo
    echo -e "${blue}Root subvolume      ${nc}       : $SELECTED_SUB_VOL"
    echo -e "${blue}Snapshots subvolume ${nc}       : $SELECTED_SNAP_VOL"
    echo -e "${blue}Snapshot Description ${nc}      : $SNAP_DESC$RUNTIME"
    echo
    echo -e "${blue}Snapshot mount directory ${nc}  : $SNAP_MOUNT_DIR"
}

if [ $DEBUG != 0 ]; then debug ; fi

mount_root_partition

cd $SNAP_MOUNT_DIR || exit 1

CheckSubVolExists

if [ -z "$ROOT_SUBVOL" ]; then
    echo -e "${red}[x]${nc} Subvolume $SELECTED_SUB_VOL does not exist on $ROOT_PARTITION"
    cd - >/dev/null || exit 0
    unmount_root_partition
fi

CheckSnapshotVolExists

if [ -z "$SNAP_SUBVOL" ]; then
    echo -e "${red}[x]${nc} Snapshots subvolume $SELECTED_SNAP_VOL does not exist on $ROOT_PARTITION"
    cd - >/dev/null || exit 0
    unmount_root_partition
fi

GetSnapshotCount
echo -e "${green}[i]${nc} Current Snapshot count    : $SNAPSHOT_COUNT (${yellow}$KEEP_SNAPSHOTS${nc} max)"

GetOldestSnapshot
echo -e "${green}[i]${nc} Oldest Snapshot           : $OLDEST_SNAPSHOT"

DeleteSnapshot
CreateSnapshot
cd - >/dev/null || exit 0
unmount_root_partition

if [ "$ROOT_MOUNT_POINT" == "/" ];then
	# Clean the file
	custom_file="/etc/grub.d/40_custom"
	sed -i '/submenu/,$d' "$custom_file"
	/usr/local/bin/snap_grub > $custom_file
	echo -e "${green}[i]${nc} Updating grub ..."
	update-grub 2> /dev/null
	echo -e "\n${green}==>${nc} Process complete\n"
fi

exit 0
