#!/bin/bash
# migrate-vm-luks — Convert an existing Proxmox VM disk to LUKS encryption
# Part of proxmox-kms-bridge (45Drives SecureVM for Proxmox)
#
# Usage: migrate-vm-luks [--dry-run] <VMID>
#
# Detects the VM's primary boot disk, creates a LUKS2 container, copies
# the disk contents into it, provisions a key in OpenBao, repoints the
# VM config, and attaches the hookscript. Saves rollback state.
set -euo pipefail

CONF_FILE="/etc/qxvault/qxvault.conf"

# ── Defaults (overridden by config) ────────────────────────────────
LUKS_IMAGE_DIR="/var/lib/vz/images"
LUKS_IMAGE_TEMPLATE="vm-{vmid}-disk-qxvault.luks"
LUKS_MAPPER_TEMPLATE="vm{vmid}_crypt"
OPENBAO_ADDR="https://127.0.0.1:8200"
OPENBAO_CACERT=""
OPENBAO_SKIP_VERIFY="false"
OPENBAO_TIMEOUT="10"
KEY_PREFIX=""

if [ -f "$CONF_FILE" ]; then
  # shellcheck source=/dev/null
  . "$CONF_FILE"
fi

# ── Arguments ──────────────────────────────────────────────────────
DRY_RUN=false
VMID=""

for arg in "$@"; do
  case "$arg" in
    --dry-run) DRY_RUN=true ;;
    -h|--help)
      echo "Usage: migrate-vm-luks [--dry-run] <VMID>"
      echo "  --dry-run   Show what would happen without making changes"
      exit 0
      ;;
    *)
      if [ -z "$VMID" ]; then
        VMID="$arg"
      else
        echo "error: unexpected argument '$arg'" >&2; exit 1
      fi
      ;;
  esac
done

[ -n "$VMID" ] || { echo "Usage: migrate-vm-luks [--dry-run] <VMID>" >&2; exit 1; }
[[ "$VMID" =~ ^[0-9]+$ ]] || { echo "error: VMID must be numeric" >&2; exit 1; }

# TARGET_DIR and TARGET_FILE are set after detecting the source disk's storage location
MAPPER="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$VMID}"
HOOKSCRIPT="local:snippets/qxvault-luks-hook.sh"
STATE_DIR="/root/qxvault-migrate-${VMID}"
LUKS_OVERHEAD_BYTES=$((16 * 1024 * 1024))

# ── Helpers ────────────────────────────────────────────────────────
log() { echo "[migrate-vm-luks] $*"; }
fail() { echo "[migrate-vm-luks] ERROR: $*" >&2; exit 1; }

need_cmd() {
  command -v "$1" >/dev/null 2>&1 || fail "Missing command: $1"
}

resolve_source_path() {
  local vol="$1"
  if [[ "$vol" = /* ]]; then
    printf '%s\n' "$vol"
  else
    pvesm path "$vol"
  fi
}

# ── Auto-detect cluster name for key prefix ────────────────────────
detect_cluster_name() {
  # Try corosync config first (works even if corosync isn't running)
  if [ -f /etc/pve/corosync.conf ]; then
    local name
    name=$(grep -oP '^\s*cluster_name:\s*\K\S+' /etc/pve/corosync.conf 2>/dev/null) || true
    if [ -n "$name" ]; then
      printf '%s\n' "$name"
      return 0
    fi
  fi
  # Fallback: use hostname
  hostname -s 2>/dev/null || printf '%s\n' "$HOSTNAME"
}

# Build the vault transit key name for a VMID.
# If KEY_PREFIX is set (or auto-detected), the name is "{prefix}-{vmid}".
# Otherwise, falls back to just "{vmid}" for backward compatibility.
build_key_name() {
  local vmid="$1"
  local prefix="${KEY_PREFIX}"
  if [ -z "$prefix" ]; then
    prefix="$(detect_cluster_name)"
  fi
  if [ -n "$prefix" ]; then
    printf '%s\n' "${prefix}-${vmid}"
  else
    printf '%s\n' "$vmid"
  fi
}

# Check whether a transit key already exists in QxVault.
transit_key_exists() {
  local key_name="$1"
  local http_code
  http_code=$(curl "${CURL_ARGS[@]}" -o /dev/null -w '%{http_code}' \
    -H "X-Vault-Token: $(cat /run/qxvault/token 2>/dev/null || true)" \
    "${OPENBAO_ADDR}/v1/transit/keys/${key_name}" 2>/dev/null) || true
  [ "$http_code" = "200" ]
}

# Resolve the key name for a new encryption, avoiding collisions.
# Writes the chosen name to the key-name file.
resolve_new_key_name() {
  local vmid="$1"
  local target_dir="$2"
  local candidate
  candidate="$(build_key_name "$vmid")"

  local counter=0
  while transit_key_exists "$candidate"; do
    counter=$((counter + 1))
    candidate="$(build_key_name "$vmid")-${counter}"
    if [ "$counter" -gt 50 ]; then
      fail "Too many key name collisions for VM ${vmid} (tried up to ${candidate})"
    fi
  done

  # Persist the resolved key name alongside the LUKS image
  mkdir -p "$target_dir"
  printf '%s\n' "$candidate" > "${target_dir}/key-name"
  chmod 600 "${target_dir}/key-name"
  printf '%s\n' "$candidate"
}

# Read the stored key name for an existing VM.
# Falls back to build_key_name(), then bare VMID for backward compat.
read_key_name() {
  local vmid="$1"
  local luks_dir="$2"
  # Check for stored key-name file
  if [ -f "${luks_dir}/key-name" ]; then
    cat "${luks_dir}/key-name"
    return 0
  fi
  # No stored name — use bare VMID (backward compat with pre-prefix installs)
  printf '%s\n' "$vmid"
}

# ── Preflight ──────────────────────────────────────────────────────
[ "$(id -u)" -eq 0 ] || fail "Must run as root"

for c in qm pvesm qemu-img cryptsetup jq base64 awk sed grep stat truncate blockdev; do
  need_cmd "$c"
done

# Validate OpenBao connectivity before proceeding
log "Verifying QxVault connectivity at ${OPENBAO_ADDR}..."
CURL_ARGS=(-s --connect-timeout "$OPENBAO_TIMEOUT")
if [ -n "$OPENBAO_CACERT" ]; then CURL_ARGS+=(--cacert "$OPENBAO_CACERT"); fi
if [ "$OPENBAO_SKIP_VERIFY" = "true" ]; then CURL_ARGS+=(-k); fi
HEALTH_CODE=$(curl "${CURL_ARGS[@]}" -o /dev/null -w '%{http_code}' "${OPENBAO_ADDR}/v1/sys/health" 2>/dev/null) || true
if [ "$HEALTH_CODE" != "200" ]; then
  fail "Cannot reach QxVault at ${OPENBAO_ADDR} (HTTP ${HEALTH_CODE:-000}). Configure qxvault.conf first."
fi
log "QxVault is reachable and healthy."

mkdir -p "$STATE_DIR"

CONFIG_MODIFIED=false

cleanup_on_error() {
  echo "[migrate-vm-luks] Migration failed. State saved in ${STATE_DIR}" >&2
  # Close mapper if we opened it
  if cryptsetup status "$MAPPER" >/dev/null 2>&1; then
    cryptsetup close "$MAPPER" 2>/dev/null || true
  fi
  # Remove the failed LUKS target
  if [ "$USE_LVM_THIN" = "true" ]; then
    if [ -e "${TARGET_FILE:-}" ]; then
      echo "[migrate-vm-luks] Removing incomplete LUKS block device: ${TARGET_FILE}" >&2
      if [[ "$TARGET_FILE" == /dev/zvol/* ]]; then
        # ZFS zvol — extract dataset name from /dev/zvol/pool/name
        local zvol_ds="${TARGET_FILE#/dev/zvol/}"
        zfs destroy -f "$zvol_ds" 2>/dev/null || true
      elif [ "${RBD_MODE:-}" = "true" ]; then
        # RBD — unmap then remove the image
        rbd unmap "$TARGET_FILE" 2>/dev/null || true
        rbd rm "${RBD_TARGET_SPEC}" "${RBD_MAP_ARGS[@]}" 2>/dev/null || true
      else
        lvremove -f "$TARGET_FILE" 2>/dev/null || true
      fi
    fi
    # Unmap source RBD if it was mapped
    if [ "${RBD_MODE:-}" = "true" ] && [ -n "${SRC_PATH:-}" ] && [[ "${SRC_PATH:-}" == /dev/rbd* ]]; then
      rbd unmap "$SRC_PATH" 2>/dev/null || true
    fi
  else
    if [ -n "${TARGET_FILE:-}" ] && [ -f "$TARGET_FILE" ]; then
      echo "[migrate-vm-luks] Removing incomplete LUKS file: ${TARGET_FILE}" >&2
      rm -f "$TARGET_FILE"
    fi
  fi
  if [ -n "${TARGET_DIR:-}" ] && [ -f "${TARGET_DIR}/key-name" ]; then
    rm -f "${TARGET_DIR}/key-name"
  fi
  if [ -n "${TARGET_DIR:-}" ] && [ -f "${TARGET_DIR}/luks-device" ]; then
    rm -f "${TARGET_DIR}/luks-device"
  fi
  # Restore original VM config if we modified it
  if [ "$CONFIG_MODIFIED" = "true" ]; then
    echo "[migrate-vm-luks] Restoring original VM config..." >&2
    qm set "$VMID" "--${BOOT_SLOT}" "$DISK_SPEC" 2>/dev/null || true
    qm set "$VMID" --delete hookscript 2>/dev/null || true
  fi
}
trap cleanup_on_error ERR

# ── Check VM state ─────────────────────────────────────────────────
if qm status "$VMID" 2>/dev/null | grep -q running; then
  fail "VM ${VMID} is running. Shut it down first."
fi

# Close any stale mapper from a previous failed attempt
if cryptsetup status "$MAPPER" >/dev/null 2>&1; then
  log "Closing stale mapper ${MAPPER} from previous attempt..."
  cryptsetup close "$MAPPER" 2>/dev/null || fail "Cannot close existing mapper ${MAPPER} — is it still in use?"
fi

# ── Detect boot disk ──────────────────────────────────────────────
CONFIG="$(qm config "$VMID")"
echo "$CONFIG" > "${STATE_DIR}/qm-config.before.txt"

BOOT_SLOT=""
for s in scsi0 virtio0 sata0 ide0 scsi1 virtio1 sata1; do
  if echo "$CONFIG" | grep -q "^${s}: "; then
    BOOT_SLOT="$s"
    break
  fi
done
[ -n "$BOOT_SLOT" ] || fail "Could not find a disk slot (scsi0/virtio0/sata0/ide0)."

DISK_SPEC="$(echo "$CONFIG" | awk -F': ' -v s="$BOOT_SLOT" '$1==s {print $2}')"
DISK_VOLID="$(echo "$DISK_SPEC" | cut -d',' -f1)"
echo "$BOOT_SLOT" > "${STATE_DIR}/boot-slot.txt"
echo "$DISK_SPEC" > "${STATE_DIR}/disk-spec.txt"
echo "$DISK_VOLID" > "${STATE_DIR}/disk-volid.txt"

log "Detected boot disk slot: ${BOOT_SLOT}"
log "Detected disk volume: ${DISK_VOLID}"

# ── Detect and clean up leftover partial encryption ────────────────
# If the VM config still points to a LUKS mapper device or the metadata
# directory has stale artifacts from a previous encryption attempt,
# clean everything up and continue with a fresh encryption.
META_DIR="${LUKS_IMAGE_DIR}/${VMID}"

if [[ "$DISK_VOLID" == /dev/mapper/*_crypt ]]; then
  log "WARNING: VM config points to LUKS mapper device: ${DISK_VOLID}"
  log "Cleaning up leftover encryption state from a previous attempt..."

  # Close the mapper if it's still open
  cryptsetup close "$MAPPER" 2>/dev/null || true

  # Try to find the original disk spec from the saved migration state
  ORIG_SPEC=""
  if [ -f "${STATE_DIR}/disk-spec.txt" ]; then
    SAVED_SPEC="$(cat "${STATE_DIR}/disk-spec.txt")"
    SAVED_VOLID="$(echo "$SAVED_SPEC" | cut -d',' -f1)"
    if [[ "$SAVED_VOLID" != /dev/mapper/* ]]; then
      ORIG_SPEC="$SAVED_SPEC"
      log "Found original disk spec in migration state: ${ORIG_SPEC}"
    fi
  fi
  # Fallback: check the pre-migration config snapshot
  if [ -z "$ORIG_SPEC" ] && [ -f "${STATE_DIR}/qm-config.before.txt" ]; then
    SAVED_SPEC="$(grep "^${BOOT_SLOT}:" "${STATE_DIR}/qm-config.before.txt" | head -1 | sed 's/^[^:]*: //')"
    SAVED_VOLID="$(echo "$SAVED_SPEC" | cut -d',' -f1)"
    if [[ "$SAVED_VOLID" != /dev/mapper/* ]] && [ -n "$SAVED_VOLID" ]; then
      ORIG_SPEC="$SAVED_SPEC"
      log "Found original disk spec in config snapshot: ${ORIG_SPEC}"
    fi
  fi

  if [ -z "$ORIG_SPEC" ]; then
    fail "VM config points to LUKS mapper ${DISK_VOLID} but cannot determine the original disk.
  Manual fix required:
    1. Identify the original disk volume (check 'pvesm list <storage>' or 'zfs list -t volume')
    2. Restore VM config: qm set ${VMID} --${BOOT_SLOT} <original-volid>,size=<size>
    3. Remove hookscript: qm set ${VMID} --delete hookscript
    4. Remove stale metadata: rm -rf ${META_DIR}/
    5. Close mapper if open: cryptsetup close ${MAPPER}
    6. Re-run: migrate-vm-luks ${VMID}"
  fi

  # Remove the stale LUKS target volume if it still exists
  if [ -f "${META_DIR}/luks-device" ]; then
    STALE_DEV="$(cat "${META_DIR}/luks-device")"
    if [ -n "$STALE_DEV" ]; then
      if [[ "$STALE_DEV" == /dev/rbd* ]] && [ -b "$STALE_DEV" ]; then
        log "Unmapping stale RBD device: ${STALE_DEV}"
        rbd unmap "$STALE_DEV" 2>/dev/null || true
      fi
      if [ -f "${META_DIR}/rbd-spec" ]; then
        STALE_RBD_SPEC="$(cat "${META_DIR}/rbd-spec")"
        log "Removing stale RBD image via pvesm free..."
        # Determine storage name from conf path or storage.cfg
        STALE_STORAGE=""
        if [ -f "${META_DIR}/rbd-conf" ]; then
          STALE_CONF="$(cat "${META_DIR}/rbd-conf")"
          STALE_STORAGE="$(basename "$STALE_CONF" .conf)"
        fi
        if [ -n "$STALE_STORAGE" ] && [ -n "$STALE_RBD_SPEC" ]; then
          STALE_RBD_IMAGE="${STALE_RBD_SPEC##*/}"
          pvesm free "${STALE_STORAGE}:${STALE_RBD_IMAGE}" 2>/dev/null || true
        fi
      elif [[ "$STALE_DEV" == /dev/zvol/* ]]; then
        STALE_ZVOL="${STALE_DEV#/dev/zvol/}"
        log "Destroying stale ZFS zvol: ${STALE_ZVOL}"
        zfs destroy -f "$STALE_ZVOL" 2>/dev/null || true
      elif [[ "$STALE_DEV" == /dev/* ]] && [ -e "$STALE_DEV" ]; then
        log "Removing stale LVM volume: ${STALE_DEV}"
        lvremove -f "$STALE_DEV" 2>/dev/null || true
      fi
    fi
  fi
  # Also check for stale .luks file
  STALE_LUKS_FILE="${META_DIR}/${LUKS_IMAGE_TEMPLATE//\{vmid\}/$VMID}"
  if [ -f "$STALE_LUKS_FILE" ]; then
    log "Removing stale LUKS file: ${STALE_LUKS_FILE}"
    rm -f "$STALE_LUKS_FILE"
  fi

  # Clean up all metadata files
  rm -f "${META_DIR}/luks-device" "${META_DIR}/key-name" \
        "${META_DIR}/rbd-spec" "${META_DIR}/rbd-conf" \
        "${META_DIR}/rbd-id" "${META_DIR}/rbd-keyring" \
        "${META_DIR}/rbd-monhost" 2>/dev/null || true

  # Restore VM config to original disk
  log "Restoring VM config: ${BOOT_SLOT} → ${ORIG_SPEC}"
  qm set "$VMID" "--${BOOT_SLOT}" "$ORIG_SPEC"
  qm set "$VMID" --delete hookscript 2>/dev/null || true

  # Re-read the now-restored config
  DISK_SPEC="$ORIG_SPEC"
  DISK_VOLID="$(echo "$DISK_SPEC" | cut -d',' -f1)"
  echo "$DISK_SPEC" > "${STATE_DIR}/disk-spec.txt"
  echo "$DISK_VOLID" > "${STATE_DIR}/disk-volid.txt"
  log "Cleanup complete — continuing with fresh encryption of ${DISK_VOLID}"

elif [ -d "$META_DIR" ]; then
  # Config is clean but metadata artifacts may exist from a previous partial run
  HAS_STALE=false
  for f in luks-device key-name rbd-spec; do
    [ -f "${META_DIR}/$f" ] && HAS_STALE=true && break
  done
  STALE_LUKS_FILE="${META_DIR}/${LUKS_IMAGE_TEMPLATE//\{vmid\}/$VMID}"
  [ -f "$STALE_LUKS_FILE" ] && HAS_STALE=true

  if [ "$HAS_STALE" = "true" ]; then
    log "WARNING: Found leftover encryption metadata in ${META_DIR}/"
    log "Cleaning up stale artifacts before proceeding..."

    # Remove stale LUKS target
    if [ -f "${META_DIR}/luks-device" ]; then
      STALE_DEV="$(cat "${META_DIR}/luks-device")"
      if [ -n "$STALE_DEV" ]; then
        if [[ "$STALE_DEV" == /dev/rbd* ]] && [ -b "$STALE_DEV" ]; then
          rbd unmap "$STALE_DEV" 2>/dev/null || true
        fi
        if [ -f "${META_DIR}/rbd-spec" ]; then
          STALE_RBD_SPEC="$(cat "${META_DIR}/rbd-spec")"
          STALE_STORAGE=""
          if [ -f "${META_DIR}/rbd-conf" ]; then
            STALE_CONF="$(cat "${META_DIR}/rbd-conf")"
            STALE_STORAGE="$(basename "$STALE_CONF" .conf)"
          fi
          if [ -n "$STALE_STORAGE" ] && [ -n "$STALE_RBD_SPEC" ]; then
            STALE_RBD_IMAGE="${STALE_RBD_SPEC##*/}"
            pvesm free "${STALE_STORAGE}:${STALE_RBD_IMAGE}" 2>/dev/null || true
          fi
        elif [[ "$STALE_DEV" == /dev/zvol/* ]]; then
          STALE_ZVOL="${STALE_DEV#/dev/zvol/}"
          zfs destroy -f "$STALE_ZVOL" 2>/dev/null || true
        elif [[ "$STALE_DEV" == /dev/* ]] && [ -e "$STALE_DEV" ]; then
          lvremove -f "$STALE_DEV" 2>/dev/null || true
        fi
      fi
    fi
    [ -f "$STALE_LUKS_FILE" ] && rm -f "$STALE_LUKS_FILE"

    rm -f "${META_DIR}/luks-device" "${META_DIR}/key-name" \
          "${META_DIR}/rbd-spec" "${META_DIR}/rbd-conf" \
          "${META_DIR}/rbd-id" "${META_DIR}/rbd-keyring" \
          "${META_DIR}/rbd-monhost" 2>/dev/null || true
    log "Stale artifacts cleaned up — continuing with fresh encryption"
  fi
fi

SRC_PATH="$(resolve_source_path "$DISK_VOLID")"
echo "$SRC_PATH" > "${STATE_DIR}/source-path.txt"

# ── Determine LUKS image location from source disk storage ─────────
# Place the LUKS container on the same storage as the source disk.
# For LVM-thin: create a thin LV (native block device, no file on boot drive)
# For ZFS: create a zvol in the same pool
# For RBD: create a new RBD image in the same Ceph pool
# For directory-based: create a .luks file alongside the source
USE_LVM_THIN=false
LVM_VG=""
LVM_THINPOOL=""
LVM_LV_NAME="vm-${VMID}-disk-qxvault"
RBD_MODE=false
RBD_POOL=""
RBD_SRC_IMAGE=""
RBD_TARGET_IMAGE=""
RBD_CONF=""
RBD_ID=""
RBD_KEYRING=""

if [[ "$DISK_VOLID" != /* ]]; then
  SRC_STORAGE="${DISK_VOLID%%:*}"

  # Detect storage type from pvesm
  STORAGE_TYPE="$(pvesm status 2>/dev/null | awk -v s="$SRC_STORAGE" '$1==s {print $2}')" || true

  if [ "$STORAGE_TYPE" = "lvmthin" ]; then
    # Parse VG and thinpool from /etc/pve/storage.cfg
    # Format:
    #   lvmthin: local-lvm
    #       thinpool data
    #       vgname pve
    LVM_VG="$(awk -v s="$SRC_STORAGE" '
      /^lvmthin:/ && $2 == s { found=1; next }
      found && /^\s+vgname/ { print $2; exit }
      /^[a-z]/ && found { exit }
    ' /etc/pve/storage.cfg)" || true

    LVM_THINPOOL="$(awk -v s="$SRC_STORAGE" '
      /^lvmthin:/ && $2 == s { found=1; next }
      found && /^\s+thinpool/ { print $2; exit }
      /^[a-z]/ && found { exit }
    ' /etc/pve/storage.cfg)" || true

    if [ -n "$LVM_VG" ] && [ -n "$LVM_THINPOOL" ]; then
      USE_LVM_THIN=true
      log "Source is LVM-thin: VG=${LVM_VG}, thinpool=${LVM_THINPOOL}"
    else
      log "WARNING: Could not parse LVM-thin VG/thinpool from storage.cfg, falling back to file-based"
    fi

  elif [ "$STORAGE_TYPE" = "zfspool" ]; then
    # ZFS zvol storage — derive pool name from the resolved source path
    # Source path is always /dev/zvol/{pool}/{volname}
    if [[ "$SRC_PATH" == /dev/zvol/* ]]; then
      ZFS_POOL="$(echo "$SRC_PATH" | cut -d'/' -f4)"
    fi
    # Fallback: try parsing from storage.cfg
    if [ -z "${ZFS_POOL:-}" ]; then
      ZFS_POOL="$(awk -v s="$SRC_STORAGE" '
        /^zfspool:/ && $2 == s { found=1; next }
        found && /^\s+pool/ { print $2; exit }
        /^[a-z]/ && found { exit }
      ' /etc/pve/storage.cfg)" || true
    fi

    if [ -n "${ZFS_POOL:-}" ]; then
      USE_LVM_THIN=true  # Reuse the block-device path (same logic applies)
      LVM_VG=""  # Not used for ZFS
      LVM_THINPOOL=""  # Not used for ZFS
      TARGET_DIR="${LUKS_IMAGE_DIR}/${VMID}"
      TARGET_FILE="/dev/zvol/${ZFS_POOL}/${LVM_LV_NAME}"
      log "Source is ZFS zvol: pool=${ZFS_POOL}"
    else
      log "WARNING: Could not determine ZFS pool name, falling back to file-based"
    fi

  elif [ "$STORAGE_TYPE" = "rbd" ]; then
    # Ceph RBD storage — pvesm path returns a QEMU protocol string, not a local path.
    # We map the RBD image locally to get a block device for the copy.
    RBD_POOL="$(awk -v s="$SRC_STORAGE" '
      /^rbd:/ && $2 == s { found=1; next }
      found && /^\s+pool/ { print $2; exit }
      /^[a-z]/ && found { exit }
    ' /etc/pve/storage.cfg)" || true
    # Fallback: extract pool from pvesm path output (format: rbd:pool/image:...)
    if [ -z "$RBD_POOL" ] && [[ "$SRC_PATH" == rbd:* ]]; then
      RBD_POOL="$(echo "$SRC_PATH" | sed 's|^rbd:||; s|:.*||; s|/.*||')"
    fi
    [ -z "$RBD_POOL" ] && fail "Cannot determine Ceph pool for storage '${SRC_STORAGE}'"

    RBD_CONF="/etc/pve/priv/ceph/${SRC_STORAGE}.conf"
    RBD_KEYRING="/etc/pve/priv/ceph/${SRC_STORAGE}.keyring"
    RBD_ID="$(awk -v s="$SRC_STORAGE" '
      /^rbd:/ && $2 == s { found=1; next }
      found && /^\s+username/ { print $2; exit }
      /^[a-z]/ && found { exit }
    ' /etc/pve/storage.cfg)" || true
    [ -z "$RBD_ID" ] && RBD_ID="admin"

    # Extract monitor hosts from storage.cfg (space- or comma-separated)
    RBD_MONHOST="$(awk -v s="$SRC_STORAGE" '
      /^rbd:/ && $2 == s { found=1; next }
      found && /^\s+monhost/ { $1=""; gsub(/^[[:space:]]+/, ""); print; exit }
      /^[a-z]/ && found { exit }
    ' /etc/pve/storage.cfg)" || true

    # Parse the full RBD spec from pvesm path output.
    # pvesm path returns: rbd:pool/[namespace/]image:conf=...:mon_host=...
    # We extract the pool/[namespace/]image part for use with rbd commands.
    RBD_FULL_SPEC="${SRC_PATH#rbd:}"
    RBD_FULL_SPEC="${RBD_FULL_SPEC%%:*}"
    # RBD_FULL_SPEC is now e.g. "rbdfast/rbdfast/vm-250-disk-0" (with namespace)
    # or "rbdfast/vm-250-disk-0" (without namespace)

    RBD_SRC_SPEC="$RBD_FULL_SPEC"
    RBD_SRC_IMAGE="${DISK_VOLID#*:}"
    # Build target spec: same pool/namespace prefix, different image name
    RBD_TARGET_IMAGE="${LVM_LV_NAME}"
    RBD_TARGET_SPEC="${RBD_FULL_SPEC%/*}/${RBD_TARGET_IMAGE}"

    log "Source is Ceph RBD: pool=${RBD_POOL}, spec=${RBD_SRC_SPEC}"
    log "Mapping source RBD image..."

    RBD_MAP_ARGS=()
    [ -f "$RBD_CONF" ] && RBD_MAP_ARGS+=(--conf "$RBD_CONF")
    [ -f "$RBD_KEYRING" ] && RBD_MAP_ARGS+=(--keyring "$RBD_KEYRING")
    [ -n "$RBD_ID" ] && RBD_MAP_ARGS+=(--id "$RBD_ID")
    # Proxmox stores monitors in storage.cfg, not always in ceph.conf
    if [ -n "$RBD_MONHOST" ]; then
      # Convert space-separated to comma-separated for -m flag
      RBD_MAP_ARGS+=(-m "$(echo "$RBD_MONHOST" | tr ' ' ',')")
    fi

    SRC_PATH="$(rbd map "${RBD_SRC_SPEC}" "${RBD_MAP_ARGS[@]}")" \
      || fail "Failed to map source RBD image: ${RBD_SRC_SPEC}"
    log "Source mapped to: ${SRC_PATH}"
    echo "$SRC_PATH" > "${STATE_DIR}/source-path.txt"
    echo "${RBD_SRC_SPEC}" > "${STATE_DIR}/rbd-source-spec.txt"

    RBD_MODE=true
    USE_LVM_THIN=true  # Block-device mode
    LVM_VG=""
    LVM_THINPOOL=""
    TARGET_DIR="${LUKS_IMAGE_DIR}/${VMID}"
    TARGET_FILE="/dev/rbd/${RBD_TARGET_SPEC}"
    log "LUKS target will be RBD image: ${RBD_TARGET_SPEC}"
  fi

  if [ "$USE_LVM_THIN" = "false" ]; then
    STORAGE_PATH="$(pvesm path "${SRC_STORAGE}:." 2>/dev/null | sed 's|/\.$||')" || true
    if [ -z "$STORAGE_PATH" ] || [ ! -d "$STORAGE_PATH" ]; then
      STORAGE_PATH="$(dirname "$(dirname "$SRC_PATH")")"
    fi
  fi
else
  # Raw filesystem path — use the parent directory's parent (storage root)
  STORAGE_PATH="$(dirname "$(dirname "$SRC_PATH")")"
fi

if [ "$USE_LVM_THIN" = "true" ] && [ -z "${TARGET_FILE:-}" ]; then
  # LVM-thin: LUKS goes on a thin LV; metadata directory stays at LUKS_IMAGE_DIR
  TARGET_DIR="${LUKS_IMAGE_DIR}/${VMID}"
  TARGET_FILE="/dev/${LVM_VG}/${LVM_LV_NAME}"
  log "LUKS will be a thin LV: ${TARGET_FILE}"
elif [ "$USE_LVM_THIN" != "true" ]; then
  # File-based path
  if [[ "${STORAGE_PATH:-}" == /dev* ]] || [ ! -d "${STORAGE_PATH:-}" ]; then
    # Source is block-based but we couldn't auto-detect a native block target.
    # Refuse to silently dump a 32GB+ file on the boot drive.
    if [[ "$SRC_PATH" == /dev/* ]]; then
      fail "Source disk is block-based ($SRC_PATH) but no native block target could be determined. Set LUKS_IMAGE_DIR in /etc/qxvault/qxvault.conf to a directory with sufficient space."
    fi
    STORAGE_PATH="$LUKS_IMAGE_DIR"
    log "Using LUKS_IMAGE_DIR=${LUKS_IMAGE_DIR}"
  fi

  TARGET_DIR="${STORAGE_PATH}/${VMID}"
  TARGET_FILE="${TARGET_DIR}/${LUKS_IMAGE_TEMPLATE//\{vmid\}/$VMID}"
  log "LUKS image will be stored at: ${TARGET_DIR}/"

  # Update LUKS_IMAGE_DIR in config if it differs from what's configured
  if [ "$STORAGE_PATH" != "$LUKS_IMAGE_DIR" ]; then
    log "Updating LUKS_IMAGE_DIR in config to: ${STORAGE_PATH}"
    if [ -f "$CONF_FILE" ]; then
      sed -i "s|^LUKS_IMAGE_DIR=.*|LUKS_IMAGE_DIR=\"${STORAGE_PATH}\"|" "$CONF_FILE"
    fi
  fi
fi

# Verify source path exists (after potential RBD mapping)
[ -e "$SRC_PATH" ] || fail "Resolved source path does not exist: ${SRC_PATH}"

SRC_INFO="$(qemu-img info --output=json "$SRC_PATH")"
printf '%s\n' "$SRC_INFO" > "${STATE_DIR}/source-qemu-img-info.json"

VIRTUAL_SIZE="$(printf '%s\n' "$SRC_INFO" | jq -r '.["virtual-size"]')"
if [ -z "$VIRTUAL_SIZE" ] || [ "$VIRTUAL_SIZE" = "null" ]; then
  fail "Could not read source virtual size."
fi

SRC_FORMAT="$(printf '%s\n' "$SRC_INFO" | jq -r '.format')"
if [ -z "$SRC_FORMAT" ] || [ "$SRC_FORMAT" = "null" ]; then
  fail "Could not read source image format."
fi

case "$SRC_FORMAT" in
  raw|qcow2|qcow|vmdk|vdi|vhdx) ;;
  *) fail "Unsupported source format: ${SRC_FORMAT}" ;;
esac

TARGET_SIZE=$((VIRTUAL_SIZE + LUKS_OVERHEAD_BYTES))
echo "$VIRTUAL_SIZE" > "${STATE_DIR}/source-virtual-size.txt"
echo "$TARGET_SIZE" > "${STATE_DIR}/target-file-size.txt"
echo "$SRC_FORMAT" > "${STATE_DIR}/source-format.txt"

log "Source path: ${SRC_PATH}"
log "Source format: ${SRC_FORMAT}"
log "Virtual size: ${VIRTUAL_SIZE} bytes"
log "Target file size: ${TARGET_SIZE} bytes"
log "Target LUKS: ${TARGET_FILE}"

# ── Dry run ────────────────────────────────────────────────────────
if [ "$DRY_RUN" = "true" ]; then
  KEY_NAME="$(build_key_name "$VMID")"
  log "DRY RUN — no changes made. Summary:"
  log "  VM ${VMID}: ${BOOT_SLOT} = ${DISK_VOLID}"
  log "  Source: ${SRC_PATH} (${SRC_FORMAT}, ${VIRTUAL_SIZE} bytes)"
  log "  Target: ${TARGET_FILE} (${TARGET_SIZE} bytes)"
  log "  Mode: $([ "$USE_LVM_THIN" = "true" ] && echo "LVM-thin LV" || echo "file-based")"
  log "  Hookscript: ${HOOKSCRIPT}"
  log "  Key will be provisioned in QxVault transit engine (key name: ${KEY_NAME})"
  exit 0
fi

# ── Create metadata directory ─────────────────────────────────────
install -d -m 700 "$TARGET_DIR"

# ── Resolve unique key name (collision-safe) ──────────────────────
KEY_NAME="$(resolve_new_key_name "$VMID" "$TARGET_DIR")"
log "Resolved key name: ${KEY_NAME}"

# ── Operation lock file (for UI progress tracking) ───────────────
OP_LOCK="/run/qxvault-op-${VMID}.json"
echo "{\"op\":\"encrypt\",\"vmid\":\"${VMID}\",\"startTime\":$(date +%s)}" > "$OP_LOCK"
trap 'rm -f "$OP_LOCK"' EXIT

if [ -e "$TARGET_FILE" ]; then
  fail "Target already exists: ${TARGET_FILE}"
fi
# For RBD, also check if the image already exists in the pool
if [ "${RBD_MODE:-}" = "true" ]; then
  if rbd info "${RBD_TARGET_SPEC}" "${RBD_MAP_ARGS[@]}" >/dev/null 2>&1; then
    fail "Target RBD image already exists: ${RBD_TARGET_SPEC}"
  fi
fi

# ── Disk space check ──────────────────────────────────────────────
if [ "$USE_LVM_THIN" = "true" ]; then
  if [ -n "$LVM_VG" ] && [ -n "$LVM_THINPOOL" ]; then
    # LVM-thin: check thin pool usage
    POOL_USED_PCT="$(lvs --noheadings -o data_percent "${LVM_VG}/${LVM_THINPOOL}" 2>/dev/null | tr -d ' ')" || true
    log "Thin pool ${LVM_VG}/${LVM_THINPOOL}: ${POOL_USED_PCT:-?}% used"
    if [ -n "$POOL_USED_PCT" ]; then
      POOL_PCT_INT="${POOL_USED_PCT%%.*}"
      if [ "${POOL_PCT_INT:-0}" -gt 80 ]; then
        log "WARNING: Thin pool is ${POOL_USED_PCT}% full — migration will still proceed"
      fi
    fi
  elif [ "${RBD_MODE:-}" = "true" ]; then
    # RBD: check pool free space via ceph df (best-effort)
    RBD_AVAIL="$(ceph df --format=json 2>/dev/null | jq -r --arg p "$RBD_POOL" '.pools[] | select(.name==$p) | .stats.max_avail')" || true
    if [ -n "$RBD_AVAIL" ] && [ "$RBD_AVAIL" != "null" ]; then
      log "Ceph pool ${RBD_POOL}: $(numfmt --to=iec "$RBD_AVAIL") available"
      if [ "$TARGET_SIZE" -gt "$RBD_AVAIL" ]; then
        fail "Insufficient Ceph pool space: need $(numfmt --to=iec "$TARGET_SIZE"), have $(numfmt --to=iec "$RBD_AVAIL")"
      fi
    fi
  elif [ -n "${ZFS_POOL:-}" ]; then
    # ZFS: check pool free space
    ZFS_AVAIL="$(zfs get -Hp -o value available "${ZFS_POOL}" 2>/dev/null)" || true
    if [ -n "$ZFS_AVAIL" ]; then
      log "ZFS pool ${ZFS_POOL}: $(numfmt --to=iec "$ZFS_AVAIL") available"
      if [ "$TARGET_SIZE" -gt "$ZFS_AVAIL" ]; then
        fail "Insufficient ZFS pool space: need $(numfmt --to=iec "$TARGET_SIZE"), have $(numfmt --to=iec "$ZFS_AVAIL")"
      fi
    fi
  fi
else
  AVAIL_BYTES=$(df --output=avail -B1 "$TARGET_DIR" | tail -1 | tr -d ' ')
  REQUIRED_BYTES=$(( TARGET_SIZE + TARGET_SIZE / 10 ))  # 10% safety margin
  SRC_FS="$(df --output=target "$SRC_PATH" 2>/dev/null | tail -1 | tr -d ' ')"
  TGT_FS="$(df --output=target "$TARGET_DIR" 2>/dev/null | tail -1 | tr -d ' ')"
  if [ "$SRC_FS" = "$TGT_FS" ]; then
    log "Source and target are on the same filesystem (${TGT_FS})"
    log "Available: $(numfmt --to=iec "$AVAIL_BYTES"), need: $(numfmt --to=iec "$REQUIRED_BYTES") for LUKS file"
  fi
  if [ "$REQUIRED_BYTES" -gt "$AVAIL_BYTES" ]; then
    fail "Insufficient disk space: need $(numfmt --to=iec "$REQUIRED_BYTES") (incl. 10% margin), have $(numfmt --to=iec "$AVAIL_BYTES") on ${TGT_FS}"
  fi
fi

# Verify key retrieval works
KEY_BYTES="$(/usr/libexec/vault-dmkey get-key --name "$KEY_NAME" | base64 -d | wc -c | tr -d ' ')"
[ "$KEY_BYTES" = "32" ] || fail "Vault key ${KEY_NAME} is not 32 bytes after decode."

# ── Create LUKS target ────────────────────────────────────────────
if [ "$USE_LVM_THIN" = "true" ]; then
  # Convert bytes to megabytes (round up)
  TARGET_SIZE_MB=$(( (TARGET_SIZE + 1048575) / 1048576 ))

  if [ -n "$LVM_VG" ] && [ -n "$LVM_THINPOOL" ]; then
    log "Creating thin LV: ${LVM_VG}/${LVM_LV_NAME} (${TARGET_SIZE_MB}M)"
    lvcreate -T "${LVM_VG}/${LVM_THINPOOL}" -n "$LVM_LV_NAME" -V "${TARGET_SIZE_MB}m"
  elif [ "${RBD_MODE}" = "true" ]; then
    log "Creating RBD image: ${RBD_TARGET_SPEC} (${TARGET_SIZE_MB}M)"
    RBD_CREATE_ARGS=()
    [ -f "$RBD_CONF" ] && RBD_CREATE_ARGS+=(--conf "$RBD_CONF")
    [ -f "$RBD_KEYRING" ] && RBD_CREATE_ARGS+=(--keyring "$RBD_KEYRING")
    [ -n "$RBD_ID" ] && RBD_CREATE_ARGS+=(--id "$RBD_ID")
    if [ -n "$RBD_MONHOST" ]; then
      RBD_CREATE_ARGS+=(-m "$(echo "$RBD_MONHOST" | tr ' ' ',')")
    fi
    rbd create "${RBD_TARGET_SPEC}" --size "${TARGET_SIZE_MB}" "${RBD_CREATE_ARGS[@]}" \
      || fail "Failed to create RBD image: ${RBD_TARGET_SPEC}"
    log "Mapping target RBD image..."
    RBD_TARGET_DEV="$(rbd map "${RBD_TARGET_SPEC}" "${RBD_MAP_ARGS[@]}")" \
      || fail "Failed to map target RBD image: ${RBD_TARGET_SPEC}"
    log "Target mapped to: ${RBD_TARGET_DEV}"
    # Use the actual mapped device for LUKS operations
    TARGET_FILE="$RBD_TARGET_DEV"
    # Store RBD metadata for hookscript
    printf '%s\n' "$RBD_TARGET_SPEC" > "${TARGET_DIR}/rbd-spec"
    printf '%s\n' "$RBD_CONF" > "${TARGET_DIR}/rbd-conf"
    printf '%s\n' "$RBD_ID" > "${TARGET_DIR}/rbd-id"
    printf '%s\n' "$RBD_KEYRING" > "${TARGET_DIR}/rbd-keyring"
    if [ -n "$RBD_MONHOST" ]; then
      printf '%s\n' "$RBD_MONHOST" > "${TARGET_DIR}/rbd-monhost"
    fi
    chmod 600 "${TARGET_DIR}"/rbd-*
  elif [ -n "${ZFS_POOL:-}" ]; then
    log "Creating ZFS zvol: ${ZFS_POOL}/${LVM_LV_NAME} (${TARGET_SIZE_MB}M)"
    zfs create -V "${TARGET_SIZE_MB}m" "${ZFS_POOL}/${LVM_LV_NAME}"
    # Wait for device node to appear
    ZVOL_RETRIES=0
    while [ ! -e "$TARGET_FILE" ] && [ $ZVOL_RETRIES -lt 10 ]; do
      sleep 0.5
      ZVOL_RETRIES=$((ZVOL_RETRIES + 1))
    done
    [ -e "$TARGET_FILE" ] || fail "ZFS zvol created but device not found at ${TARGET_FILE}"
  fi
  # Store the device path for other scripts to find
  printf '%s\n' "$TARGET_FILE" > "${TARGET_DIR}/luks-device"
  chmod 600 "${TARGET_DIR}/luks-device"
else
  log "Creating target file: ${TARGET_FILE}"
  truncate -s "$TARGET_SIZE" "$TARGET_FILE"
  chmod 600 "$TARGET_FILE"
fi

log "Formatting target as LUKS2"
/usr/libexec/vault-dmkey get-key --name "$KEY_NAME" \
  | base64 -d \
  | cryptsetup luksFormat --type luks2 --key-file=- "$TARGET_FILE"

log "Opening mapper: ${MAPPER}"
/usr/libexec/vault-dmkey get-key --name "$KEY_NAME" \
  | base64 -d \
  | cryptsetup open --type luks --key-file=- "$TARGET_FILE" "$MAPPER"

MAPPER_SIZE="$(blockdev --getsize64 "/dev/mapper/${MAPPER}")"
echo "$MAPPER_SIZE" > "${STATE_DIR}/mapper-size.txt"
log "Mapped device size: ${MAPPER_SIZE} bytes"

if [ "$MAPPER_SIZE" -lt "$VIRTUAL_SIZE" ]; then
  cryptsetup close "$MAPPER"
  fail "Mapped device is smaller than source image."
fi

log "Copying source disk into /dev/mapper/${MAPPER} ..."
qemu-img convert -t none -p -n -f "$SRC_FORMAT" -O raw "$SRC_PATH" "/dev/mapper/${MAPPER}"

log "Verifying copied data..."
# Verify the written data is readable and has correct size
WRITTEN_SIZE="$(blockdev --getsize64 "/dev/mapper/${MAPPER}")"
if [ "$WRITTEN_SIZE" -lt "$VIRTUAL_SIZE" ]; then
  fail "Post-copy verification failed: written ${WRITTEN_SIZE} bytes but expected at least ${VIRTUAL_SIZE}"
fi
log "Post-copy verification passed (${WRITTEN_SIZE} bytes on mapper)"

log "Configuring VM for encrypted boot..."

log "Verifying Proxmox snippets storage"
install -d -m 755 /var/lib/vz/snippets

if [ ! -f /var/lib/vz/snippets/qxvault-luks-hook.sh ]; then
  fail "Missing /var/lib/vz/snippets/qxvault-luks-hook.sh — is proxmox-kms-bridge installed?"
fi

log "Attaching hookscript..."
CONFIG_MODIFIED=true
qm set "$VMID" --hookscript "$HOOKSCRIPT"

log "Switching VM disk on ${BOOT_SLOT} to /dev/mapper/${MAPPER}..."
qm set "$VMID" "--${BOOT_SLOT}" "/dev/mapper/${MAPPER}"

log "Saving rollback state..."
NEWCFG="$(qm config "$VMID")"
echo "$NEWCFG" > "${STATE_DIR}/qm-config.after.txt"

log "=========================================="
log "  Migration complete for VMID ${VMID}"
log "=========================================="
log "  LUKS target:    ${TARGET_FILE}$([ "$USE_LVM_THIN" = "true" ] && echo " (LVM-thin LV)")"
log "  Mapper device:  /dev/mapper/${MAPPER}"
log "  Hookscript:     ${HOOKSCRIPT}"
log "  Key stored in:  QxVault transit engine (key name: ${KEY_NAME})"
log "  Rollback state: ${STATE_DIR}/"
log ""
log "  Start the VM:   qm start ${VMID}"
log "=========================================="

# Lock the LUKS image by closing the mapper
log "Locking LUKS image (closing mapper ${MAPPER})..."
if cryptsetup close "$MAPPER" 2>&1; then
    log "✓ LUKS image is now locked"
else
    log "WARNING: Failed to close mapper, but encryption is complete"
fi

# Unmap RBD images if Ceph storage was used
if [ "$RBD_MODE" = "true" ]; then
    log "Unmapping target RBD image..."
    rbd unmap "$TARGET_FILE" 2>/dev/null \
        && log "✓ Unmapped target RBD: ${TARGET_FILE}" \
        || log "WARNING: Failed to unmap target RBD: ${TARGET_FILE}"
    log "Unmapping source RBD image..."
    rbd unmap "$SRC_PATH" 2>/dev/null \
        && log "✓ Unmapped source RBD: ${SRC_PATH}" \
        || log "WARNING: Failed to unmap source RBD: ${SRC_PATH}"
fi

# Remove the original disk image now that encryption succeeded
log "Removing original disk image..."
if [[ "$DISK_VOLID" == /* ]]; then
    # Raw filesystem path
    rm -f "$DISK_VOLID" && log "✓ Removed original image: ${DISK_VOLID}" \
        || log "WARNING: Failed to remove original image: ${DISK_VOLID}"
else
    # Proxmox storage volume — free via pvesm
    pvesm free "$DISK_VOLID" 2>/dev/null \
        && log "✓ Freed original volume: ${DISK_VOLID}" \
        || log "WARNING: Failed to free original volume: ${DISK_VOLID} (may need manual cleanup)"
fi

# Clean up rollback state since migration succeeded
rm -rf "$STATE_DIR"
log "✓ Cleaned up rollback state"

log "=========================================="
log "✓ Task Complete!"
log "VM ${VMID} is now encrypted and locked."
log "=========================================="

exit 0