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
  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!"