#!/usr/bin/env bash
###################################################################
#
#	File    : snap_grub
#	Author  : adb <adb>
#	Created : 24-08-09
#
###################################################################
#
if [ $(id -u) != 0 ]; then exit 1; fi

# Set the temp directory for mounting the root partition
TEMP_MOUNT_POINT="/tmp/root_dev/"

# Create the temporary mount point dir
if [ ! -d "$TEMP_MOUNT_POINT" ]; then
    #printf "[i] Creating $TEMP_MOUNT_POINT.."
    mkdir -p "$TEMP_MOUNT_POINT"
    #printf "done\n"
fi

# Get existing snapshots
SNAPSHOTS=$(btrfs subvolume list / | grep "snapshots/" | awk '{print $NF}' | sort -r)

GRUB_CONFIG_FILE="/boot/grub/grub.cfg"
GRUB_CUSTOM_FILE="/etc/grub.d/40_custom"

# Check if the grub configuration file exists
if [ ! -f "$GRUB_CONFIG_FILE" ]; then 
    #printf "[x] Cannot proceed. $GRUB_CONFIG_FILE does not exist\n"
    exit 0
fi

# Get the default menu entry data
DEFAULT_MENU_ENTRY=$(sed -n '/^menuentry /,/^}/{p;/^}/q}' $GRUB_CONFIG_FILE)
DISTRO_NAME="$(echo $DEFAULT_MENU_ENTRY | cut -d "'" -f2)"

# Check if snapshots exist
if [ -z "$SNAPSHOTS" ]; then 
    #printf "[!] No snapshots currently exist\n"
    exit 0
fi

# Get the partition layout for the root and boot partitions
ROOT_PARTITION=$(mount | grep " / " | head -n1 | awk '{print $1}')
ROOT_DEVICE=${ROOT_PARTITION::-1}
ROOT_UUID=$(lsblk -f --noheadings $ROOT_PARTITION | head -n1 | awk '{ print $4}')
ROOT_FS=$(lsblk -f --noheadings $ROOT_PARTITION | head -n1 | awk '{ print $2}')

BOOT_PARTITION=$(mount | grep "/boot" | head -n1 | awk '{print $1}')
BOOT_UUID=$(lsblk -f --noheadings $BOOT_PARTITION | head -n1 | awk '{ print $4}')

if [ -z "$ROOT_UUID" ]; then
    #printf "[x] Cannot proceed. No UUID found for root device\n"
    exit 0
fi

# Check if the root (/) filesystem is btrfs
if [ "$ROOT_FS" != "btrfs" ]; then
    #printf "[!] Operation aborted. \"/\" is not a btrfs filesystem\n"
    #echo $ROOT_FS
    exit 0
fi

if [ -z "$BOOT_UUID" ]; then
    # /boot is on the same volume as root
    BOOT_UUID=$ROOT_UUID
    BOOT_PARTITION=$ROOT_PARTITION
    BOOT_DEVICE=$ROOT_DEVICE
else
    BOOT_PARTITION=$(blkid | grep "$BOOT_UUID" | sed 's/://g' | awk '{print $1}')
    BOOT_DEVICE=${ROOT_PARTITION::-1}
    BOOT_FS=$(lsblk -f --noheadings $BOOT_PARTITION | head -n1 | awk '{ print $2}')
fi

function debug(){
    printf "=======================================================\n"
    printf "Root UUID\t: $ROOT_UUID\n"
    printf "Root Device\t: $ROOT_DEVICE\n"
    printf "Root Partiton\t: $ROOT_PARTITION\n"
    printf "Root filesystem\t: $ROOT_FS\n"
    printf "\n"
    printf "Boot UUID\t: $BOOT_UUID\n"
    printf "Boot Device\t: $BOOT_DEVICE\n"
    printf "Boot Partiton\t: $BOOT_PARTITION\n"
    printf "Boot filesystem\t: $BOOT_FS\n"
    printf "=======================================================\n"
    printf "\n"
    
    printf "[i] Current snapshots :\n"
    for SNAPSHOT in $SNAPSHOTS; do
        printf "\t$SNAPSHOT\n"
    done
}

function ClearGrubCustomFile() {

    if [ -f "$GRUB_CUSTOM_FILE" ];then
        CONTENT='#!/bin/sh\nexec tail -n +3 $0'
        printf "$CONTENT\n" > $GRUB_CUSTOM_FILE
    fi
}

# Clear the grub custom file 
ClearGrubCustomFile

if [ ! -z "$SNAPSHOTS" ]; then

    printf "\nsubmenu 'System Snapshots' {\n" >> $GRUB_CUSTOM_FILE

    for SNAP in ${SNAPSHOTS}; do
    
        mount $ROOT_PARTITION $TEMP_MOUNT_POINT
    
        cd $TEMP_MOUNT_POINT
    
        SNAPSHOT="$(echo $SNAP | awk -F@ '{ print $NF}')"

        DIST_ESC=$(echo $DISTRO_NAME | sed 's/[\/\\.]/\\&/g')

        ENTRY_NAME="$(echo "$DEFAULT_MENU_ENTRY" | sed "s/$DIST_ESC/Snapshot : $SNAPSHOT/g")"
    
        if [ -f "$SNAP/vmlinuz" ]; then

            KERNEL_VERSION="$(uname -r)"
    
            SNAPSHOT_MENU_ENTRY=$(echo "$ENTRY_NAME" | sed "s/Loading Linux.*/Loading Linux $KERNEL_VERSION'/g") 
            SNAPSHOT_MENU_ENTRY="$(echo "$SNAPSHOT_MENU_ENTRY" | sed "s/\@\/boot\/vmlinuz-[^ ]*/@snapshots\/@$SNAPSHOT\/boot\/vmlinuz-$KERNEL_VERSION/g")"
            SNAPSHOT_MENU_ENTRY="$(echo "$SNAPSHOT_MENU_ENTRY" | sed "s/subvol=@/subvol=@snapshots\/@$SNAPSHOT/")"
            SNAPSHOT_MENU_ENTRY="$(echo "$SNAPSHOT_MENU_ENTRY" | sed "s/\@\/boot\/initrd.img-[^ ]*/@snapshots\/@$SNAPSHOT\/boot\/initrd.img-$KERNEL_VERSION/g")"
    
            printf "\t$SNAPSHOT_MENU_ENTRY\n\n" >> $GRUB_CUSTOM_FILE
    
        else
            # There is no kernel in the snapshot
            KERNEL_VERSION="$(uname -r)"
            SNAPSHOT_MENU_ENTRY=$(echo "$ENTRY_NAME" | sed "s/Loading Linux.*/Loading Linux $KERNEL_VERSION'/g") 
    
            SET_SUBVOL="$(echo "$SNAPSHOT_MENU_ENTRY" | sed "s/subvol=@/subvol=@snapshots\/@$SNAPSHOT/")"
    
            printf "\t$SET_SUBVOL\n\n" >> $GRUB_CUSTOM_FILE
    
        fi
        
    
        cd 
        umount $TEMP_MOUNT_POINT
    
    done
    
    printf "}\n" >> $GRUB_CUSTOM_FILE

fi


