#!/bin/sh

set -eu

image_dir=/home/partimag

ask_disk() {
  disks=$(awk '$4 ~ /^((md|mmcblk)[0-9]+|nvme[0-9]+n[0-9]+|[sv]d[a-z]|xvd[a-z])$/ {
                 print $4
               }' /proc/partitions)
  [ -n "$disks" ] || return 1
  printf %s "$disks" | fzf -1 --height=1 --layout=reverse-list
}

ask_restore_image() {
  available_images=$(
    find "$image_dir" -mindepth 1 -maxdepth 1 -type d ! -name '.*' -print0 \
      | xargs -0 --no-run-if-empty -n 1 basename)
  [ -n "$available_images" ] || return 1
  printf %s "$available_images" | fzf -1 --height=5 --layout=reverse-list --color=bg+:#FFFFFF,fg+:#000000
}

bootserver_address=$(cat /run/puavo/nbd-server 2>/dev/null) || true

if [ -z "$bootserver_address" ]; then
  echo 'Not running on netboot device, exiting...' >&2
  exit 1
fi

mkdir -p "$image_dir"
if ! mountpoint -q "$image_dir"; then
  sshfs "clonezilla.images@${bootserver_address}:" /home/partimag
fi

mode=${1:-}

case "$mode" in
  restoredisk)
    if ! restore_image=$(ask_restore_image); then
      echo 'No images to restore from' >&2
      exit 1
    fi
    echo 'Choose a target disk for restore:'
    if ! target_disk=$(ask_disk); then
      echo 'No disk to restore to' >&2
      exit 1
    fi
    echo "Chosen << ${target_disk} >>"
    echo "Secure Erasing the disk first, please wait."

    # We want to try and maintain data security as well as possible, so erase disk first before writing, even if it's empty already.
    # Always prefer the most secure way of erasing. If that fails, puavo-disk-erase will fallback to less secure erasing, one at a time until something works.
    /usr/sbin/puavo-disk-erase --force --normal "$target_disk"

    # XXX some patch to ocs-sr for this would be nice
    # We must disable partition resizing to the LVM partitions because we
    # want to keep that in our control, but we do like the resizing for the
    # Windows partitions.  Thus disable the "ocs-expand-lvm" script in this
    # crude way.
    ln -fns /usr/bin/true /usr/sbin/ocs-expand-lvm
    # Then run the actual restore:
    ocs-sr -batch -g auto -e1 auto -e2 -r -j2 -scr -k1 -p true \
           restoredisk "$restore_image" "$target_disk"

    # Resize LVM partition to match the physical partition size
    # and extend puavo-home accordingly
    lvm_part=$(lsblk -n -l -o NAME,PARTTYPENAME "/dev/$target_disk" | awk '/Linux LVM/ {print $1}')
    restore_size=$(awk '$4 == "disk" && /^(nvme[0-9]+n[0-9]+|[sv]d[a-z]) / {print $3}' "$image_dir/$restore_image/blkdev.list")
    target_size=$(lsblk -n -l -o SIZE,NAME "/dev/$target_disk" | awk '/'"$target_disk"'$/ {print $1}')

    if [ "${restore_size}" != "${target_size}" ]; then
      if [ -n "${lvm_part}" ]; then
        printf "%s" "Linux LVM found, resizing to occupy unclaimed free space.."
        pvresize "/dev/$lvm_part"
        lvextend -l +100%FREE /dev/mapper/puavo-home
        vgchange -a y puavo
        e2fsck -f /dev/mapper/puavo-home
        resize2fs /dev/mapper/puavo-home
      fi
    fi
    ;;
  savedisk)
    product_name=$(dmidecode -s system-product-name | tr ' /.' ___)

    echo 'Choose a source disk to save image from:'
    if ! source_disk=$(ask_disk); then
      echo 'No disk to save' >&2
      exit 1
    fi
    echo "Chosen << ${source_disk} >>"

    if ! disk_size="$(facter "blockdevice_${source_disk}_size")"; then
      disk_size="0"
    fi
    if [ "$disk_size" -gt 0 ]; then
      image_name="${product_name}_${disk_size}-$(date -u +%Y-%m-%d-%H%M%S)"
    else
      image_name="${product_name}-$(date -u +%Y-%m-%d-%H%M%S)"
    fi

    read -p "Use the following disk image name: [${image_name}] " answer
    if [ -n "$answer" ]; then
      image_name="$answer"
    fi

    ocs-sr -rm-win-swap-hib -q2 -j2 -z1p -i 4096 -fsck-src-part -scs -p true \
           savedisk "$image_name" "$source_disk"
    ;;
  *)
    echo "Unknown mode '${mode}'" >&2
    exit 1
    ;;
esac
