if [[ "$EUID" -ne 0 ]]; then
  echo "Please run the script as root."
  exit 1
fi

usage() {
  echo "Usage: $0 [-m partition] [-b backup_location]"
  exit 1
}

cleanup() {
  if [ -d "/persist.bak" ]; then btrfs -q subvolume delete "/persist.bak"; fi
  if [ -n "$backup_location" ]; then rm -f "$backup_location.tmp"; fi

  if [ -n "$mount_location" ]; then
    if mount | grep -q "$mount_location"; then umount "$mount_location"; fi
    if [ -d "$mount_location" ]; then rmdir "$mount_location"; fi
  fi
}

partition=""
backup_location=""
mount_location=""

trap cleanup EXIT

while getopts "m:b:" opt; do
  case "$opt" in
    m) partition="$OPTARG" ;;
    b) backup_location="$OPTARG" ;;
    *) usage ;;
  esac
done

if [ -n "$partition" ]; then
  mkdir -p "/mnt"
  mount_location=$(mktemp -d /mnt/backup.XXXXXX)
  echo "Mounting $partition at $mount_location..."
  mount "$partition" "$mount_location"
fi

if [ -z "$mount_location" ]; then
  if [[ "$backup_location" != /* ]]; then
    backup_location="$(realpath "$backup_location")"
  fi
else
  if [[ "$backup_location" = /* ]]; then
    echo "Error: When a partition is mounted, backup_location must be relative."
    exit 1
  fi

  backup_location="$(realpath "$mount_location/$backup_location")"
fi

backup_location="$backup_location/$(hostname)-$(date +%Y-%m-%d-%H-%M-%S).btrfs.gz"

echo "Creating /persist snapshot..."
btrfs -q subvolume snapshot -r "/persist" "/persist.bak"

echo "Creating backup at $backup_location..."
btrfs -q send "/persist.bak" | gzip > "$backup_location.tmp"

mv "$backup_location.tmp" "$backup_location"

echo "Backup completed successfully!"