#!/usr/bin/env bash
#
# omv-microvm-snapshot-create — snapshot a VM's disk into <snapshot-dir>.
# If it's running, also pauses it and dumps memory/device state first
# (warm snapshot); if not, just copies the disk (cold snapshot).
#
# Data disks are copied alongside as disk-<name>.img, except those with
# "include in snapshots and backups" turned off — those are left out and
# keep whatever contents they have at restore time.
#
# A jailed VM's Firecracker can only write inside its chroot, so its
# memory/device state goes to the staging directory omv-microvm-run
# bind-mounts in, and is moved here afterwards. Such a snapshot is marked
# with a "jailed" file: its device state refers to chroot paths, so it can
# only be loaded by a jailed VM again (see omv-microvm-snapshot-restore).
#
# Usage: omv-microvm-snapshot-create <vm-name> <snapshot-dir>

set -eu

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

VMNAME="${1:?usage: omv-microvm-snapshot-create <vm-name> <snapshot-dir>}"
SNAPSHOT_DIR="${2:?usage: omv-microvm-snapshot-create <vm-name> <snapshot-dir>}"

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

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

# Must match the derivation in omv-microvm-run.
RUN_DIR="/run/openmediavault-microvm/${VMNAME}"
SOCKET_PATH="${RUN_DIR}/firecracker.socket"

# 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}"
ROOTFS_PATH="${VM_DIR}/rootfs.ext4"
[ -f "$ROOTFS_PATH" ] || die "Rootfs missing: ${ROOTFS_PATH}"

DISK_ROWS=$(omv-microvm-list-disks "$VMNAME") || die "Could not list data disks."

copy_disks() {
    cp --reflink=auto "$ROOTFS_PATH" "${SNAPSHOT_DIR}/rootfs.ext4"
    local name path backup
    while IFS=$'\t' read -r name path _ backup _; do
        [ -n "$name" ] || continue
        if [ "$backup" -ne 1 ]; then
            _log "Skipping data disk '${name}' (excluded from snapshots and backups)."
            continue
        fi
        [ -f "$path" ] || die "Data disk '${name}' missing: ${path}"
        _log "Copying data disk '${name}' ..."
        cp --reflink=auto "$path" "${SNAPSHOT_DIR}/disk-${name}.img"
    done <<< "$DISK_ROWS"
}

mkdir -p "$SNAPSHOT_DIR"

if [ ! -S "$SOCKET_PATH" ]; then
    _log "VM not running — copying disks only (cold snapshot)."
    copy_disks
    _log "Snapshot complete."
    exit 0
fi

_fc_api() {
    # _fc_api <method> <path> <json-body>
    local method=$1 path=$2 body=$3 http_code resp
    resp=$(mktemp)
    http_code=$(curl -sS --unix-socket "$SOCKET_PATH" -o "$resp" \
        -w '%{http_code}' -X "$method" "http://localhost${path}" \
        -H 'Content-Type: application/json' -d "$body") || { rm -f "$resp"; die "Firecracker API request failed: ${method} ${path}"; }
    if [ "$http_code" -lt 200 ] || [ "$http_code" -ge 300 ]; then
        _log "Firecracker API error (${http_code}): $(cat "$resp")"
        rm -f "$resp"
        die "Firecracker API request rejected: ${method} ${path} (HTTP ${http_code})"
    fi
    rm -f "$resp"
}

# Always resume on exit if we paused — a failure partway through must
# never leave the VM frozen. Cleared once the explicit resume succeeds.
PAUSED=0
on_exit() {
    if [ "$PAUSED" -eq 1 ]; then
        _log "Resuming VM after a failed snapshot attempt ..."
        curl -sS --unix-socket "$SOCKET_PATH" -X PATCH 'http://localhost/vm' \
            -H 'Content-Type: application/json' -d '{"state":"Resumed"}' >/dev/null 2>&1 || true
    fi
}
trap on_exit EXIT

_log "Pausing VM ..."
_fc_api PATCH /vm '{"state":"Paused"}'
PAUSED=1

# Must match omv-microvm-run.
JAILED=0
JAIL_STAGE="${VM_DIR}/.jail-snapshot"
if [ -f "${RUN_DIR}/jailed" ]; then
    JAILED=1
    rm -f "${JAIL_STAGE}/memfile" "${JAIL_STAGE}/vmstate"
    FC_SNAPSHOT_DIR="/snapshot-out"
else
    FC_SNAPSHOT_DIR="$SNAPSHOT_DIR"
fi

_log "Writing snapshot to ${SNAPSHOT_DIR} ..."
_fc_api PUT /snapshot/create "$(jq -n \
    --arg mem "${FC_SNAPSHOT_DIR}/memfile" \
    --arg state "${FC_SNAPSHOT_DIR}/vmstate" \
    '{"mem_file_path": $mem, "snapshot_path": $state, "snapshot_type": "Full"}')"

# Before resuming, so the disk copies stay consistent with the memory dump.
_log "Copying disks ..."
copy_disks

_log "Resuming VM ..."
_fc_api PATCH /vm '{"state":"Resumed"}'
PAUSED=0

# After resuming: when the destination is on another filesystem (a backup
# target) this is a full copy of the memory dump, which the VM needn't sit
# paused through.
if [ "$JAILED" -eq 1 ]; then
    _log "Moving memory/device state out of the jail ..."
    mv "${JAIL_STAGE}/memfile" "${JAIL_STAGE}/vmstate" "${SNAPSHOT_DIR}/" \
        || die "Could not move the snapshot out of the jail staging directory."
    chown root:root "${SNAPSHOT_DIR}/memfile" "${SNAPSHOT_DIR}/vmstate"
    touch "${SNAPSHOT_DIR}/jailed"
fi

_log "Snapshot complete."
exit 0
