#!/usr/bin/env bash
#
# omv-microvm-snapshot-restore-hook — ExecStartPost hook that loads a
# snapshot into the "empty" Firecracker process omv-microvm-run just
# started in restore mode. No-op on every normal start.
#
# A failed load here fails the whole start job (KillMode=control-group
# kills Firecracker too), landing the unit in "failed" — which
# getActiveState() already maps to the UI's "error" state.
#
# Usage: omv-microvm-snapshot-restore-hook <vm-name>

set -eu

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

VMNAME="${1:?usage: omv-microvm-snapshot-restore-hook <vm-name>}"

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

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

# Must match the derivation in omv-microvm-run.
SF_UUID=$(omv_config_get "/config/services/microvm/sharedfolderref")
omv_isuuid "$SF_UUID" || exit 0
SF_PATH=$(omv_get_sharedfolder_path "$SF_UUID")
[ -n "$SF_PATH" ] || exit 0
VM_DIR="${SF_PATH}/vms/${VMNAME}"
MARKER="${VM_DIR}/.restore_from"

[ -f "$MARKER" ] || exit 0

SNAPSHOT_DIR=$(cat "$MARKER")
_log "Restore marker present — loading snapshot from ${SNAPSHOT_DIR} ..."

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

for _ in $(seq 1 50); do
    [ -S "$SOCKET_PATH" ] && break
    sleep 0.1
done
[ -S "$SOCKET_PATH" ] || die "API socket never appeared: ${SOCKET_PATH}"

# A jailed VM sees the snapshot at /snapshot-in (bind-mounted by
# omv-microvm-run, which also handed its files to the jail user).
JAILED=0
FC_SNAPSHOT_DIR="$SNAPSHOT_DIR"
if [ -f "${RUN_DIR}/jailed" ]; then
    JAILED=1
    FC_SNAPSHOT_DIR="/snapshot-in"
fi
# Firecracker keeps the files open once loaded, so they can go back to root
# right away, whether the load worked or not.
_restore_owner() {
    if [ "$JAILED" -eq 1 ]; then
        chown root:root "${SNAPSHOT_DIR}/memfile" "${SNAPSHOT_DIR}/vmstate" || true
    fi
}

BODY=$(jq -n \
    --arg mem "${FC_SNAPSHOT_DIR}/memfile" \
    --arg state "${FC_SNAPSHOT_DIR}/vmstate" \
    '{
        "mem_backend": {"backend_type": "File", "backend_path": $mem},
        "snapshot_path": $state,
        "resume_vm": true
    }')

resp=$(mktemp)
http_code=$(curl -sS --unix-socket "$SOCKET_PATH" -o "$resp" -w '%{http_code}' \
    -X PUT 'http://localhost/snapshot/load' \
    -H 'Content-Type: application/json' -d "$BODY") || { rm -f "$resp"; _restore_owner; die "Firecracker API request failed: PUT /snapshot/load"; }
_restore_owner

if [ "$http_code" -lt 200 ] || [ "$http_code" -ge 300 ]; then
    _log "Firecracker API error (${http_code}): $(cat "$resp")"
    rm -f "$resp"
    die "Snapshot load rejected (HTTP ${http_code}) — leaving marker in place for diagnosis."
fi
rm -f "$resp"

# Only removed on success, so a failed load stays diagnosable instead of
# silently falling back to a cold boot next time.
rm -f "$MARKER"
_log "Snapshot loaded and VM resumed."
exit 0
