#!/usr/bin/env bash
#
# omv-microvm-snapshot-restore — overwrite a VM's disk with a snapshot's
# saved disk, and mark it for a warm (memory-state) boot if the snapshot
# has one. Called after the VM's unit is stopped; does not start it back up.
#
# --cold restores the disks only, ignoring any memory/device state. Needed
# when restoring onto a different VM than the snapshot was taken from: the
# saved device state names that VM's own disk files and TAP device, so
# loading it would attach the other VM's disks.
#
# Usage: omv-microvm-snapshot-restore <vm-name> <snapshot-dir> [--cold]

set -eu

. /usr/share/openmediavault/scripts/helper-functions

VMNAME="${1:?usage: omv-microvm-snapshot-restore <vm-name> <snapshot-dir>}"
SNAPSHOT_DIR="${2:?usage: omv-microvm-snapshot-restore <vm-name> <snapshot-dir>}"
FORCE_COLD=0
[ "${3:-}" = "--cold" ] && FORCE_COLD=1

_log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${VMNAME}: $*"
}

die() {
    _log "ERROR: $*"
    exit 1
}

[ -f "${SNAPSHOT_DIR}/rootfs.ext4" ] || die "Snapshot is incomplete: missing ${SNAPSHOT_DIR}/rootfs.ext4"

# Must match the derivation in omv-microvm-run.
SF_UUID=$(omv_config_get "/config/services/microvm/sharedfolderref")
omv_isuuid "$SF_UUID" || die "No shared folder configured for VM storage."
SF_PATH=$(omv_get_sharedfolder_path "$SF_UUID")
[ -n "$SF_PATH" ] || die "Could not resolve shared folder path."
VM_DIR="${SF_PATH}/vms/${VMNAME}"
mkdir -p "$VM_DIR"

WARM=0
if [ -f "${SNAPSHOT_DIR}/vmstate" ] && [ -f "${SNAPSHOT_DIR}/memfile" ]; then
    WARM=1
    if [ "$FORCE_COLD" -eq 1 ]; then
        _log "Restoring onto a different VM — skipping the saved memory state; the VM will cold boot."
        WARM=0
    fi
fi

# A warm snapshot's device state holds the drive paths Firecracker saw:
# chroot paths when jailed, host paths otherwise. So it only loads with the
# jailer setting it was taken with.
if [ "$WARM" -eq 1 ]; then
    SNAP_JAILED=0
    [ -f "${SNAPSHOT_DIR}/jailed" ] && SNAP_JAILED=1
    VM_JAILED=0
    omv_checkyesno "$(omv_config_get "/config/services/microvm/vms/vm[name='${VMNAME}']/jailer")" && VM_JAILED=1
    if [ "$SNAP_JAILED" -ne "$VM_JAILED" ]; then
        if [ "$SNAP_JAILED" -eq 1 ]; then
            die "This snapshot was taken while the VM ran in the jailer, which is now turned off for it. Turn \"Run in jailer\" back on to restore it."
        fi
        die "This snapshot was taken while the VM ran without the jailer, which is now turned on for it. Turn \"Run in jailer\" off to restore it."
    fi
fi

# Checked before anything is overwritten. Saved data disks go back to
# wherever that disk lives now. A warm snapshot's device state has the
# disk's path baked in, so one whose disk has since been removed can't be
# loaded — a cold one just leaves it out.
DISK_ROWS=$(omv-microvm-list-disks "$VMNAME") || die "Could not list data disks."
declare -A DISK_PATHS=()
while IFS=$'\t' read -r name path _ _ _; do
    [ -n "$name" ] || continue
    DISK_PATHS["$name"]="$path"
done <<< "$DISK_ROWS"

declare -a RESTORE_DISKS=()
for saved in "${SNAPSHOT_DIR}"/disk-*.img; do
    [ -f "$saved" ] || continue
    name=$(basename "$saved" .img)
    name="${name#disk-}"
    if [ -z "${DISK_PATHS[$name]:-}" ]; then
        if [ "$WARM" -eq 1 ]; then
            die "Snapshot includes data disk '${name}', which no longer exists. Recreate a disk with that name on this VM first."
        fi
        _log "WARNING: snapshot includes data disk '${name}', which no longer exists — skipping it."
        continue
    fi
    RESTORE_DISKS+=("$name")
done

_log "Overwriting current disk with snapshot's saved disk ..."
cp --reflink=auto "${SNAPSHOT_DIR}/rootfs.ext4" "${VM_DIR}/rootfs.ext4"

for name in "${RESTORE_DISKS[@]}"; do
    _log "Overwriting data disk '${name}' with snapshot's saved copy ..."
    mkdir -p "$(dirname "${DISK_PATHS[$name]}")"
    cp --reflink=auto "${SNAPSHOT_DIR}/disk-${name}.img" "${DISK_PATHS[$name]}"
done

# Only warm snapshots (with memory/device state) need the next start to
# load them over the API — a cold snapshot is just a disk, so a normal
# cold boot afterward is already correct.
# A stale marker (e.g. from an earlier failed warm load) would otherwise
# make a cold restore's next start wait for a snapshot load.
if [ "$WARM" -eq 1 ]; then
    printf '%s' "$SNAPSHOT_DIR" > "${VM_DIR}/.restore_from"
else
    rm -f "${VM_DIR}/.restore_from"
fi

_log "Ready to restore from ${SNAPSHOT_DIR}."
exit 0
