#!/bin/bash

# 1. Create uuid for backup mount
# 2. Unlock luks-uuid
# 3. Create /tmp/uuid
# 4. Mount /tmp/uuid
# 5. btrfs send
# 5.5 Update 
# 6. umount
# 7. rm /tmp/uuid
# 8. luksclose

function exit_success {
    # Unmount /tmp/uuid
    log "INFO" "Unmounting $BACKUP_DRIVE_MNT"
    umount $BACKUP_DRIVE_MNT

    # Exit
    exit 0
}

function exit_fail {
    # Unmount /tmp/uuid
    log "INFO" "Unmounting $BACKUP_DRIVE_MNT"
    umount $BACKUP_DRIVE_MNT

    # Exit
    exit 1
}

function get_latest {
    DIR=$1
    if [ -f $DIR/$LATEST ]; then
        echo $(cat $DIR/$LATEST)
    else
        echo ""
    fi
}

function update_latest {
    DIR=$1
    NAME=$2
    echo $2 > $DIR/$LATEST
}

function log {
    LEVEL=$1
    MESSAGE=$2
    echo "$LEVEL: $MESSAGE"
}

function notify {
    LEVEL=$1
    MESSAGE=$2
    log "$LEVEL" "$MESSAGE"
    sudo -E -u $USER notify-send "$LEVEL" "$MESSAGE"
}

# Backup info
export BACKUP_DRIVE_UUID={{ disk.uuid }}
export BACKUP_DRIVE_PASSWORD={{ disk.password }}
export BACKUP_DRIVE_TMP_UUID=$(uuidgen)
export BACKUP_DRIVE_NAME=luks-$BACKUP_DRIVE_UUID

# For notifications
export DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/{{ notifications.user.uid }}/bus
export USER={{ notifications.user.name }}

export SOURCE_DIR=${SOURCE_DIR:=/}
# Fix basename / showing up as "/" -> change to "root"
if [ $(basename $SOURCE_DIR) = / ]; then
    export SNAPSHOT_PREFIX=${SNAPSHOT_PREFIX:=root}
else
    export SNAPSHOT_PREFIX=${SNAPSHOT_PREFIX:=$(basename $SOURCE_DIR)}
fi
# Set snapshot prefix based on basename
export SNAPSHOT_TIME=$(date +"%y_%m_%d-%H.%M")
export SNAPSHOT_NAME=$SNAPSHOT_PREFIX-$SNAPSHOT_TIME
export SNAPSHOT_DIR=${SNAPSHOT_DIR:=/.snapshots}
export LATEST=$SNAPSHOT_PREFIX-latest
export BACKUP_DRIVE_MNT=/tmp/$BACKUP_DRIVE_TMP_UUID
export BACKUP_DIR=${BACKUP_DIR:=$BACKUP_DRIVE_MNT/$(hostname)}

# Show snapshot settings
echo "SOURCE_DIR" "$SOURCE_DIR"
echo "SNAPSHOT_PREFIX" "$SNAPSHOT_PREFIX"
echo "SNAPSHOT_TIME" "$SNAPSHOT_TIME"
echo "SNAPSHOT_NAME" "$SNAPSHOT_NAME"
echo "SNAPSHOT_DIR" "$SNAPSHOT_DIR"
echo "LATEST" "$LATEST"
echo "BACKUP_DRIVE_MNT" "$BACKUP_DRIVE_MNT"
echo "BACKUP_DIR" "$BACKUP_DIR"

# Create readonly snapshot
log "INFO" "Creating snapshot from $SOURCE_DIR as $SNAPSHOT_DIR/$SNAPSHOT_NAME"
if [ -d $SNAPSHOT_DIR/$SNAPSHOT_NAME ]; then
    log "WARN" "Snapshot $SNAPSHOT_DIR/$SNAPSHOT_NAME already created. Skipping"
else
    btrfs subvolume snapshot -r $SOURCE_DIR $SNAPSHOT_DIR/$SNAPSHOT_NAME
fi

# Update latest in snapshot dir
log "INFO" "Updating latest in $SNAPSHOT_DIR to $SNAPSHOT_NAME."
update_latest $SNAPSHOT_DIR $SNAPSHOT_NAME

# Unlock backup drive
if [ -L /dev/disk/by-uuid/$BACKUP_DRIVE_UUID ]; then
    cryptsetup luksOpen /dev/disk/by-uuid/$BACKUP_DRIVE_UUID $BACKUP_DRIVE_NAME --key-file=$BACKUP_DRIVE_PASSWORD
    cryptsetup status /dev/mapper/$BACKUP_DRIVE_NAME
else
    log "INFO" "Backup drive $BACKUP_DRIVE_UUID not found"
    log "INFO" "Snapshot $SNAPSHOT_NAME completed successfully."
    notify "WARN" "Drive $BACKUP_DRIVE_UUID could not be found. Snapshot completed without backup."
    exit 0
fi

if [ $? = 0 ]; then
    log "INFO" "Drive $BACKUP_DRIVE_UUID unlocked"
else
    notify "ERROR" "Drive $BACKUP_DRIVE_UUID could not be unlocked."
    exit_fail
fi

# Create /tmp/uuid
log "INFO" "Creating $BACKUP_DRIVE_MNT"
mkdir $BACKUP_DRIVE_MNT

# Mount /tmp/uuid
log "INFO" "Mounting /dev/mapper/$BACKUP_DRIVE_NAME"
mount -t btrfs -o compress=zstd /dev/mapper/$BACKUP_DRIVE_NAME $BACKUP_DRIVE_MNT

if [ $? = 0 ]; then
    log "INFO" "Drive $BACKUP_DRIVE_UUID mounted at $BACKUP_DRIVE_MNT"
else
    notify "ERROR" "Drive $BACKUP_DRIVE_NAME could not be mounted."
    exit_fail
fi

# First check if the snapshot dir has a "latest" snapshot
# This will be needed to send an incremental snapshot
LATEST_SNAPSHOT="$(get_latest $SNAPSHOT_DIR)"
log "INFO" "Latest snapshot is $LATEST_SNAPSHOT"

# Next, check if the backup drive has a "latest" snapshot
LATEST_BACKUP="$(get_latest $BACKUP_DIR)"
log "INFO" "Latest backup is $LATEST_BACKUP"

# Now check if the "latest" snapshots match
# btrfs requires both the sending drive and receiving drive have
# matching parent snapshots.
#
# There are a few scenarios to cover
# 1. Neither the backup drive nor the local snapshot dir have a "latest"
#       This can happen if the backup occurs before any snapshots are
#       taken. Don't send anything.
# 2. The backup drive has a "latest" but the snapshot dir doesn't
#       This can happen when the local drive is restored from backup
#       but the snapshot dir didn't copy over. nothing to send.
# 3. The backup drive and snapshot dir have a "latest" and they are the
#    same.
#       Send backup with parent as normal.
# 4. The snapshot dir has a "latest" but the backup drive doesn't
#       This can happen when backing up for the first time. Send the
#       snapshot without a parent
# 5. Both the snapshot dir and backup drive have a latest, but they are
#    out of sync.
#       This can happen when snapshots are taken with the backup drive
#       disconnected. There's a few sub-scenarios here:
#   a. The snapshot dir has the "latest" snapshot from the backup dir,
#      it's just older than the "latest" snapshot in the snapshot dir
#           Re-sync the "latest" snapshot dir with the one in the
#           backup dir. Send as normal with parents.
#   b. The snapshot dir does not have the "latest" snapshot from the
#      backup dir. 
#           Here be dragons. Something went wrong and will likely need
#           to be manually reconfigured. Raise a critical alert.

# Scenario 1 and 2
if [ "$LATEST_SNAPSHOT" = "" ]; then
    notify "WARN" "Neither the snapshot dir nor the backup drive has a 'latest' snapshot."
    exit_success
fi

# Scenario 3
if [ "$LATEST_SNAPSHOT" = "$LATEST_BACKUP" ]; then
    log "INFO" "Proceeding with backups as normal."
    # Send incremental snapshot
    btrfs send -p $SNAPSHOT_DIR/$LATEST_SNAPSHOT $SNAPSHOT_DIR/$SNAPSHOT_NAME | btrfs receive $BACKUP_DIR
    if [ $? != 0 ]; then
        notify "ERROR" "btrfs send -p $SNAPSHOT_DIR/$LATEST_SNAPSHOT $SNAPSHOT_DIR/$SNAPSHOT_NAME failed."
        exit_fail
    fi

    # Update latest in backup dir
    update_latest $BACKUP_DIR $SNAPSHOT_NAME

    # Update latest in snapshot dir
    update_latest $SNAPSHOT_DIR $SNAPSHOT_NAME

    # Exit
    sudo -E -u $USER notify-send "Backup completed" "INFO: Backup $SNAPSHOT_NAME completed successfully."
    exit_success
fi

# Scenario 4
if [ "$LATEST_BACKUP" = "" ]; then
    log "INFO" "No prior backups detected. Sending full backup."
    # Send incremental snapshot
    btrfs send $SNAPSHOT_DIR/$SNAPSHOT_NAME | btrfs receive $BACKUP_DIR
    if [ $? != 0 ]; then
        notify "ERROR" "btrfs send $SNAPSHOT_DIR/$SNAPSHOT_NAME failed."
        exit_fail
    fi

    # Update latest in backup dir
    update_latest $BACKUP_DIR $SNAPSHOT_NAME

    # Update latest in snapshot dir
    update_latest $SNAPSHOT_DIR $SNAPSHOT_NAME

    # Exit
    notify "INFO" "Backup $SNAPSHOT_NAME completed successfully."
    exit_success
fi

# Scenario 5a
log "INFO" "Detected drift. Attempting to synchronize latest snapshot with backup. Set to $LATEST_BACKUP."
if [ -d $SNAPSHOT_DIR/$LATEST_BACKUP ]; then
    log "INFO" "$LATEST_BACKUP found in snapshot dir. Synchronizing and proceeding."

    btrfs send -p $SNAPSHOT_DIR/$LATEST_BACKUP $SNAPSHOT_DIR/$SNAPSHOT_NAME | btrfs receive $BACKUP_DIR
    if [ $? != 0 ]; then
        notify "ERROR" "btrfs send -p $SNAPSHOT_DIR/$LATEST_SNAPSHOT $SNAPSHOT_DIR/$SNAPSHOT_NAME failed."
        exit_fail
    fi
    
    # Update latest in backup dir
    update_latest $BACKUP_DIR $SNAPSHOT_NAME

    # Update latest in snapshot dir
    update_latest $SNAPSHOT_DIR $SNAPSHOT_NAME

    # Exit
    notify "INFO" "Backup $SNAPSHOT_NAME completed successfully."
    exit_success
# Scenario 5b
else
    log "ERROR" "Something went wrong. $LATEST_BACKUP not found in $SNAPSHOT_DIR."
    notify "ERROR" "$LATEST_BACKUP not found in $SNAPSHOT_DIR."
    exit_fail
fi