#!/usr/bin/env bash
# ============================================================================
# qcrows-export — Export a QCrows image to qcow2, ISO, or PXE format
#
# Converts a .qcrows archive into:
#   - qcow2: QEMU disk image with rootfs installed, ready for VM boot
#   - iso:   Bootable ISO for live USB media (hybrid MBR/EFI)
#   - pxe:   PXE/TFTP directory for network boot via dnsmasq
#
# Usage:
#   qcrows-export --format qcow2 image.qcrows -o disk.qcow2
#   qcrows-export --format iso image.qcrows -o live.iso
#   qcrows-export --format pxe image.qcrows -o /srv/tftp
#
# Prerequisites:
#   qcow2: qemu-img, parted, mkfs.ext4, mount, tar
#   iso:   xorriso (or genisoimage), syslinux/mtools (for MBR boot),
#          grub-mkrescue (for EFI boot)
#   pxe:   dnsmasq (for DHCP+TFTP), squashfs-tools (for rootfs packing)
# ============================================================================
set -euo pipefail

# ─── Version ────────────────────────────────────────────────────────────────
VERSION="0.2.0"

# ─── Colors ─────────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BLUE='\033[0;34m'
NC='\033[0m'

log()   { echo -e "${GREEN}[qcrows-export]${NC} $*"; }
warn()  { echo -e "${YELLOW}[qcrows-export]${NC} WARNING: $*"; }
error() { echo -e "${RED}[qcrows-export]${NC} ERROR: $*" >&2; }
die()   { error "$@"; exit 1; }

# ─── Defaults ───────────────────────────────────────────────────────────────
FORMAT=""
INPUT=""
OUTPUT=""
DISK_SIZE=""        # calculated from rootfs size if empty
BOOT_LABEL="QCROWS"
ROOT_LABEL="ROOTFS"
EFI_PARTITION=false
COMPRESS_QCOW2=true
VERBOSE=false
PXE_SERVER_IP=""
PXE_DEFAULT=true
PXE_APPEND_CMDLINE=""

# ─── Temp directory ────────────────────────────────────────────────────────
TMPDIR=""
WORKDIR=""
cleanup() {
  if [[ -n "$WORKDIR" && -d "$WORKDIR" ]]; then
    mountpoint -q "${WORKDIR}/mnt/rootfs" 2>/dev/null && umount -l "${WORKDIR}/mnt/rootfs" 2>/dev/null || true
    mountpoint -q "${WORKDIR}/mnt/efi" 2>/dev/null && umount -l "${WORKDIR}/mnt/efi" 2>/dev/null || true
    losetup -D 2>/dev/null || true
  fi
  [[ -n "$TMPDIR" && -d "$TMPDIR" ]] && rm -rf "$TMPDIR"
}
trap cleanup EXIT

# ─── Helpers ────────────────────────────────────────────────────────────────

# Return the first existing file from candidate paths
candidate_file() {
  local p
  for p in "$@"; do
    [[ -f "$p" ]] && { echo "$p"; return 0; }
  done
  return 1
}

# Return the first existing file from candidate paths (supports glob expansion)
first_existing_file() {
  local candidate match
  for candidate in "$@"; do
    for match in $candidate; do
      if [[ -f "$match" ]]; then
        echo "$match"
        return 0
      fi
    done
  done
  return 1
}

# Copy source to destination only if source exists
copy_if_exists() {
  local src="$1" dest="$2"
  [[ -f "$src" ]] && cp "$src" "$dest"
}

# Construct kernel command line from base params + optional boot-params.conf
build_cmdline() {
  local base="$1"
  local params_file="$2"
  if [[ -n "$params_file" && -f "$params_file" ]]; then
    local extra
    extra=$(tr '\n' ' ' < "$params_file" | sed 's/  */ /g; s/^ *//; s/ *$//')
    echo "${base} ${extra}"
  else
    echo "$base"
  fi
}

# ─── Usage ──────────────────────────────────────────────────────────────────
usage() {
  cat <<EOF
qcrows-export v${VERSION} — Export QCrows image to qcow2 or ISO

USAGE:
  qcrows-export --format FORMAT INPUT -o OUTPUT [OPTIONS]

REQUIRED:
  --format FORMAT     Export format: qcow2 | iso | pxe
  INPUT               Path to .qcrows archive
  -o, --output PATH   Output file (qcow2/iso) or TFTP directory (pxe)

QCOW2 OPTIONS:
  --disk-size SIZE    Disk size (e.g., 2G, 500M). Default: rootfs size + 256M
  --no-compress       Do not compress the qcow2 image
  --efi               Create an EFI-compatible partition layout

ISO OPTIONS:
  --boot-label LABEL  Boot partition label (default: QCROWS)
  --root-label LABEL  Root filesystem label (default: ROOTFS)
  --efi               Create EFI-bootable ISO (requires grub-mkrescue)

PXE OPTIONS:
  --server-ip IP      PXE/DHCP server IP (required for pxe format)
  --no-default        Do not set as default PXE entry
  --append-cmdline S  Extra kernel command-line parameters

GENERAL OPTIONS:
  -v, --verbose       Verbose output
  -h, --help          Show this help
  --version           Show version

PREREQUISITES:
  qcow2:  qemu-img, parted, mkfs.ext4, mount, tar
  iso:    xorriso or genisoimage, isolinux (for MBR), grub-mkrescue (for EFI)
  pxe:    dnsmasq, squashfs-tools (mksquashfs)

EXAMPLES:
  # Export to qcow2 (sized from rootfs)
  qcrows-export --format qcow2 alpine-kata.qcrows -o alpine-kata.qcow2

  # Export to 2GB qcow2 with EFI support
  qcrows-export --format qcow2 --disk-size 2G --efi alpine-kata.qcrows -o alpine-efi.qcow2

  # Export to bootable ISO for USB
  qcrows-export --format iso alpine-kata.qcrows -o alpine-live.iso

  # Write ISO to USB
  dd if=alpine-live.iso of=/dev/sdX bs=4M status=progress && sync

  # Export to PXE/TFTP directory
  qcrows-export --format pxe --server-ip 192.168.1.1 alpine-kata.qcrows -o /srv/tftp
EOF
}

# ─── Parse Arguments ────────────────────────────────────────────────────────
parse_args() {
  local positional_args=()
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --format)        FORMAT="$2"; shift 2 ;;
      --output|-o)     OUTPUT="$2"; shift 2 ;;
      --disk-size)     DISK_SIZE="$2"; shift 2 ;;
      --no-compress)   COMPRESS_QCOW2=false; shift ;;
      --efi)           EFI_PARTITION=true; shift ;;
      --boot-label)    BOOT_LABEL="$2"; shift 2 ;;
      --root-label)    ROOT_LABEL="$2"; shift 2 ;;
      --verbose|-v)    VERBOSE=true; shift ;;
      --server-ip)     PXE_SERVER_IP="$2"; shift 2 ;;
      --no-default)    PXE_DEFAULT=false; shift ;;
      --append-cmdline) PXE_APPEND_CMDLINE="$2"; shift 2 ;;
      --help|-h)       usage; exit 0 ;;
      --version)       echo "qcrows-export v${VERSION}"; exit 0 ;;
      -*)              die "Unknown option: $1 (use --help)" ;;
      *)               positional_args+=("$1"); shift ;;
    esac
  done

  INPUT="${positional_args[0]:-}"

  [[ -z "$FORMAT" ]] && die "Format is required (--format qcow2|iso|pxe)"
  [[ -z "$INPUT" ]]  && die "Input .qcrows file is required"
  [[ -z "$OUTPUT" ]] && die "Output path is required (--output or -o)"
  [[ ! -f "$INPUT" ]] && die "Input file not found: $INPUT"

  case "$FORMAT" in
    qcow2|iso|pxe) ;;
    *)              die "Unsupported format: ${FORMAT} (use qcow2, iso, or pxe)" ;;
  esac

  [[ "$FORMAT" == "pxe" && -z "$PXE_SERVER_IP" ]] && die "PXE format requires --server-ip"
}

# ─── Prerequisite Checks ───────────────────────────────────────────────────
check_prereqs_qcow2() {
  local -a missing=()
  command -v qemu-img &>/dev/null   || missing+=(qemu-img)
  command -v parted &>/dev/null     || missing+=(parted)
  command -v mkfs.ext4 &>/dev/null  || missing+=(mkfs.ext4)
  command -v mount &>/dev/null      || missing+=(mount)
  command -v tar &>/dev/null        || missing+=(tar)

  [[ ${#missing[@]} -gt 0 ]] && die "Missing prerequisites for qcow2 export: ${missing[*]}. Install qemu-utils, parted, e2fsprogs."
}

check_prereqs_iso() {
  local -a missing=()

  # ISO tool detection — ordered preference
  local -a iso_tools=(xorriso genisoimage mkisofs)
  ISO_TOOL=""
  for tool in "${iso_tools[@]}"; do
    if command -v "$tool" &>/dev/null; then
      ISO_TOOL="$tool"
      break
    fi
  done
  [[ -z "$ISO_TOOL" ]] && missing+=("xorriso or genisoimage")

  # MBR boot requires isolinux
  if [[ "$EFI_PARTITION" == "false" ]]; then
    candidate_file /usr/lib/ISOLINUX/isolinux.bin /usr/share/syslinux/isolinux.bin &>/dev/null \
      || missing+=(isolinux/syslinux)
  fi

  # EFI boot requires grub-mkrescue
  [[ "$EFI_PARTITION" == "true" ]] && { command -v grub-mkrescue &>/dev/null || missing+=(grub-mkrescue); }

  [[ ${#missing[@]} -gt 0 ]] && die "Missing prerequisites for ISO export: ${missing[*]}. Install xorriso, syslinux, and/or grub2."
}

# ─── Extract QCrows Archive ────────────────────────────────────────────────
extract_qcrows() {
  local archive="$1"
  local dest="$2"

  log "Extracting QCrows archive: $(basename "$archive")"

  case "$archive" in
    *.gz)  tar xzf "$archive" -C "$dest" 2>/dev/null || tar xf "$archive" -C "$dest" 2>/dev/null ;;
    *)     tar xf "$archive" -C "$dest" 2>/dev/null ;;
  esac

  # Descend into subdirectory if the archive wraps everything in one
  local -a top_entries
  mapfile -t top_entries < <(find "$dest" -maxdepth 1 -type d ! -path "$dest")
  if [[ ${#top_entries[@]} -eq 1 && -f "${top_entries[0]}/metadata.toml" ]]; then
    mv "${top_entries[0]}"/* "$dest/" 2>/dev/null || true
    rmdir "${top_entries[0]}" 2>/dev/null || true
  fi

  # Validate required components
  [[ -f "${dest}/metadata.toml" ]]   || die "metadata.toml not found in archive"
  [[ -f "${dest}/hashes.sha256" ]]   || die "hashes.sha256 not found in archive"

  # Find rootfs
  ROOTFS_FILE=""
  local -a rootfs_candidates=()
  mapfile -t rootfs_candidates < <(find "$dest" -maxdepth 1 -name 'rootfs*' -type f 2>/dev/null)
  [[ ${#rootfs_candidates[@]} -gt 0 ]] || die "rootfs not found in archive"
  ROOTFS_FILE="${rootfs_candidates[0]}"
  log "Found rootfs: $(basename "$ROOTFS_FILE")"

  # Find kernel
  KERNEL_FILE="$(first_existing_file "${dest}/kernel/vmlinuz" "${dest}/kernel/vmlinux" || true)"
  [[ -n "$KERNEL_FILE" ]] || die "kernel not found in archive (required for bootable export)"
  log "Found kernel: $(basename "$KERNEL_FILE")"

  # Find initrd
  INITRD_FILE=""
  local -a initrd_candidates=()
  mapfile -t initrd_candidates < <(find "$dest" -maxdepth 1 \( -name 'initrd*' -o -name 'initramfs*' \) -type f 2>/dev/null)
  if [[ ${#initrd_candidates[@]} -gt 0 ]]; then
    INITRD_FILE="${initrd_candidates[0]}"
    log "Found initrd: $(basename "$INITRD_FILE")"
  else
    warn "No initrd found — ISO boot may require initramfs"
  fi

  # Find boot-params
  BOOT_PARAMS_FILE=""
  [[ -f "${dest}/boot-params.conf" ]] && { BOOT_PARAMS_FILE="${dest}/boot-params.conf"; log "Found boot-params.conf"; }

  # Read kernel version from metadata
  KERNEL_VERSION=$(awk '/\[ kernel \]/{found=1} found && /^version/{print; exit}' "$dest/metadata.toml" | sed 's/.*= *//' | tr -d '"' || echo "unknown")
  if [[ -z "$KERNEL_VERSION" || "$KERNEL_VERSION" == "unknown" ]]; then
    KERNEL_VERSION=$(strings "$KERNEL_FILE" 2>/dev/null | grep -oP '^\d+\.\d+\.\d+' | head -1 || echo "unknown")
  fi
  log "Kernel version: ${KERNEL_VERSION}"

  # Verify hashes
  log "Verifying archive integrity..."
  (cd "$dest" && sha256sum --quiet -c hashes.sha256 2>/dev/null) || warn "SHA-256 verification reported mismatches — proceed with caution"
}

# ─── Compute disk size from rootfs ─────────────────────────────────────────
compute_disk_size() {
  if [[ -n "$DISK_SIZE" ]]; then
    echo "$DISK_SIZE"
    return
  fi

  local rootfs_bytes
  rootfs_bytes=$(stat --format='%s' "$ROOTFS_FILE" 2>/dev/null || stat -f '%z' "$ROOTFS_FILE" 2>/dev/null || echo 0)
  # Assume 2.5x compression ratio for uncompressed rootfs size estimate
  local estimated_mb=$(( (rootfs_bytes * 5 / 2) / (1024 * 1024) ))
  # Add 256MB headroom, minimum 512MB
  local total_mb=$(( estimated_mb + 256 ))
  [[ $total_mb -lt 512 ]] && total_mb=512
  # Round up to next 128MB boundary
  total_mb=$(( ((total_mb / 128) + 1) * 128 ))
  echo "${total_mb}M"
}

# ─── Export to qcow2 ───────────────────────────────────────────────────────
export_qcow2() {
  local archive="$1"
  local output="$2"

  check_prereqs_qcow2

  WORKDIR="$(mktemp -d /tmp/qcrows-export-qcow2-XXXXXX)"
  local extract_dir="${WORKDIR}/extracted"
  mkdir -p "$extract_dir"

  extract_qcrows "$archive" "$extract_dir"

  local disk_size
  disk_size=$(compute_disk_size)
  log "Disk size: ${disk_size}"

  # ── Step 1: Create raw disk image ─────────────────────────────────────
  local raw_img="${WORKDIR}/disk.raw"
  log "Creating raw disk image (${disk_size})..."
  qemu-img create -f raw "$raw_img" "$disk_size"

  # ── Step 2: Partition ─────────────────────────────────────────────────
  log "Partitioning disk..."

  local efi_part="" root_part=""
  if [[ "$EFI_PARTITION" == "true" ]]; then
    parted -s "$raw_img" mklabel gpt
    parted -s "$raw_img" mkpart ESP fat32 1MiB 513MiB
    parted -s "$raw_img" set 1 esp on
    parted -s "$raw_img" mkpart root ext4 513MiB 100%
  else
    parted -s "$raw_img" mklabel msdos
    parted -s "$raw_img" mkpart primary ext4 1MiB 100%
    parted -s "$raw_img" set 1 boot on
  fi

  # ── Step 3: Set up loop device ────────────────────────────────────────
  log "Setting up loop device..."
  local loop_dev
  loop_dev=$(losetup --find --show --partscan "$raw_img")

  sleep 1
  partprobe "$loop_dev" 2>/dev/null || true
  sleep 1

  local -a partitions=()
  mapfile -t partitions < <(ls "${loop_dev}p"* 2>/dev/null || echo "")

  if [[ "$EFI_PARTITION" == "true" ]]; then
    efi_part="${partitions[0]}"
    root_part="${partitions[1]}"
    log "Formatting EFI partition..."
    mkfs.vfat -F 32 -n "$BOOT_LABEL" "$efi_part" 2>/dev/null || die "Failed to format EFI partition"
  else
    root_part="${partitions[0]}"
  fi

  log "Formatting root partition..."
  mkfs.ext4 -L "$ROOT_LABEL" -q "$root_part" || die "Failed to format root partition"

  # ── Step 4: Mount and install rootfs ──────────────────────────────────
  local mnt="${WORKDIR}/mnt"
  mkdir -p "${mnt}/rootfs"

  log "Mounting root partition..."
  mount "$root_part" "${mnt}/rootfs"

  log "Extracting rootfs onto disk..."
  case "$ROOTFS_FILE" in
    *.tar.gz|*.tgz)  tar xzf "$ROOTFS_FILE" -C "${mnt}/rootfs" ;;
    *.tar.xz|*.txz)  tar xJf "$ROOTFS_FILE" -C "${mnt}/rootfs" ;;
    *.tar.zst)        tar --zstd -xf "$ROOTFS_FILE" -C "${mnt}/rootfs" ;;
    *)                tar xf "$ROOTFS_FILE" -C "${mnt}/rootfs" ;;
  esac

  # ── Step 5: Install kernel ───────────────────────────────────────────
  log "Installing kernel into /boot..."
  mkdir -p "${mnt}/rootfs/boot"

  cp "$KERNEL_FILE" "${mnt}/rootfs/boot/vmlinuz-${KERNEL_VERSION}"
  copy_if_exists "${extract_dir}/kernel/config" "${mnt}/rootfs/boot/config-${KERNEL_VERSION}"
  [[ -n "$INITRD_FILE" ]] && cp "$INITRD_FILE" "${mnt}/rootfs/boot/initrd-${KERNEL_VERSION}"

  # ── Step 6: Create fstab ─────────────────────────────────────────────
  log "Generating /etc/fstab..."
  local -a fstab_lines=("LABEL=${ROOT_LABEL}  /          ext4  defaults,noatime  0 1")
  [[ "$EFI_PARTITION" == "true" ]] && fstab_lines+=("LABEL=${BOOT_LABEL}  /boot/efi  vfat  defaults,noatime  0 2")
  printf '%s\n' "# /etc/fstab — QCrows export disk layout" "${fstab_lines[@]}" > "${mnt}/rootfs/etc/fstab"

  # ── Step 7: Install bootloader ───────────────────────────────────────
  local cmdline
  cmdline=$(build_cmdline "root=LABEL=${ROOT_LABEL} rw console=ttyS0 console=tty0" "$BOOT_PARAMS_FILE")

  if [[ "$EFI_PARTITION" == "false" ]]; then
    # MBR boot with extlinux
    log "Installing extlinux bootloader..."
    mkdir -p "${mnt}/rootfs/boot/extlinux"

    local -a extlinux_lines=(
      "# extlinux.conf — QCrows boot configuration"
      "DEFAULT kata"
      "TIMEOUT 30"
      ""
      "LABEL kata"
      "  KERNEL /boot/vmlinuz-${KERNEL_VERSION}"
    )
    [[ -n "$INITRD_FILE" ]] && extlinux_lines+=("  INITRD /boot/initrd-${KERNEL_VERSION}")
    extlinux_lines+=("  APPEND ${cmdline}")
    printf '%s\n' "${extlinux_lines[@]}" > "${mnt}/rootfs/boot/extlinux/extlinux.conf"

    # Install extlinux if available
    command -v extlinux &>/dev/null && extlinux --install "${mnt}/rootfs/boot/extlinux" 2>/dev/null \
      || warn "extlinux install failed — install syslinux-extlinux"

    # Install MBR boot code
    local mbr_bin
    mbr_bin=$(candidate_file /usr/lib/syslinux/mbr/mbr.bin /usr/share/syslinux/mbr.bin || true)
    [[ -n "$mbr_bin" ]] && dd if="$mbr_bin" of="$raw_img" bs=440 count=1 conv=notrunc 2>/dev/null \
      || warn "MBR boot code install failed — install syslinux"
  else
    # EFI boot with GRUB
    log "Setting up EFI bootloader..."
    mkdir -p "${mnt}/efi"
    mount "$efi_part" "${mnt}/efi"
    mkdir -p "${mnt}/efi/EFI/BOOT"

    cp "$KERNEL_FILE" "${mnt}/efi/EFI/BOOT/vmlinuz"
    [[ -n "$INITRD_FILE" ]] && cp "$INITRD_FILE" "${mnt}/efi/EFI/BOOT/initrd"

    # Build GRUB config
    local -a grub_lines=(
      "# grub.cfg — QCrows EFI boot configuration"
      "set timeout=3"
      "set default=0"
      ""
      "menuentry \"Kata Container Image (${KERNEL_VERSION})\" {"
      "  linux /EFI/BOOT/vmlinuz ${cmdline}"
    )
    [[ -n "$INITRD_FILE" ]] && grub_lines+=("  initrd /EFI/BOOT/initrd")
    grub_lines+=("}")
    printf '%s\n' "${grub_lines[@]}" > "${mnt}/efi/EFI/BOOT/grub.cfg"

    # Locate and install GRUB EFI binary
    local grub_x64
    grub_x64=$(candidate_file \
      "/usr/lib/grub/x86_64-efi/grubx64.efi" \
      "/usr/share/grub/x86_64-efi/grubx64.efi" \
      "/boot/efi/EFI/fedora/grubx64.efi" \
      "/boot/efi/EFI/debian/grubx64.efi" \
      "/boot/efi/EFI/ubuntu/grubx64.efi" \
    || true)

    if [[ -n "$grub_x64" ]]; then
      cp "$grub_x64" "${mnt}/efi/EFI/BOOT/grubx64.efi"
      log "Installed GRUB EFI bootloader"
    else
      warn "No GRUB EFI binary found — EFI boot requires a bootloader"
      warn "Install grub2-efi-x64 or similar package"
    fi

    umount "${mnt}/efi"
  fi

  # ── Step 8: Unmount ──────────────────────────────────────────────────
  log "Syncing and unmounting..."
  sync
  umount "${mnt}/rootfs"

  # Detach loop device
  losetup -d "$loop_dev" 2>/dev/null || true

  # ── Step 9: Convert to qcow2 ────────────────────────────────────────
  log "Converting to qcow2..."
  local -a qcow2_opts=()
  [[ "$COMPRESS_QCOW2" == "true" ]] && qcow2_opts+=(-c)

  qemu-img convert -f raw -O qcow2 "${qcow2_opts[@]}" "$raw_img" "$output"

  local final_size
  final_size=$(du -sh "$output" | cut -f1)

  # ── Summary ──────────────────────────────────────────────────────────
  echo ""
  echo -e "${GREEN}╔══════════════════════════════════════════════╗${NC}"
  echo -e "${GREEN}║   QCrows → qcow2 export complete            ║${NC}"
  echo -e "${GREEN}╚══════════════════════════════════════════════╝${NC}"
  echo ""
  echo -e "  Output:      ${CYAN}${output}${NC}"
  echo -e "  Size:        ${CYAN}${final_size}${NC}"
  echo -e "  Disk size:   ${CYAN}${disk_size}${NC}"
  echo -e "  Kernel:      ${CYAN}${KERNEL_VERSION}${NC}"
  echo -e "  Boot:        ${CYAN}$([ "$EFI_PARTITION" == "true" ] && echo "EFI/GPT" || echo "MBR/extlinux")${NC}"
  echo ""
  echo -e "  Boot with QEMU:"
  echo -e "  ${CYAN}qemu-system-x86_64 -m 1G -smp 2 -drive file=${output},format=qcow2 -enable-kvm${NC}"
  echo ""
}

# ─── Export to ISO ──────────────────────────────────────────────────────────
export_iso() {
  local archive="$1"
  local output="$2"

  check_prereqs_iso

  WORKDIR="$(mktemp -d /tmp/qcrows-export-iso-XXXXXX)"
  local extract_dir="${WORKDIR}/extracted"
  local iso_root="${WORKDIR}/iso"
  mkdir -p "$extract_dir" "$iso_root"

  extract_qcrows "$archive" "$extract_dir"

  # ── Step 1: Build ISO directory structure ─────────────────────────────
  log "Building ISO directory structure..."
  mkdir -p "${iso_root}/boot/isolinux"
  mkdir -p "${iso_root}/boot/grub"
  mkdir -p "${iso_root}/live"
  mkdir -p "${iso_root}/kata"

  # ── Step 2: Copy kernel and initrd ───────────────────────────────────
  log "Installing kernel for ISO boot..."
  cp "$KERNEL_FILE" "${iso_root}/boot/vmlinuz"
  cp "$KERNEL_FILE" "${iso_root}/live/vmlinuz"

  if [[ -n "$INITRD_FILE" ]]; then
    cp "$INITRD_FILE" "${iso_root}/boot/initrd"
    cp "$INITRD_FILE" "${iso_root}/live/initrd"
  fi

  copy_if_exists "${extract_dir}/kernel/config" "${iso_root}/boot/"

  # ── Step 3: Copy rootfs and Kata deployment bundle ───────────────────
  log "Bundling rootfs for live boot..."
  cp "$ROOTFS_FILE" "${iso_root}/live/rootfs.tar.gz"

  # Place components in /kata for direct Kata deployment
  cp "$ROOTFS_FILE" "${iso_root}/kata/"
  [[ -n "$INITRD_FILE" ]] && cp "$INITRD_FILE" "${iso_root}/kata/"
  cp "$KERNEL_FILE" "${iso_root}/kata/"
  copy_if_exists "${extract_dir}/kernel/config" "${iso_root}/kata/"

  # Copy QCrows metadata for in-ISO reference
  cp "${extract_dir}/metadata.toml" "${iso_root}/kata/"
  copy_if_exists "${extract_dir}/menu.toml" "${iso_root}/kata/"
  copy_if_exists "${extract_dir}/build.toml" "${iso_root}/kata/"
  copy_if_exists "${extract_dir}/boot-params.conf" "${iso_root}/kata/"

  # ── Step 4: Build kernel command line ────────────────────────────────
  local cmdline
  cmdline=$(build_cmdline "boot=live components quiet" "$BOOT_PARAMS_FILE")

  # ── Step 5: Create isolinux config (MBR boot) ────────────────────────
  log "Creating isolinux configuration..."

  local -a isolinux_lines=(
    "# isolinux.cfg — QCrows live boot menu"
    "DEFAULT kata"
    "TIMEOUT 30"
    "PROMPT 0"
    ""
    "LABEL kata"
    "  KERNEL /boot/vmlinuz"
  )
  [[ -n "$INITRD_FILE" ]] && isolinux_lines+=("  INITRD /boot/initrd")
  isolinux_lines+=("  APPEND ${cmdline}")
  printf '%s\n' "${isolinux_lines[@]}" > "${iso_root}/boot/isolinux/isolinux.cfg"

  # ── Step 6: Create GRUB config (EFI boot) ───────────────────────────
  log "Creating GRUB configuration..."

  local -a grub_lines=(
    "# grub.cfg — QCrows EFI boot menu"
    "set timeout=3"
    "set default=0"
    "set graphics=off"
    ""
    "menuentry \"Kata Container Image (${KERNEL_VERSION})\" {"
    "  linux /boot/vmlinuz ${cmdline}"
  )
  [[ -n "$INITRD_FILE" ]] && grub_lines+=("  initrd /boot/initrd")
  grub_lines+=("}")

  # Secondary entry for Kata deployment mode
  grub_lines+=(
    ""
    "menuentry \"Kata Deployment Mode (load rootfs from /kata/)\" {"
    "  linux /boot/vmlinuz kata-deploy=live quiet"
  )
  [[ -n "$INITRD_FILE" ]] && grub_lines+=("  initrd /boot/initrd")
  grub_lines+=("}")
  printf '%s\n' "${grub_lines[@]}" > "${iso_root}/boot/grub/grub.cfg"

  # ── Step 7: Create a README on the ISO ───────────────────────────────
  cat > "${iso_root}/README.txt" <<README
QCrows Live Image — $(date -u +%Y-%m-%d)
=========================================
Kernel: ${KERNEL_VERSION}
Format: QCrows live ISO

This ISO contains a Kata Container image usable in two modes:

1. LIVE BOOT: Boot this ISO directly (via USB or virtual CD-ROM).
   The kernel and initrd load with the rootfs at /live/rootfs.tar.gz

2. KATA DEPLOY: Copy files from the /kata/ directory on this ISO
   to /usr/share/kata-containers/ on your host system:
     - vmlinuz (guest kernel)
     - rootfs.tar.gz (guest filesystem)
     - initrd (initial ramdisk)
     - config (kernel .config)
     - metadata.toml (image manifest)
     - menu.toml (Cockpit UI entry)
     - boot-params.conf (kernel parameters)

   Then update your configuration.toml to point to the new paths.

Verify:  qcrows-verify <archive.qcrows>
Inspect: qcrows-inspect <archive.qcrows>
README

  # ── Step 8: Build ISO ────────────────────────────────────────────────
  log "Building ISO image..."

  if [[ "$EFI_PARTITION" == "true" ]]; then
    log "Creating EFI-bootable ISO with grub-mkrescue..."
    grub-mkrescue -o "$output" "$iso_root" 2>/dev/null || die "grub-mkrescue failed"
  else
    # MBR-bootable ISO using isolinux + xorriso/genisoimage
    log "Creating MBR-bootable ISO..."

    local isolinux_bin
    isolinux_bin=$(candidate_file \
      /usr/lib/ISOLINUX/isolinux.bin \
      /usr/share/syslinux/isolinux.bin \
      /usr/lib/syslinux/isolinux.bin \
    || true)
    [[ -z "$isolinux_bin" ]] && die "isolinux.bin not found — install syslinux or isolinux"

    # Copy isolinux binaries
    cp "$isolinux_bin" "${iso_root}/boot/isolinux/isolinux.bin"

    # Copy syslinux modules
    local syslinux_dir
    syslinux_dir=$(dirname "$isolinux_bin")
    find "$syslinux_dir" -name '*.c32' -exec cp {} "${iso_root}/boot/isolinux/" \; 2>/dev/null || true

    case "$ISO_TOOL" in
      xorriso)
        local isohdpfx
        isohdpfx=$(candidate_file "${syslinux_dir}/isohdpfx.bin" || true)
        local -a xorriso_args=(
          -as mkisofs
          -o "$output"
          -c boot/isolinux/boot.cat
          -b boot/isolinux/isolinux.bin
          -no-emul-boot
          -boot-load-size 4
          -boot-info-table
          -V "${BOOT_LABEL}"
          -J -R
        )
        [[ -n "$isohdpfx" ]] && xorriso_args=(-isohybrid-mbr "$isohdpfx" "${xorriso_args[@]}")
        xorriso "${xorriso_args[@]}" "$iso_root" 2>/dev/null || die "xorriso ISO creation failed"
        ;;
      genisoimage|mkisofs)
        "${ISO_TOOL}" \
          -o "$output" \
          -c boot/isolinux/boot.cat \
          -b boot/isolinux/isolinux.bin \
          -no-emul-boot \
          -boot-load-size 4 \
          -boot-info-table \
          -V "${BOOT_LABEL}" \
          -J -R \
          "$iso_root" 2>/dev/null || die "${ISO_TOOL} ISO creation failed"

        # Make hybrid (dd-able to USB)
        command -v isohybrid &>/dev/null && isohybrid "$output" 2>/dev/null \
          || warn "isohybrid failed — ISO may not be USB-bootable"
        ;;
    esac
  fi

  local final_size
  final_size=$(du -sh "$output" | cut -f1)

  # ── Summary ──────────────────────────────────────────────────────────
  echo ""
  echo -e "${GREEN}╔══════════════════════════════════════════════╗${NC}"
  echo -e "${GREEN}║   QCrows → ISO export complete              ║${NC}"
  echo -e "${GREEN}╚══════════════════════════════════════════════╝${NC}"
  echo ""
  echo -e "  Output:      ${CYAN}${output}${NC}"
  echo -e "  Size:        ${CYAN}${final_size}${NC}"
  echo -e "  Kernel:      ${CYAN}${KERNEL_VERSION}${NC}"
  echo -e "  Boot mode:   ${CYAN}$([ "$EFI_PARTITION" == "true" ] && echo "EFI/GRUB" || echo "MBR/isolinux")${NC}"
  echo -e "  Hybrid:      ${CYAN}$([ "$EFI_PARTITION" == "false" ] && echo "yes (isohybrid)" || echo "EFI only")${NC}"
  echo ""
  echo -e "  Write to USB:"
  echo -e "  ${CYAN}dd if=${output} of=/dev/sdX bs=4M status=progress && sync${NC}"
  echo ""
  echo -e "  Boot with QEMU:"
  echo -e "  ${CYAN}qemu-system-x86_64 -m 1G -smp 2 -cdrom ${output} -enable-kvm${NC}"
  echo ""
  echo -e "  Kata files on ISO:"
  echo -e "  ${CYAN}/kata/vmlinuz${NC}        — Guest kernel"
  echo -e "  ${CYAN}/kata/rootfs.tar.gz${NC}  — Guest rootfs"
  echo -e "  ${CYAN}/kata/initrd${NC}         — Initial ramdisk"
  echo -e "  ${CYAN}/kata/metadata.toml${NC}  — Image manifest"
  echo -e "  ${CYAN}/kata/menu.toml${NC}      — Cockpit UI entry"
  echo ""
}

# ─── PXE Export ─────────────────────────────────────────────────────────────

export_pxe() {
  local archive="$1"
  local tftp_root="$2"

  log "Exporting QCrows image to PXE/TFTP directory"

  # Prerequisite checks
  local -a missing=()
  command -v mksquashfs &>/dev/null || missing+=(mksquashfs)
  command -v dnsmasq &>/dev/null   || missing+=(dnsmasq)
  [[ ${#missing[@]} -gt 0 ]] && warn "Missing tools: ${missing[*]} — PXE boot may not function until installed"

  # Extract archive to temp directory
  TMPDIR=$(mktemp -d "${TMPDIR:-/tmp}/qcrows-export-pxe.XXXXXX")
  extract_qcrows "$archive" "$TMPDIR"

  # Derive image slug from metadata
  local image_name
  image_name=$(awk '/^name/{print $3}' "${TMPDIR}/metadata.toml" | tr -d '"' | tr ' ' '-' || echo "kata-image")
  [[ -z "$image_name" ]] && image_name="kata-image"

  local pxe_dir="${tftp_root}/kata/${image_name}"
  local pxe_cfg_dir="${tftp_root}/pxelinux.cfg"

  log "Creating PXE directory structure: ${pxe_dir}"
  mkdir -p "$pxe_dir" "$pxe_cfg_dir"

  # Copy kernel to TFTP
  log "Copying kernel to ${pxe_dir}/vmlinuz"
  cp "$KERNEL_FILE" "${pxe_dir}/vmlinuz"

  # Copy initrd if present
  local initrd_dest=""
  if [[ -n "$INITRD_FILE" ]]; then
    log "Copying initrd to ${pxe_dir}/initrd.img"
    cp "$INITRD_FILE" "${pxe_dir}/initrd.img"
    initrd_dest="kata/${image_name}/initrd.img"
  fi

  # Pack rootfs as squashfs for efficient network boot
  local rootfs_dest="${pxe_dir}/rootfs.squashfs"
  if command -v mksquashfs &>/dev/null; then
    log "Packing rootfs as squashfs to ${rootfs_dest}"
    local rootfs_unpack_dir
    rootfs_unpack_dir=$(mktemp -d "${TMPDIR}/rootfs-unpack.XXXXXX")
    tar xf "$ROOTFS_FILE" -C "$rootfs_unpack_dir" 2>/dev/null || {
      warn "Rootfs extraction failed — copying raw archive instead"
      cp "$ROOTFS_FILE" "${pxe_dir}/rootfs.tar.gz"
      rootfs_dest="${pxe_dir}/rootfs.tar.gz"
    }
    if [[ -d "$rootfs_unpack_dir" && ! -f "${pxe_dir}/rootfs.tar.gz" ]]; then
      mksquashfs "$rootfs_unpack_dir" "$rootfs_dest" -noappend -comp zstd -Xcompression-level 3 2>/dev/null \
        || mksquashfs "$rootfs_unpack_dir" "$rootfs_dest" -noappend -comp gzip 2>/dev/null \
        || { warn "Squashfs creation failed — using raw rootfs"; cp "$ROOTFS_FILE" "${pxe_dir}/rootfs.tar.gz"; rootfs_dest="${pxe_dir}/rootfs.tar.gz"; }
      rm -rf "$rootfs_unpack_dir"
    fi
  else
    warn "mksquashfs not available — copying raw rootfs archive"
    cp "$ROOTFS_FILE" "${pxe_dir}/rootfs.tar.gz"
    rootfs_dest="${pxe_dir}/rootfs.tar.gz"
  fi

  # Build kernel command line
  local base_cmdline="root=/dev/nfs nfsroot=${PXE_SERVER_IP}:${pxe_dir} ip=dhcp console=ttyS0,115200"
  local cmdline
  cmdline=$(build_cmdline "$base_cmdline" "$BOOT_PARAMS_FILE")
  [[ -n "$PXE_APPEND_CMDLINE" ]] && cmdline="${cmdline} ${PXE_APPEND_CMDLINE}"

  # Generate pxelinux configuration
  local pxe_config_file="${pxe_cfg_dir}/${image_name}"
  log "Writing PXE configuration to ${pxe_config_file}"

  {
    [[ "$PXE_DEFAULT" == "true" ]] && echo "DEFAULT ${image_name}"
    echo "LABEL ${image_name}"
    echo "    MENU LABEL Kata ${image_name}"
    echo "    KERNEL kata/${image_name}/vmlinuz"
    [[ -n "$initrd_dest" ]] && echo "    INITRD ${initrd_dest}"
    echo "    APPEND ${cmdline}"
  } > "$pxe_config_file"

  # Update dnsmasq configuration
  local dnsmasq_conf="/etc/dnsmasq.d/kata-pxe.conf"
  if [[ -d "/etc/dnsmasq.d" ]]; then
    log "Updating dnsmasq configuration at ${dnsmasq_conf}"
    if [[ ! -f "$dnsmasq_conf" ]] || ! grep -q "pxe-service" "$dnsmasq_conf" 2>/dev/null; then
      {
        echo ""
        echo "# Kata Containers PXE — managed by qcrows-export"
        echo "dhcp-range=192.168.1.100,192.168.1.200,255.255.255.0,12h"
        echo "dhcp-boot=pxelinux.0"
        echo "pxe-service=x86PC,\"Kata QCrows Boot\",pxelinux"
        echo "enable-tftp"
        echo "tftp-root=${tftp_root}"
        echo ""
      } >> "$dnsmasq_conf"
    fi
    log "dnsmasq config updated — restart dnsmasq to activate"
  else
    warn "/etc/dnsmasq.d not found — configure dnsmasq manually for DHCP+TFTP"
  fi

  # ── Summary ──────────────────────────────────────────────────────────
  echo ""
  echo -e "${GREEN}╔══════════════════════════════════════════════╗${NC}"
  echo -e "${GREEN}║   QCrows → PXE export complete              ║${NC}"
  echo -e "${GREEN}╚══════════════════════════════════════════════╝${NC}"
  echo ""
  echo -e "  TFTP Root:     ${CYAN}${tftp_root}${NC}"
  echo -e "  Image Dir:     ${CYAN}${pxe_dir}${NC}"
  echo -e "  Kernel:        ${CYAN}${pxe_dir}/vmlinuz${NC}"
  [[ -n "$initrd_dest" ]] && echo -e "  Initrd:        ${CYAN}${pxe_dir}/initrd.img${NC}"
  echo -e "  Rootfs:        ${CYAN}${rootfs_dest}${NC}"
  echo -e "  PXE Config:    ${CYAN}${pxe_config_file}${NC}"
  echo -e "  Server IP:     ${CYAN}${PXE_SERVER_IP}${NC}"
  echo -e "  Default Entry: ${CYAN}$([ "$PXE_DEFAULT" == "true" ] && echo "yes" || echo "no")${NC}"
  echo ""
  echo -e "  Next steps:"
  echo -e "  ${CYAN}1. Ensure pxelinux.0 exists in ${tftp_root}/${NC}"
  echo -e "  ${CYAN}2. Restart dnsmasq: systemctl restart dnsmasq${NC}"
  echo -e "  ${CYAN}3. Boot target machines via PXE network boot${NC}"
  echo ""
}

# ─── Entry Point ────────────────────────────────────────────────────────────
parse_args "$@"

case "$FORMAT" in
  qcow2) export_qcow2 "$INPUT" "$OUTPUT" ;;
  iso)   export_iso   "$INPUT" "$OUTPUT" ;;
  pxe)   export_pxe   "$INPUT" "$OUTPUT" ;;
  *)     die "Unhandled format: ${FORMAT}" ;;
esac