#!/bin/bash
# proxmox-kms-bridge — CLI for 45Drives SecureVM for Proxmox
# Part of proxmox-kms-bridge
#
# Usage: proxmox-kms-bridge <command> [args]
#
# Commands:
#   status           Show OpenBao connectivity and encrypted VM summary
#   verify <VMID>    Test key retrieval for a VM without unlocking
#   revert <VMID>    Decrypt and revert a VM to its original unencrypted state
#   cleanup          Remove orphaned LUKS images for deleted VMs
set -euo pipefail

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

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

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

# ── Helpers ────────────────────────────────────────────────────────
log() { echo "[proxmox-kms-bridge] $*"; }
die() { echo "[proxmox-kms-bridge] ERROR: $*" >&2; exit 1; }

vault_curl() {
  local -a curl_args=( -s --connect-timeout "$OPENBAO_TIMEOUT" --max-time "$((OPENBAO_TIMEOUT * 2))" )
  if [ -n "$OPENBAO_CACERT" ]; then curl_args+=( --cacert "$OPENBAO_CACERT" ); fi
  if [ "$OPENBAO_SKIP_VERIFY" = "true" ]; then curl_args+=( -k ); fi
  curl "${curl_args[@]}" "$@"
}

usage() {
  cat <<'EOF'
45Drives SecureVM for Proxmox — Built by 45Drives in partnership with Crypto4A

Usage: proxmox-kms-bridge <command> [args]

Commands:
  status           Show QxVault connectivity and encrypted VM summary
  verify <VMID>    Test key retrieval for a VM (does not unlock)
  revert <VMID>    Decrypt and revert a VM to unencrypted state
  cleanup          Remove orphaned LUKS images for deleted VMs
  help             Show this help

EOF
  exit 0
}

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

# Find LUKS image for a VM — checks configured path first, then searches storage
find_luks_image() {
  local vmid="$1"
  local meta_dir="${LUKS_IMAGE_DIR}/${vmid}"

  # Check for RBD spec (Ceph block device — may not be mapped)
  if [ -f "${meta_dir}/rbd-spec" ]; then
    # Check if luks-device points to a currently mapped device
    if [ -f "${meta_dir}/luks-device" ]; then
      local dev
      dev="$(cat "${meta_dir}/luks-device")"
      if [ -b "$dev" ]; then
        printf '%s\n' "$dev"
        return 0
      fi
    fi
    # RBD not mapped — return the spec prefixed with rbd: so callers know
    printf 'rbd:%s\n' "$(cat "${meta_dir}/rbd-spec")"
    return 0
  fi

  # Check for LVM-thin / ZFS device pointer
  if [ -f "${meta_dir}/luks-device" ]; then
    local dev
    dev="$(cat "${meta_dir}/luks-device")"
    if [ -b "$dev" ]; then
      printf '%s\n' "$dev"
      return 0
    fi
  fi

  # File-based: check configured path
  local filename="${LUKS_IMAGE_TEMPLATE//\{vmid\}/$vmid}"
  local primary="${meta_dir}/${filename}"
  if [ -f "$primary" ]; then
    printf '%s\n' "$primary"
    return 0
  fi
  # Search common storage mount points
  local found
  found="$(timeout 3 find /var/lib/vz /mnt -maxdepth 4 -path "*/${vmid}/${filename}" -print -quit 2>/dev/null)" || true
  if [ -n "$found" ] && [ -f "$found" ]; then
    printf '%s\n' "$found"
    return 0
  fi
  # Not found — return the primary path anyway (caller checks existence)
  printf '%s\n' "$primary"
  return 1
}

# Read the stored key name for a VM.
# Checks for key-name file next to the LUKS image, falls back to bare VMID.
read_key_name() {
  local vmid="$1"
  local image_path="$2"
  local image_dir

  # For RBD, ZFS zvol, and LVM-thin: use the metadata directory
  if [[ "$image_path" == rbd:* ]] || [[ "$image_path" == /dev/* ]]; then
    image_dir="${LUKS_IMAGE_DIR}/${vmid}"
  else
    image_dir="$(dirname "$image_path")"
  fi

  if [ -f "${image_dir}/key-name" ]; then
    cat "${image_dir}/key-name"
    return 0
  fi
  # Backward compat: no key-name file → use bare VMID
  printf '%s\n' "$vmid"
}

# ── status command ─────────────────────────────────────────────────
cmd_status() {
  echo "45Drives SecureVM for Proxmox — Status"
  echo "======================================="
  echo

  # OpenBao connectivity
  echo "QxVault Endpoint: ${OPENBAO_ADDR}"
  local health_code
  health_code=$(vault_curl -o /dev/null -w '%{http_code}' "${OPENBAO_ADDR}/v1/sys/health" 2>/dev/null) || true
  case "$health_code" in
    200) echo "QxVault Status:   Connected (unsealed)" ;;
    429) echo "QxVault Status:   Connected (standby)" ;;
    472) echo "QxVault Status:   Connected (DR secondary)" ;;
    473) echo "QxVault Status:   Connected (performance standby)" ;;
    501) echo "QxVault Status:   Connected (NOT INITIALIZED)" ;;
    503) echo "QxVault Status:   Connected (SEALED)" ;;
    *)   echo "QxVault Status:   UNREACHABLE (HTTP ${health_code:-000})" ;;
  esac

  # Config
  echo "Config File:      ${CONF_FILE}"
  if [ -f "$CONF_FILE" ]; then
    echo "Config Status:    Present"
  else
    echo "Config Status:    MISSING"
  fi
  echo

  # Scan for encrypted VMs
  echo "Encrypted VMs:"
  echo "──────────────────────────────────────────────────────────"
  printf "%-8s %-12s %-10s %s\n" "VMID" "LUKS State" "Hookscript" "Image"
  echo "──────────────────────────────────────────────────────────"

  local found=0
  local vmid
  # Use qm list for reliable VM enumeration across all storage types
  while IFS= read -r vmid; do
    [ -n "$vmid" ] || continue
    local conf_file="/etc/pve/qemu-server/${vmid}.conf"
    [ -f "$conf_file" ] || continue
    if grep -q "qxvault-luks-hook" "$conf_file" 2>/dev/null; then
      local mapper="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$vmid}"
      local image
      image="$(find_luks_image "$vmid")"
      local luks_state="closed"
      local hook_state="attached"

      if cryptsetup status "$mapper" >/dev/null 2>&1; then
        luks_state="OPEN"
      fi

      if [[ "$image" != rbd:* ]] && [ ! -f "$image" ] && [ ! -b "$image" ]; then
        luks_state="NO IMAGE"
      fi

      printf "%-8s %-12s %-10s %s\n" "$vmid" "$luks_state" "$hook_state" "$image"
      found=$((found + 1))
    fi
  done < <(qm list 2>/dev/null | awk 'NR>1 {print $1}')

  if [ "$found" -eq 0 ]; then
    echo "  (none found)"
  fi
  echo "──────────────────────────────────────────────────────────"
  echo "Total encrypted VMs: $found"

  # Scan for orphaned LUKS images (VMs that no longer exist)
  local orphan_count=0
  local orphan_list=""
  local luks_files=()
  while IFS= read -r -d '' f; do
    luks_files+=("$f")
  done < <(timeout 5 find "$LUKS_IMAGE_DIR" /var/lib/vz /mnt -maxdepth 4 -name 'vm-*-disk-qxvault.luks' -print0 2>/dev/null | sort -z -u)

  for luks_file in "${luks_files[@]}"; do
    local obase
    obase=$(basename "$luks_file")
    local ovmid
    ovmid=$(echo "$obase" | sed -n 's/^vm-\([0-9]\+\)-disk-qxvault\.luks$/\1/p')
    [ -n "$ovmid" ] || continue
    if ! qm status "$ovmid" >/dev/null 2>&1; then
      orphan_count=$((orphan_count + 1))
      local osize
      osize=$(stat -c%s "$luks_file" 2>/dev/null || echo 0)
      orphan_list="${orphan_list}$(printf '%-8s %s (%s)' "$ovmid" "$luks_file" "$(numfmt --to=iec "$osize")")\n"
    fi
  done

  echo
  if [ "$orphan_count" -gt 0 ]; then
    echo "Orphaned LUKS Images: $orphan_count"
    echo "──────────────────────────────────────────────────────────"
    printf "%-8s %s\n" "VMID" "Image (Size)"
    echo "──────────────────────────────────────────────────────────"
    echo -e "$orphan_list"
    echo "Run 'proxmox-kms-bridge cleanup' to remove."
  else
    echo "Orphaned LUKS Images: 0"
  fi
}

# ── verify command ─────────────────────────────────────────────────
cmd_verify() {
  local vmid="${1:-}"
  [ -n "$vmid" ] || die "Usage: proxmox-kms-bridge verify <VMID>"
  [[ "$vmid" =~ ^[0-9]+$ ]] || die "VMID must be numeric"

  echo "Verifying key retrieval for VMID ${vmid}..."

  # Check LUKS image first — if not encrypted, don't query vault (which auto-creates keys)
  local image is_rbd=false
  image="$(find_luks_image "$vmid")"

  # Handle RBD images (not mapped — find_luks_image returns rbd:pool/image)
  if [[ "$image" == rbd:* ]]; then
    is_rbd=true
    # For display, strip the rbd: prefix
    local rbd_spec="${image#rbd:}"
    image="rbd:${rbd_spec}"
  fi

  if [ "$is_rbd" = "false" ] && [ ! -f "$image" ] && [ ! -b "$image" ]; then
    echo "  LUKS image:    ${image} (NOT FOUND)"
    echo "  Status:        VM is not encrypted — no key to verify"
    echo "  Verification complete."
    return 0
  fi

  # 2FA status
  local tfa_conf="/etc/qxvault/2fa/${vmid}.conf"
  if [ -f "$tfa_conf" ]; then
    local tfa_mode=""
    # shellcheck source=/dev/null
    . "$tfa_conf"
    echo "  2FA:           ENABLED (mode: ${TFA_MODE:-unknown})"
    echo "  Auto-unlock:   BLOCKED (manual unlock required)"
  else
    echo "  2FA:           disabled"
    echo "  Auto-unlock:   enabled"
  fi

  # Test key retrieval (stderr shows debug info)
  local key_name
  key_name="$(read_key_name "$vmid" "$image")"
  echo "  Key name:      ${key_name}"
  local key_b64
  if key_b64=$(/usr/libexec/vault-dmkey get-key --name "$key_name" 2>&1); then
    if [ -n "$key_b64" ]; then
      local key_len
      key_len=$(echo "$key_b64" | base64 -d 2>/dev/null | wc -c)
      echo "  Key retrieval: OK"
      echo "  Key length:    ${key_len} bytes"
    else
      echo "  Key retrieval: OK (but could not parse output)"
    fi
  else
    echo "  Key retrieval: FAILED"
    # shellcheck disable=SC2001
    echo "$key_b64" | sed 's/^/    /'
    exit 1
  fi

  # LUKS image details
  echo "  LUKS image:    ${image} (exists)"
  if [ "$is_rbd" = "false" ]; then
    local luks_info
    if luks_info=$(cryptsetup luksDump "$image" 2>/dev/null | head -5); then
      # shellcheck disable=SC2001
      echo "$luks_info" | sed 's/^/    /'
    fi
  fi

  # Check mapper state
  local mapper="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$vmid}"
  if cryptsetup status "$mapper" >/dev/null 2>&1; then
    echo "  Mapper:        /dev/mapper/${mapper} (OPEN)"
  else
    echo "  Mapper:        /dev/mapper/${mapper} (closed)"
  fi

  echo "  Verification complete."
}

# ── revert command ─────────────────────────────────────────────────
cmd_revert() {
  local vmid="${1:-}"
  [ -n "$vmid" ] || die "Usage: proxmox-kms-bridge revert <VMID>"
  [[ "$vmid" =~ ^[0-9]+$ ]] || die "VMID must be numeric"
  [ "$(id -u)" -eq 0 ] || die "Must run as root"

  local mapper="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$vmid}"
  local image
  image="$(find_luks_image "$vmid")"
  local state_dir="/root/qxvault-migrate-${vmid}"

  # ── RBD detection ──────────────────────────────────────────────
  local is_rbd=false rbd_spec="" rbd_mapped_dev=""
  local rbd_conf="" rbd_id="" rbd_keyring="" rbd_monhost=""
  local -a rbd_map_args=()
  if [[ "$image" == rbd:* ]]; then
    is_rbd=true
    rbd_spec="${image#rbd:}"
    local rbd_meta="${LUKS_IMAGE_DIR}/${vmid}"
    rbd_conf="$(cat "${rbd_meta}/rbd-conf" 2>/dev/null)" || true
    rbd_id="$(cat "${rbd_meta}/rbd-id" 2>/dev/null)" || true
    rbd_keyring="$(cat "${rbd_meta}/rbd-keyring" 2>/dev/null)" || true
    rbd_monhost="$(cat "${rbd_meta}/rbd-monhost" 2>/dev/null)" || true
    [ -n "$rbd_conf" ] && [ -f "$rbd_conf" ] && rbd_map_args+=(--conf "$rbd_conf")
    [ -n "$rbd_keyring" ] && [ -f "$rbd_keyring" ] && rbd_map_args+=(--keyring "$rbd_keyring")
    [ -n "$rbd_id" ] && rbd_map_args+=(--id "$rbd_id")
    [ -n "$rbd_monhost" ] && rbd_map_args+=(-m "$(echo "$rbd_monhost" | tr ' ' ',')")
  fi

  # Check VM is stopped
  if command -v qm >/dev/null 2>&1 && qm status "$vmid" 2>/dev/null | grep -q running; then
    die "VM ${vmid} is running. Stop it first: qm stop ${vmid}"
  fi

  # Check mapper — close it if VM is stopped
  if cryptsetup status "$mapper" >/dev/null 2>&1; then
    log "Closing LUKS mapper $mapper..."
    local retry=0
    local max_retries=5
    while [ $retry -lt $max_retries ]; do
      if cryptsetup close "$mapper" 2>/dev/null; then
        log "LUKS mapper closed successfully"
        break
      fi
      retry=$((retry + 1))
      if [ $retry -lt $max_retries ]; then
        log "Mapper still in use, waiting 1s before retry ($retry/$max_retries)..."
        sleep 1
      else
        die "LUKS mapper $mapper is still in use after $max_retries attempts."
      fi
    done
  fi

  [ -f "$image" ] || [ -b "$image" ] || [ "$is_rbd" = "true" ] || die "LUKS image not found: $image"

  # ── 2FA check ──────────────────────────────────────────────────
  local tfa_conf="/etc/qxvault/2fa/${vmid}.conf"
  local tfa_mode=""
  if [ -f "$tfa_conf" ]; then
    # shellcheck source=/dev/null
    . "$tfa_conf"
    tfa_mode="${TFA_MODE:-}"
    if [ "$tfa_mode" = "passphrase" ]; then
      die "2FA (passphrase mode) is enabled for VM ${vmid}. Disable 2FA first, then revert."
    fi
    # TOTP mode: vault key opens LUKS directly — proceed and clean up 2FA after
    if [ "$tfa_mode" = "totp" ]; then
      log "2FA (TOTP mode) is enabled — vault key is sufficient, will clean up 2FA config"
    fi
  fi

  local boot_slot="" original_spec=""

  # Try saved rollback state first
  if [ -d "$state_dir" ]; then
    [ -f "${state_dir}/boot-slot.txt" ] && boot_slot=$(<"${state_dir}/boot-slot.txt")
    [ -f "${state_dir}/disk-spec.txt" ] && original_spec=$(<"${state_dir}/disk-spec.txt")
  fi

  # If no saved state, reconstruct from current VM config
  if [ -z "$boot_slot" ]; then
    log "No rollback state found — reconstructing from VM config..."
    local config
    config="$(qm config "$vmid")"

    # Find the disk slot pointing to /dev/mapper (the encrypted disk)
    for s in scsi0 virtio0 sata0 ide0 scsi1 virtio1 sata1; do
      local slot_val
      slot_val="$(echo "$config" | awk -F': ' -v s="$s" '$1==s {print $2}')"
      if [ -n "$slot_val" ] && echo "$slot_val" | grep -q "/dev/mapper/"; then
        boot_slot="$s"
        break
      fi
    done
    [ -n "$boot_slot" ] || die "Cannot find a disk slot using /dev/mapper in VM config"
    log "Found encrypted disk on slot: ${boot_slot}"
  fi

  # Operation lock file (for UI progress tracking)
  local op_lock="/run/qxvault-op-${vmid}.json"
  echo "{\"op\":\"revert\",\"vmid\":\"${vmid}\",\"startTime\":$(date +%s)}" > "$op_lock"

  # Cleanup trap: unmap RBD and close LUKS on error
  # shellcheck disable=SC2064
  trap "
    rm -f '${op_lock}'
    cryptsetup close '${mapper}' 2>/dev/null || true
    if [ -n \"\${rbd_mapped_dev:-}\" ] && [ -b \"\${rbd_mapped_dev:-}\" ]; then
      rbd unmap \"\$rbd_mapped_dev\" 2>/dev/null || true
    fi
  " EXIT

  # Unlock VM
  log "Unlocking VM ${vmid}..."
  qm unlock "$vmid" 2>/dev/null || true

  if [ -n "$original_spec" ]; then
    # ── Fast path: saved state has the original disk spec ──────────
    log "Reverting VMID ${vmid} to saved original disk"
    log "  Original disk: ${boot_slot} = ${original_spec}"

    log "Removing hookscript..."
    qm set "$vmid" --delete hookscript

    log "Restoring VM disk config..."
    qm set "$vmid" "--${boot_slot}" "$original_spec"
  else
    # ── Full path: decrypt LUKS back to a new disk ─────────────────
    log "Original disk was removed. Decrypting LUKS back to a new disk image..."

    # Resolve key name
    local key_name
    key_name="$(read_key_name "$vmid" "$image")"

    # For RBD, map the LUKS image before opening
    local luks_device="$image"
    if [ "$is_rbd" = "true" ]; then
      log "Mapping RBD LUKS image for decryption: ${rbd_spec}"
      rbd_mapped_dev="$(rbd map "${rbd_spec}" "${rbd_map_args[@]}")" \
        || die "Failed to map RBD image: ${rbd_spec}"
      log "RBD mapped to: ${rbd_mapped_dev}"
      luks_device="$rbd_mapped_dev"
    fi

    # Open the LUKS mapper to read from it
    log "Opening LUKS mapper for decryption..."
    /usr/libexec/vault-dmkey get-key --name "$key_name" \
      | base64 -d \
      | cryptsetup open --type luks --key-file=- "$luks_device" "$mapper"

    # Safety: verify the mapper actually opened before we proceed to copy+destroy
    if ! [ -b "/dev/mapper/${mapper}" ]; then
      die "LUKS mapper /dev/mapper/${mapper} did not open — aborting revert to protect data"
    fi

    local mapper_size
    mapper_size="$(blockdev --getsize64 "/dev/mapper/${mapper}")"
    log "Mapper device size: ${mapper_size} bytes ($(numfmt --to=iec "$mapper_size"))"

    if [ "$is_rbd" = "true" ]; then
      # ── RBD full path: create new RBD image via Proxmox storage ──
      local rbd_storage=""
      # Extract storage name from conf path: /etc/pve/priv/ceph/<storage>.conf
      if [ -n "$rbd_conf" ]; then
        rbd_storage="$(basename "$rbd_conf" .conf)"
      fi
      # Fallback: search storage.cfg for the pool
      if [ -z "$rbd_storage" ]; then
        local rbd_pool="${rbd_spec%%/*}"
        rbd_storage="$(awk -v p="$rbd_pool" '
          /^rbd:/ { name=$2; pool=""; next }
          name && /^[[:space:]]+pool/ { pool=$2; next }
          /^[a-z]/ && name { if (pool == p || (pool == "" && name == p)) { print name; exit }; name="" }
          END { if (name && (pool == p || (pool == "" && name == p))) print name }
        ' /etc/pve/storage.cfg)" || true
      fi
      [ -n "$rbd_storage" ] || die "Cannot determine Proxmox storage name for RBD pool"
      log "Using Proxmox storage: ${rbd_storage}"

      local mapper_size_mb=$(( (mapper_size + 1048575) / 1048576 ))
      local new_volid="" alloc_out=""
      for disk_suffix in 0 1 2 3 4 5 6 7 8 9; do
        alloc_out="$(pvesm alloc "$rbd_storage" "$vmid" "vm-${vmid}-disk-${disk_suffix}" "${mapper_size_mb}M" 2>&1)" && break
        alloc_out=""
      done
      [ -n "$alloc_out" ] || die "Cannot allocate new RBD disk in storage ${rbd_storage}"
      # pvesm alloc returns e.g. "successfully created 'rbd_fast:vm-250-disk-0'"
      new_volid="$(echo "$alloc_out" | grep -oP "'\K[^']+" || echo "$alloc_out")"
      log "Allocated new volume: ${new_volid}"

      # Parse the new volume's RBD spec from pvesm path
      local new_pvesm_path new_rbd_spec new_rbd_dev
      new_pvesm_path="$(pvesm path "$new_volid")"
      new_rbd_spec="${new_pvesm_path#rbd:}"
      new_rbd_spec="${new_rbd_spec%%:*}"

      log "Mapping new RBD image: ${new_rbd_spec}"
      new_rbd_dev="$(rbd map "${new_rbd_spec}" "${rbd_map_args[@]}")" \
        || die "Failed to map new RBD image: ${new_rbd_spec}"
      log "New RBD mapped to: ${new_rbd_dev}"

      log "Copying decrypted data (this may take a while)..."
      dd bs=4M if="/dev/mapper/${mapper}" of="$new_rbd_dev" status=progress conv=fsync \
        || die "Failed to copy data to new RBD image"

      log "Unmapping new RBD image..."
      rbd unmap "$new_rbd_dev" 2>/dev/null || true

      log "Closing LUKS mapper..."
      cryptsetup close "$mapper" 2>/dev/null || true

      log "Unmapping LUKS RBD image..."
      rbd unmap "$rbd_mapped_dev" 2>/dev/null || true
      rbd_mapped_dev=""

      log "Removing hookscript..."
      qm set "$vmid" --delete hookscript

      log "Attaching new disk: ${new_volid}"
      qm set "$vmid" "--${boot_slot}" "${new_volid}"
    else
      # ── Non-RBD full path: ZFS, LVM-thin, or file-based ─────────
      local revert_ok=false

      if [[ "$image" == /dev/zvol/* ]]; then
        # ── ZFS zvol: allocate new zvol via Proxmox storage ──────
        # Determine Proxmox storage name from ZFS pool
        local zfs_pool
        zfs_pool="$(echo "$image" | cut -d'/' -f4)"  # /dev/zvol/<pool>/...
        local zfs_storage=""
        zfs_storage="$(awk -v p="$zfs_pool" '
          /^zfspool:/ { name=$2; pool=""; next }
          name && /^[[:space:]]+pool/ { pool=$2; next }
          /^[a-z]/ && name { if (pool == p) { print name; exit }; name="" }
          END { if (name && pool == p) print name }
        ' /etc/pve/storage.cfg)" || true
        [ -n "$zfs_storage" ] || die "Cannot determine Proxmox storage for ZFS pool '${zfs_pool}'"
        log "Using Proxmox storage: ${zfs_storage} (ZFS pool: ${zfs_pool})"

        local mapper_size_mb=$(( (mapper_size + 1048575) / 1048576 ))
        local new_volid="" alloc_out=""
        for disk_suffix in 0 1 2 3 4 5 6 7 8 9; do
          alloc_out="$(pvesm alloc "$zfs_storage" "$vmid" "vm-${vmid}-disk-${disk_suffix}" "${mapper_size_mb}M" 2>&1)" && break
          alloc_out=""
        done
        [ -n "$alloc_out" ] || die "Cannot allocate new zvol in storage ${zfs_storage}"
        # pvesm alloc returns e.g. "successfully created 'zfs:vm-102-disk-0'"
        new_volid="$(echo "$alloc_out" | grep -oP "'\K[^']+" || echo "$alloc_out")"
        log "Allocated new volume: ${new_volid}"

        local new_dev
        new_dev="$(pvesm path "$new_volid")"
        log "New device: ${new_dev}"

        log "Copying decrypted data (this may take a while)..."
        dd bs=4M if="/dev/mapper/${mapper}" of="$new_dev" status=progress conv=fsync \
          || die "Failed to copy data to new zvol"

        log "Closing LUKS mapper..."
        cryptsetup close "$mapper" 2>/dev/null || true

        log "Removing hookscript..."
        qm set "$vmid" --delete hookscript

        log "Attaching new disk: ${new_volid}"
        qm set "$vmid" "--${boot_slot}" "${new_volid}"
        revert_ok=true

      elif [[ "$image" == /dev/* ]]; then
        # ── LVM-thin: allocate new LV via Proxmox storage ────────
        # Determine Proxmox storage from the VG name
        local lvm_vg
        lvm_vg="$(echo "$image" | awk -F'/' '{print $3}')"  # /dev/<vg>/...
        local lvm_storage=""
        lvm_storage="$(awk -v v="$lvm_vg" '
          /^lvmthin:/ { name=$2; vgname=""; next }
          name && /^[[:space:]]+vgname/ { vgname=$2; next }
          /^[a-z]/ && name { if (vgname == v) { print name; exit }; name="" }
          END { if (name && vgname == v) print name }
        ' /etc/pve/storage.cfg)" || true
        [ -n "$lvm_storage" ] || die "Cannot determine Proxmox storage for VG '${lvm_vg}'"
        log "Using Proxmox storage: ${lvm_storage} (VG: ${lvm_vg})"

        local mapper_size_mb=$(( (mapper_size + 1048575) / 1048576 ))
        local new_volid="" alloc_out=""
        for disk_suffix in 0 1 2 3 4 5 6 7 8 9; do
          alloc_out="$(pvesm alloc "$lvm_storage" "$vmid" "vm-${vmid}-disk-${disk_suffix}" "${mapper_size_mb}M" 2>&1)" && break
          alloc_out=""
        done
        [ -n "$alloc_out" ] || die "Cannot allocate new LV in storage ${lvm_storage}"
        # pvesm alloc returns e.g. "successfully created 'local-lvm:vm-102-disk-0'"
        new_volid="$(echo "$alloc_out" | grep -oP "'\K[^']+" || echo "$alloc_out")"
        log "Allocated new volume: ${new_volid}"

        local new_dev
        new_dev="$(pvesm path "$new_volid")"
        log "New device: ${new_dev}"

        log "Copying decrypted data (this may take a while)..."
        dd bs=4M if="/dev/mapper/${mapper}" of="$new_dev" status=progress conv=fsync \
          || die "Failed to copy data to new LV"

        log "Closing LUKS mapper..."
        cryptsetup close "$mapper" 2>/dev/null || true

        log "Removing hookscript..."
        qm set "$vmid" --delete hookscript

        log "Attaching new disk: ${new_volid}"
        qm set "$vmid" "--${boot_slot}" "${new_volid}"
        revert_ok=true

      else
        # ── File-based storage: decrypt to qcow2 ──────────────────
        local image_dir
        image_dir="$(dirname "$image")"
        local new_disk="${image_dir}/vm-${vmid}-disk-0.qcow2"

        # Avoid clobbering an existing file
        local suffix=0
        while [ -e "$new_disk" ]; do
          suffix=$((suffix + 1))
          new_disk="${image_dir}/vm-${vmid}-disk-${suffix}.qcow2"
        done

        log "Creating new disk image: ${new_disk}"
        log "Copying decrypted data (this may take a while)..."
        qemu-img convert -p -f raw -O qcow2 "/dev/mapper/${mapper}" "$new_disk"

        log "Closing LUKS mapper..."
        cryptsetup close "$mapper" 2>/dev/null || true

        log "Removing hookscript..."
        qm set "$vmid" --delete hookscript

        # Determine the storage name from the image path
        # e.g. /mnt/pve/rndnfs/images/243/... → storage "rndnfs"
        local storage_name=""
        if [[ "$image_dir" =~ /mnt/pve/([^/]+)/images/ ]]; then
          storage_name="${BASH_REMATCH[1]}"
        fi

        if [ -n "$storage_name" ]; then
          # Import the new disk through Proxmox storage so it gets a proper volid
          log "Importing disk into Proxmox storage '${storage_name}'..."
          local import_output
          import_output="$(qm importdisk "$vmid" "$new_disk" "$storage_name" --format qcow2 2>&1)" || true
          log "$import_output"

          # Extract the new volume ID from the import output
          local new_volid
          new_volid="$(echo "$import_output" | grep -oP "Successfully imported disk as '\\K[^']+")" || true

          if [ -n "$new_volid" ]; then
            log "Attaching imported disk: ${new_volid}"
            qm set "$vmid" "--${boot_slot}" "${new_volid}"
            # Remove the intermediate file (importdisk copies it)
            rm -f "$new_disk"
          else
            # Fallback: point directly at the qcow2 file
            log "Import did not return volid; pointing VM at file directly"
            qm set "$vmid" "--${boot_slot}" "${new_disk}"
          fi
        else
          # No Proxmox storage detected — point directly at the new file
          log "Pointing VM disk at: ${new_disk}"
          qm set "$vmid" "--${boot_slot}" "${new_disk}"
        fi
        revert_ok=true
      fi
    fi
  fi

  # Remove LUKS image/LV and metadata
  log "Removing LUKS: ${image}"
  if [ "$is_rbd" = "true" ]; then
    # Unmap if still mapped
    if [ -n "$rbd_mapped_dev" ] && [ -b "$rbd_mapped_dev" ]; then
      log "Unmapping LUKS RBD device: ${rbd_mapped_dev}"
      rbd unmap "$rbd_mapped_dev" 2>/dev/null || true
      rbd_mapped_dev=""
    fi
    # Determine the Proxmox storage volume ID for proper cleanup
    local luks_volid=""
    local luks_rbd_storage=""
    if [ -n "$rbd_conf" ]; then
      luks_rbd_storage="$(basename "$rbd_conf" .conf)"
    fi
    if [ -z "$luks_rbd_storage" ]; then
      local cleanup_pool="${rbd_spec%%/*}"
      luks_rbd_storage="$(awk -v p="$cleanup_pool" '
        /^rbd:/ { name=$2; pool=""; next }
        name && /^[[:space:]]+pool/ { pool=$2; next }
        /^[a-z]/ && name { if (pool == p || (pool == "" && name == p)) { print name; exit }; name="" }
        END { if (name && (pool == p || (pool == "" && name == p))) print name }
      ' /etc/pve/storage.cfg)" || true
    fi
    # Construct the volid: storage:image-name
    local luks_image_name="${rbd_spec##*/}"
    if [ -n "$luks_rbd_storage" ]; then
      luks_volid="${luks_rbd_storage}:${luks_image_name}"
      log "Removing LUKS RBD via pvesm: ${luks_volid}"
      pvesm free "$luks_volid" 2>/dev/null \
        || log "WARNING: Failed to remove RBD volume ${luks_volid}"
    else
      log "WARNING: Cannot determine storage for RBD cleanup — trying rbd rm"
      rbd rm "${rbd_spec}" "${rbd_map_args[@]}" 2>/dev/null \
        || log "WARNING: Failed to remove RBD image ${rbd_spec}"
    fi
  elif [[ "$image" == /dev/zvol/* ]]; then
    local zvol_ds="${image#/dev/zvol/}"
    zfs destroy -f "$zvol_ds" 2>/dev/null || log "WARNING: Failed to destroy zvol ${zvol_ds}"
  elif [[ "$image" == /dev/* ]]; then
    lvremove -f "$image" 2>/dev/null || log "WARNING: Failed to remove LV ${image}"
  else
    rm -f "$image"
  fi
  local meta_dir="${LUKS_IMAGE_DIR}/${vmid}"
  rm -f "${meta_dir}/key-name" "${meta_dir}/luks-device"
  rm -f "${meta_dir}/rbd-spec" "${meta_dir}/rbd-conf" "${meta_dir}/rbd-id" "${meta_dir}/rbd-keyring" "${meta_dir}/rbd-monhost"
  # Also check if key-name is next to the image (for NFS-based setups)
  if [ "$is_rbd" != "true" ]; then
    local image_dir2
    image_dir2="$(dirname "$image")"
    if [ "$image_dir2" != "$meta_dir" ]; then
      rm -f "${image_dir2}/key-name"
    fi
  fi

  # Clean up rollback state if it exists
  [ -d "$state_dir" ] && rm -rf "$state_dir"

  # Clean up 2FA config if it was enabled (TOTP mode only reaches here)
  if [ -f "$tfa_conf" ]; then
    log "Removing 2FA config for VMID ${vmid}"
    rm -f "$tfa_conf"
    rm -f "/etc/qxvault/2fa/${vmid}.totp"
  fi

  log "=========================================="
  log "  Revert complete for VMID ${vmid}"
  log "=========================================="
  log "  Hookscript removed"
  log "  LUKS image removed"
  [ -n "$tfa_mode" ] && log "  2FA config removed"
  log "  Note: Key still exists in QxVault (manual cleanup if needed)"
  log "=========================================="
}


cmd_cleanup() {
  local dry_run=false
  if [ "${1:-}" = "--dry-run" ] || [ "${1:-}" = "-n" ]; then
    dry_run=true
  fi

  log "Scanning for orphaned LUKS images..."

  local orphan_count=0
  local cleaned_count=0
  local luks_files=()

  # Collect all LUKS image files first (avoids set -e issues with find+while)
  while IFS= read -r -d '' f; do
    luks_files+=("$f")
  done < <(timeout 5 find "$LUKS_IMAGE_DIR" /var/lib/vz /mnt -maxdepth 4 -name 'vm-*-disk-qxvault.luks' -print0 2>/dev/null | sort -z -u)

  # Also check for LVM-thin LUKS LVs via metadata directories
  for meta_dir in "${LUKS_IMAGE_DIR}"/*/; do
    [ -d "$meta_dir" ] || continue
    if [ -f "${meta_dir}luks-device" ]; then
      local dev
      dev="$(cat "${meta_dir}luks-device")"
      if [ -b "$dev" ]; then
        luks_files+=("$dev")
      fi
    fi
  done

  for luks_target in "${luks_files[@]}"; do
    # Extract VMID from filename or device name
    local vmid=""
    local basename
    basename=$(basename "$luks_target")
    if [[ "$basename" =~ ^vm-([0-9]+)-disk-qxvault ]]; then
      vmid="${BASH_REMATCH[1]}"
    fi

    if [ -z "$vmid" ]; then
      continue
    fi

    # Check if VM still exists in Proxmox
    if qm status "$vmid" >/dev/null 2>&1; then
      continue
    fi

    orphan_count=$((orphan_count + 1))

    # Close LUKS mapper if still open
    local mapper="${LUKS_MAPPER_TEMPLATE//\{vmid\}/$vmid}"
    if cryptsetup status "$mapper" >/dev/null 2>&1; then
      if [ "$dry_run" = true ]; then
        log "  [dry-run] Would close mapper: $mapper"
      else
        log "  Closing orphaned mapper: $mapper"
        cryptsetup close "$mapper" 2>/dev/null || true
      fi
    fi

    if [ "$dry_run" = true ]; then
      log "  [dry-run] Would remove: $luks_target (VMID $vmid no longer exists)"
    else
      log "  Removing orphaned LUKS: $luks_target (VMID $vmid no longer exists)"
      if [[ "$luks_target" == /dev/zvol/* ]]; then
        local zvol_ds="${luks_target#/dev/zvol/}"
        zfs destroy -f "$zvol_ds" 2>/dev/null || true
      elif [[ "$luks_target" == /dev/* ]]; then
        lvremove -f "$luks_target" 2>/dev/null || true
      else
        rm -f "$luks_target"
      fi
      # Remove the VM's metadata directory
      local meta="${LUKS_IMAGE_DIR}/${vmid}"
      rm -f "${meta}/key-name" "${meta}/luks-device" 2>/dev/null || true
      rmdir "$meta" 2>/dev/null || true
      # Also try the image's own directory if different
      local vmdir
      vmdir=$(dirname "$luks_target")
      [ "$vmdir" != "$meta" ] && rmdir "$vmdir" 2>/dev/null || true
      cleaned_count=$((cleaned_count + 1))
    fi
  done

  if [ "$orphan_count" -eq 0 ]; then
    log "No orphaned LUKS images found."
  elif [ "$dry_run" = true ]; then
    log "Found $orphan_count orphaned LUKS image(s). Run without --dry-run to remove."
  else
    log "Cleaned up $cleaned_count orphaned LUKS image(s)."
  fi
}

# ── Dispatch ───────────────────────────────────────────────────────
COMMAND="${1:-}"
shift || true

case "$COMMAND" in
  status)   cmd_status "$@" ;;
  verify)   cmd_verify "$@" ;;
  revert)   cmd_revert "$@" ;;
  cleanup)  cmd_cleanup "$@" ;;
  help|-h|--help|"") usage ;;
  *)        die "Unknown command: $COMMAND. Run 'proxmox-kms-bridge help' for usage." ;;
esac
