#!/usr/bin/env bash
# ============================================================================
# qcrows-initrd-regen — Environment-aware initrd/ramfs/cramfs regeneration
#
# Rebuilds the initrd component of a QCrows bundle with precise awareness
# of the Kata Containers runtime environment. Unlike dracut or mkinitramfs,
# this script produces an initrd where kata-agent IS the final process
# (no pivot_root, no switch_root) and includes ONLY the kernel modules
# required by the detected VMM transport layer.
#
# Usage:
#   qcrows-initrd-regen image.qcrows [options]
#   qcrows-initrd-regen image.qcrows --vmm qemu --compress lz4 -o new-initrd.img
#   qcrows-initrd-regen image.qcrows --in-place
#
# Environment detection (automatic when --vmm is not specified):
#   1. Reads /etc/kata-containers/configuration.toml or the path in
#      KATA_CFG to determine the active hypervisor profile.
#   2. Inspects the QCrows bundle's kernel/.config to determine which
#      compression algorithms the guest kernel supports.
#   3. Examines the rootfs to detect init system type (systemd, OpenRC,
#      busybox /init) and locate the kata-agent binary.
#   4. Detects host architecture for correct module paths.
#
# Output formats:
#   cpio-gzip   — gzip-compressed cpio (default, universal)
#   cpio-lz4    — lz4-compressed cpio (fastest boot, kernel must support)
#   cpio-xz     — xz-compressed cpio (smallest size, slowest boot)
#   cpio-zstd   — zstd-compressed cpio (balanced, kernel 5.15+)
#   cramfs      — cramfs filesystem (read-only, extremely minimal)
# ============================================================================
set -euo pipefail

# ─── Version ────────────────────────────────────────────────────────────────
VERSION="0.1.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}[initrd-regen]${NC} $*"; }
warn()  { echo -e "${YELLOW}[initrd-regen]${NC} WARNING: $*"; }
error() { echo -e "${RED}[initrd-regen]${NC} ERROR: $*" >&2; }
die()   { error "$@"; exit 1; }

# ─── Defaults ───────────────────────────────────────────────────────────────
INPUT=""
OUTPUT=""
VMM=""                # auto-detected from kata config if empty
COMPRESS=""           # auto-detected from kernel config if empty
FORMAT=""             # cpio-gzip, cpio-lz4, cpio-xz, cpio-zstd, cramfs
AGENT_PATH=""         # auto-discovered from rootfs if empty
INIT_STYLE=""         # auto-detected: systemd, openrc, busybox
IN_PLACE=false
KEEP_TEMP=false
STRIP_LOCALES=true
STRIP_DOCS=true
STRIP_MAN=true
EXTRA_MODULES=""      # comma-separated additional kernel modules
EXCLUDE_MODULES=""    # comma-separated modules to exclude
KATA_CFG=""           # path to configuration.toml (auto-detected if empty)
VERBOSE=false
KERNEL_MODULES_DIR="" # auto-discovered from bundle kernel version

# ─── Temp directory ────────────────────────────────────────────────────────
TMPDIR=""
cleanup() {
  if [[ -n "$TMPDIR" && -d "$TMPDIR" && "$KEEP_TEMP" == "false" ]]; then
    rm -rf "$TMPDIR"
  elif [[ -n "$TMPDIR" && -d "$TMPDIR" && "$KEEP_TEMP" == "true" ]]; then
    log "Temp directory preserved: ${TMPDIR}"
  fi
}
trap cleanup EXIT

# ============================================================================
# VMM Kernel Module Matrix
# ============================================================================
# Each VMM requires a specific set of virtio transports and guest drivers.
# This table drives module selection — no nested ifs, just array lookup.
#
# Format: VMM:required_modules:optional_modules
#   required  — MUST be included for the VMM to function
#   optional  — included only if present in the kernel and not excluded

declare -A VMM_REQUIRED_MODULES
declare -A VMM_OPTIONAL_MODULES

VMM_REQUIRED_MODULES[qemu]="virtio_pci virtio_blk virtio_net virtio_rng virtio_balloon virtio_console virtio_gpu"
VMM_OPTIONAL_MODULES[qemu]="virtio_scsi virtio_input virtio_crypto virtio_mem virtio_pmem 9pnet_virtio"

VMM_REQUIRED_MODULES[cloud-hypervisor]="virtio_mmio virtio_blk virtio_net virtio_rng virtio_console virtio_balloon"
VMM_OPTIONAL_MODULES[cloud-hypervisor]="virtio_gpu virtio_scsi 9pnet_virtio"

VMM_REQUIRED_MODULES[firecracker]="virtio_mmio virtio_blk virtio_net virtio_rng"
VMM_OPTIONAL_MODULES[firecracker]="virtio_console"

VMM_REQUIRED_MODULES[dragonball]="virtio_blk virtio_net virtio_rng virtio_console"
VMM_OPTIONAL_MODULES[dragonball]="virtio_balloon virtio_gpu 9pnet_virtio"

# Core modules needed regardless of VMM — filesystem and device infrastructure
CORE_MODULES="ext4 vfat nls_cp437 nls_iso8859_1 dm_mod dm_bufio loop squashfs overlay"

# ============================================================================
# Compression detection from kernel .config
# ============================================================================
# The guest kernel must be compiled with support for the chosen initrd
# compression algorithm. These config symbols determine what works.

declare -A COMPRESS_CONFIG_MAP
COMPRESS_CONFIG_MAP[cpio-gzip]="CONFIG_RD_GZIP"
COMPRESS_CONFIG_MAP[cpio-lz4]="CONFIG_RD_LZ4"
COMPRESS_CONFIG_MAP[cpio-xz]="CONFIG_RD_XZ"
COMPRESS_CONFIG_MAP[cpio-zstd]="CONFIG_RD_ZSTD"

# Compression commands: how to compress the cpio stream
declare -A COMPRESS_CMD
COMPRESS_CMD[cpio-gzip]="gzip -9"
COMPRESS_CMD[cpio-lz4]="lz4 -l"
COMPRESS_CMD[cpio-xz]="xz -9 --check=crc32"
COMPRESS_CMD[cpio-zstd]="zstd -19"

# Compression file extensions
declare -A COMPRESS_EXT
COMPRESS_EXT[cpio-gzip]=".img"
COMPRESS_EXT[cpio-lz4]=".cpio.lz4"
COMPRESS_EXT[cpio-xz]=".cpio.xz"
COMPRESS_EXT[cpio-zstd]=".cpio.zst"

# Priority order for auto-detection (fastest boot first)
COMPRESS_PRIORITY=("cpio-lz4" "cpio-zstd" "cpio-gzip" "cpio-xz")

# ============================================================================
# Kata configuration.toml VMM detection
# ============================================================================
# Maps hypervisor entries in configuration.toml to our VMM identifiers.

declare -A CFG_HYPERVISOR_MAP
CFG_HYPERVISOR_MAP[qemu]="qemu"
CFG_HYPERVISOR_MAP[cloud-hypervisor]="cloud-hypervisor"
CFG_HYPERVISOR_MAP[clh]="cloud-hypervisor"
CFG_HYPERVISOR_MAP[firecracker]="firecracker"
CFG_HYPERVISOR_MAP[fc]="firecracker"
CFG_HYPERVISOR_MAP[dragonball]="dragonball"

# ============================================================================
# Usage
# ============================================================================
usage() {
  cat <<EOF
qcrows-initrd-regen v${VERSION} — Environment-aware initrd regeneration for Kata Containers

USAGE:
  qcrows-initrd-regen QCROWS_FILE [OPTIONS]

REQUIRED:
  QCROWS_FILE             Path to .qcrows archive

VMM SELECTION (auto-detected from kata config if not specified):
  --vmm VMM               Target VMM: qemu|cloud-hypervisor|firecracker|dragonball
  --kata-cfg FILE         Path to configuration.toml (default: auto-detect)

OUTPUT:
  -o, --output FILE       Output initrd file path (default: kata-initrd.img in cwd)
  --in-place              Replace the initrd inside the QCrows bundle
  --format FORMAT         Force output format: cpio-gzip|cpio-lz4|cpio-xz|cpio-zstd|cramfs
  --compress ALGO         Shorthand: gzip|lz4|xz|zstd (sets format to cpio-\$ALGO)

INITRD CONTENT:
  --agent PATH            Path to kata-agent binary (default: extract from rootfs)
  --init-style STYLE      Init system: systemd|openrc|busybox (default: auto-detect)
  --extra-modules LIST    Comma-separated additional kernel modules to include
  --exclude-modules LIST  Comma-separated kernel modules to exclude

OPTIMIZATION:
  --no-strip-locales      Keep locale files in initrd
  --no-strip-docs         Keep documentation files
  --no-strip-man          Keep man pages
  --keep-temp             Preserve temp directory for inspection

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

ENVIRONMENT:
  KATA_CFG                Default path to configuration.toml

AUTO-DETECTION:
  When --vmm is not specified, the script reads configuration.toml
  to determine the active hypervisor. When --compress is not specified,
  it reads the bundled kernel/.config to find the fastest algorithm
  the guest kernel supports.

  When --init-style is not specified, the script examines the rootfs
  to detect the init system:
    - /sbin/init → systemd → systemd init style
    - /sbin/openrc-init → OpenRC → openrc init style
    - Neither → busybox /init style (kata-agent as PID 1)

EXAMPLES:
  # Auto-detect everything, output to cwd
  qcrows-initrd-regen alpine-3.20-kata.qcrows

  # Target QEMU with lz4 compression
  qcrows-initrd-regen alpine.qcrows --vmm qemu --compress lz4

  # Update the bundle in-place
  qcrows-initrd-regen alpine.qcrows --in-place --vmm cloud-hypervisor

  # Custom agent binary and extra modules
  qcrows-initrd-regen gentoo.qcrows --agent /usr/bin/kata-agent \\
    --extra-modules "tun,veth,bridge" -o custom-initrd.img

  # Inspect what would be included without building
  qcrows-initrd-regen alpine.qcrows --keep-temp --verbose
EOF
}

# ============================================================================
# Parse Arguments
# ============================================================================
parse_args() {
  local positional_count=0
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --vmm)             VMM="$2"; shift 2 ;;
      --kata-cfg)        KATA_CFG="$2"; shift 2 ;;
      -o|--output)       OUTPUT="$2"; shift 2 ;;
      --in-place)        IN_PLACE=true; shift ;;
      --format)          FORMAT="$2"; shift 2 ;;
      --compress)        FORMAT="cpio-${2}"; shift 2 ;;
      --agent)           AGENT_PATH="$2"; shift 2 ;;
      --init-style)      INIT_STYLE="$2"; shift 2 ;;
      --extra-modules)   EXTRA_MODULES="$2"; shift 2 ;;
      --exclude-modules) EXCLUDE_MODULES="$2"; shift 2 ;;
      --no-strip-locales) STRIP_LOCALES=false; shift ;;
      --no-strip-docs)    STRIP_DOCS=false; shift ;;
      --no-strip-man)     STRIP_MAN=false; shift ;;
      --keep-temp)        KEEP_TEMP=true; shift ;;
      -v|--verbose)       VERBOSE=true; shift ;;
      -h|--help)          usage; exit 0 ;;
      --version)          echo "qcrows-initrd-regen v${VERSION}"; exit 0 ;;
      -*)                 die "Unknown option: $1 (use --help)" ;;
      *)
        positional_count=$((positional_count + 1))
        case "$positional_count" in
          1) INPUT="$1" ;;
          *) die "Unexpected positional argument: $1" ;;
        esac
        shift
        ;;
    esac
  done
}

# ============================================================================
# Validation (array-driven)
# ============================================================================
validate_inputs() {
  [[ -z "$INPUT" ]] && die "QCrows file is required. Usage: qcrows-initrd-regen <file.qcrows> [options]"
  [[ ! -f "$INPUT" ]] && die "File not found: ${INPUT}"

  # Validate --vmm value against known VMMs
  local -a known_vmms=("qemu" "cloud-hypervisor" "firecracker" "dragonball")
  if [[ -n "$VMM" ]]; then
    local vmm_valid=false
    local known
    for known in "${known_vmms[@]}"; do
      [[ "$VMM" == "$known" ]] && { vmm_valid=true; break; }
    done
    [[ "$vmm_valid" == "false" ]] && die "Unknown VMM: ${VMM}. Supported: ${known_vmms[*]}"
  fi

  # Validate --init-style
  local -a known_init_styles=("systemd" "openrc" "busybox")
  if [[ -n "$INIT_STYLE" ]]; then
    local style_valid=false
    local known
    for known in "${known_init_styles[@]}"; do
      [[ "$INIT_STYLE" == "$known" ]] && { style_valid=true; break; }
    done
    [[ "$style_valid" == "false" ]] && die "Unknown init style: ${INIT_STYLE}. Supported: ${known_init_styles[*]}"
  fi

  # Validate --format
  local -a known_formats=("cpio-gzip" "cpio-lz4" "cpio-xz" "cpio-zstd" "cramfs")
  if [[ -n "$FORMAT" ]]; then
    local fmt_valid=false
    local known
    for known in "${known_formats[@]}"; do
      [[ "$FORMAT" == "$known" ]] && { fmt_valid=true; break; }
    done
    [[ "$fmt_valid" == "false" ]] && die "Unknown format: ${FORMAT}. Supported: ${known_formats[*]}"
  fi

  # Validate --agent path if specified
  if [[ -n "$AGENT_PATH" ]]; then
    [[ ! -f "$AGENT_PATH" ]] && die "Agent binary not found: ${AGENT_PATH}"
    [[ ! -x "$AGENT_PATH" ]] && die "Agent binary not executable: ${AGENT_PATH}"
  fi

  # In-place requires output to be empty
  [[ "$IN_PLACE" == "true" && -n "$OUTPUT" ]] && die "Cannot use --in-place and --output together"

  # cramfs format cannot be used in-place (different structure)
  [[ "$IN_PLACE" == "true" && "$FORMAT" == "cramfs" ]] && die "cramfs format is not supported with --in-place (use cpio format instead)"
}

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

  log "Extracting QCrows bundle: $(basename "$archive")"
  mkdir -p "$dest"

  # Detect compression and extract
  case "$archive" in
    *.gz|*.qcrows.gz|*.tar.gz)
      tar -xzf "$archive" -C "$dest"
      ;;
    *)
      # Try gzip first, then uncompressed
      if tar -tzf "$archive" &>/dev/null; then
        tar -xzf "$archive" -C "$dest"
      else
        tar -xf "$archive" -C "$dest"
      fi
      ;;
  esac

  # Verify required components exist
  local -a required_files=("metadata.toml" "rootfs.tar.gz" "kernel/")
  local req
  for req in "${required_files[@]}"; do
    [[ -e "${dest}/${req}" ]] || die "Missing required component in bundle: ${req}"
  done

  # Verify kernel binary exists
  local kernel_found=false
  [[ -f "${dest}/kernel/vmlinuz" ]] && kernel_found=true
  [[ -f "${dest}/kernel/vmlinux" ]] && kernel_found=true
  [[ "$kernel_found" == "false" ]] && die "No kernel binary found in bundle (expected kernel/vmlinuz or kernel/vmlinux)"

  # Verify kernel config exists (QCrows v0.2 requirement)
  [[ -f "${dest}/kernel/config" ]] || die "No kernel config found in bundle — required for env-aware module selection"

  log "Bundle extracted to: ${dest}"
}

# ============================================================================
# Auto-detect VMM from kata configuration.toml
# ============================================================================
detect_vmm() {
  [[ -n "$VMM" ]] && return

  log "Auto-detecting VMM from Kata configuration..."

  # Candidate configuration.toml paths
  local -a cfg_candidates=(
    "${KATA_CFG:-}"
    "/etc/kata-containers/configuration.toml"
    "/usr/share/defaults/kata-containers/configuration.toml"
    "/usr/share/kata-containers/configuration.toml"
    "/opt/kata/share/defaults/kata-containers/configuration.toml"
  )

  local cfg_path=""
  local candidate
  for candidate in "${cfg_candidates[@]}"; do
    [[ -n "$candidate" && -f "$candidate" ]] && { cfg_path="$candidate"; break; }
  done

  if [[ -z "$cfg_path" ]]; then
    warn "No configuration.toml found — defaulting to qemu"
    VMM="qemu"
    return
  fi

  log "Found configuration: ${cfg_path}"

  # Extract the hypervisor type from configuration.toml.
  # The active hypervisor section is identified by the [hypervisor.qemu] /
  # [hypervisor.cloud-hypervisor] / etc. section headers.
  # The default runtime determines which section is active.

  local runtime_link=""
  if [[ -L "/usr/bin/kata-runtime" ]]; then
    runtime_link="$(readlink -f /usr/bin/kata-runtime 2>/dev/null || true)"
  fi

  # Check kata-runtime symlink naming convention
  local -a runtime_vmm_map=(
    "kata-qemu:qemu"
    "kata-clh:cloud-hypervisor"
    "kata-cloud-hypervisor:cloud-hypervisor"
    "kata-fc:firecracker"
    "kata-firecracker:firecracker"
    "kata-dragonball:dragonball"
  )

  local entry pattern vmm_id
  for entry in "${runtime_vmm_map[@]}"; do
    pattern="${entry%%:*}"
    vmm_id="${entry##*:}"
    if [[ "$runtime_link" == *"$pattern"* ]]; then
      VMM="$vmm_id"
      log "Detected VMM from runtime symlink: ${VMM}"
      return
    fi
  done

  # Fallback: parse configuration.toml for the first hypervisor section
  local -a section_patterns=("qemu" "cloud-hypervisor" "clh" "firecracker" "fc" "dragonball")
  local sect
  for sect in "${section_patterns[@]}"; do
    if grep -q "\\[hypervisor\\.${sect}\\]" "$cfg_path" 2>/dev/null; then
      VMM="${CFG_HYPERVISOR_MAP[$sect]:-$sect}"
      log "Detected VMM from config section: ${VMM}"
      return
    fi
  done

  # Check which runtime is actually active via kata-runtime
  local kata_runtime
  kata_runtime="$(which kata-runtime 2>/dev/null || true)"
  if [[ -n "$kata_runtime" ]]; then
    local runtime_info
    runtime_info="$("$kata_runtime" --version 2>/dev/null | head -1 || true)"
    [[ "$VERBOSE" == "true" ]] && log "kata-runtime version: ${runtime_info}"
  fi

  warn "Could not determine active VMM — defaulting to qemu"
  VMM="qemu"
}

# ============================================================================
# Auto-detect compression from kernel .config
# ============================================================================
detect_compression() {
  [[ -n "$FORMAT" && "$FORMAT" != "cramfs" ]] && return

  log "Auto-detecting compression from kernel config..."

  local kernel_config="${EXTRACT_DIR}/kernel/config"
  [[ ! -f "$kernel_config" ]] && { warn "No kernel config available — defaulting to gzip"; FORMAT="${FORMAT:-cpio-gzip}"; return; }

  # Check each compression algorithm in priority order
  local algo config_key
  for algo in "${COMPRESS_PRIORITY[@]}"; do
    config_key="${COMPRESS_CONFIG_MAP[$algo]:-}"
    [[ -z "$config_key" ]] && continue

    if grep -q "^${config_key}=y" "$kernel_config" 2>/dev/null; then
      FORMAT="$algo"
      log "Kernel supports ${algo##*-} — selected as optimal compression"
      return
    fi
  done

  # No supported compression found in kernel config — this should not happen
  warn "No initrd compression support detected in kernel config — defaulting to gzip"
  FORMAT="${FORMAT:-cpio-gzip}"
}

# ============================================================================
# Auto-detect init style from rootfs
# ============================================================================
detect_init_style() {
  [[ -n "$INIT_STYLE" ]] && return

  log "Auto-detecting init system from rootfs..."

  local rootfs_dir="${EXTRACT_DIR}/rootfs-extracted"
  mkdir -p "$rootfs_dir"

  # Extract just enough of the rootfs to detect init system.
  # We only need the directory listing, not the full content.
  tar -tzf "${EXTRACT_DIR}/rootfs.tar.gz" 2>/dev/null | head -500 > "${TMPDIR}/rootfs-listing.txt"

  if grep -q "sbin/init$\|sbin/systemd$" "${TMPDIR}/rootfs-listing.txt"; then
    INIT_STYLE="systemd"
    log "Detected init system: systemd"
  elif grep -q "sbin/openrc-init$" "${TMPDIR}/rootfs-listing.txt"; then
    INIT_STYLE="openrc"
    log "Detected init system: OpenRC"
  else
    INIT_STYLE="busybox"
    log "Detected init system: busybox (kata-agent as PID 1)"
  fi
}

# ============================================================================
# Discover kata-agent in rootfs
# ============================================================================
discover_agent() {
  [[ -n "$AGENT_PATH" && -f "$AGENT_PATH" ]] && return

  log "Discovering kata-agent in rootfs..."

  local listing="${TMPDIR}/rootfs-listing.txt"

  # Agent candidate paths in the rootfs
  local -a agent_candidates=(
    "usr/bin/kata-agent"
    "usr/local/bin/kata-agent"
    "usr/sbin/kata-agent"
    "sbin/kata-agent"
  )

  local candidate
  for candidate in "${agent_candidates[@]}"; do
    if grep -q "^\\./${candidate}$\\|^${candidate}$" "$listing" 2>/dev/null; then
      AGENT_PATH="${EXTRACT_DIR}/rootfs-extracted/${candidate}"
      log "Found kata-agent at: /${candidate}"
      return
    fi
  done

  die "kata-agent not found in rootfs. Specify path with --agent"
}

# ============================================================================
# Extract rootfs for initrd building
# ============================================================================
extract_rootfs() {
  local rootfs_dir="${EXTRACT_DIR}/rootfs-extracted"
  [[ -d "$rootfs_dir" && -f "${rootfs_dir}/init" ]] && return

  log "Extracting rootfs for initrd assembly..."
  mkdir -p "$rootfs_dir"
  tar -xzf "${EXTRACT_DIR}/rootfs.tar.gz" -C "$rootfs_dir" 2>/dev/null || \
    tar -xf "${EXTRACT_DIR}/rootfs.tar.gz" -C "$rootfs_dir" 2>/dev/null || \
    die "Failed to extract rootfs"

  log "Rootfs extracted: $(du -sh "$rootfs_dir" | cut -f1)"
}

# ============================================================================
# Resolve kernel module paths
# ============================================================================
resolve_kernel_modules() {
  log "Resolving kernel modules for VMM: ${VMM}"

  local kernel_config="${EXTRACT_DIR}/kernel/config"

  # Determine kernel version from the bundled kernel
  local kernel_version
  kernel_version="$(strings "${EXTRACT_DIR}/kernel/vmlinuz" "${EXTRACT_DIR}/kernel/vmlinux" 2>/dev/null \
    | grep -oP '^\d+\.\d+\.\d+' | head -1 || true)"

  if [[ -z "$kernel_version" ]]; then
    # Fallback: extract from kernel config
    kernel_version="$(grep "^CONFIG_VERSION_SIGNATURE=" "$kernel_config" 2>/dev/null \
      | grep -oP '\d+\.\d+\.\d+' | head -1 || true)"
  fi

  if [[ -z "$kernel_version" ]]; then
    warn "Could not determine kernel version — module inclusion will be name-only"
  else
    log "Kernel version: ${kernel_version}"
  fi

  # Build the complete module list from the matrix
  local -a required_modules=()
  local -a optional_modules=()

  # Core modules (always needed)
  local core
  for core in $CORE_MODULES; do
    required_modules+=("$core")
  done

  # VMM-specific required modules
  local vmm_req="${VMM_REQUIRED_MODULES[$VMM]:-}"
  local mod
  for mod in $vmm_req; do
    required_modules+=("$mod")
  done

  # VMM-specific optional modules (only if compiled in kernel)
  local vmm_opt="${VMM_OPTIONAL_MODULES[$VMM]:-}"
  for mod in $vmm_opt; do
    # Check if the module is compiled (=y or =m) in the kernel config
    local config_name="CONFIG_${mod^^}"
    if grep -q "^${config_name}=y\\|^${config_name}=m" "$kernel_config" 2>/dev/null; then
      optional_modules+=("$mod")
    fi
  done

  # Extra modules from --extra-modules
  if [[ -n "$EXTRA_MODULES" ]]; then
    IFS=',' read -ra extra_arr <<< "$EXTRA_MODULES"
    for mod in "${extra_arr[@]}"; do
      [[ -n "$mod" ]] && required_modules+=("$mod")
    done
  fi

  # Apply exclusion list from --exclude-modules
  local -a exclude_arr=()
  if [[ -n "$EXCLUDE_MODULES" ]]; then
    IFS=',' read -ra exclude_arr <<< "$EXCLUDE_MODULES"
  fi

  # Final module list (deduplicated, exclusions applied)
  local -A module_set=()
  local m
  for m in "${required_modules[@]}" "${optional_modules[@]}"; do
    local excluded=false
    local ex
    for ex in "${exclude_arr[@]}"; do
      [[ "$m" == "$ex" ]] && { excluded=true; break; }
    done
    [[ "$excluded" == "false" ]] && module_set[$m]=1
  done

  # Store final module list
  FINAL_MODULES=()
  for m in "${!module_set[@]}"; do
    FINAL_MODULES+=("$m")
  done

  log "Module list (${#FINAL_MODULES[@]} modules): ${FINAL_MODULES[*]}"
  KERNEL_VERSION="$kernel_version"
}

# ============================================================================
# Copy kernel modules from rootfs
# ============================================================================
copy_kernel_modules() {
  local initrd_root="$1"
  local rootfs_dir="${EXTRACT_DIR}/rootfs-extracted"

  [[ -z "$KERNEL_VERSION" ]] && { warn "No kernel version — skipping module copy"; return; }

  local modules_src="${rootfs_dir}/lib/modules/${KERNEL_VERSION}"
  local modules_dest="${initrd_root}/lib/modules/${KERNEL_VERSION}"

  # Check if modules exist in the rootfs
  if [[ ! -d "$modules_src" ]]; then
    warn "Kernel modules directory not found in rootfs: ${modules_src}"
    warn "Initrd will not include loadable kernel modules (built-in only)"
    return
  fi

  log "Copying kernel modules for ${KERNEL_VERSION}..."
  mkdir -p "$modules_dest"

  # Copy module files for selected modules
  local mod_count=0
  local mod
  for mod in "${FINAL_MODULES[@]}"; do
    # Find the module .ko file (multiple possible paths and extensions)
    local -a ko_candidates
    mapfile -t ko_candidates < <(
      find "$modules_src" -name "${mod}.ko*" -type f 2>/dev/null || true
    )

    local ko
    for ko in "${ko_candidates[@]}"; do
      local rel_path="${ko#${modules_src}/}"
      mkdir -p "${modules_dest}/$(dirname "$rel_path")"
      cp "$ko" "${modules_dest}/${rel_path}"
      mod_count=$((mod_count + 1))
      [[ "$VERBOSE" == "true" ]] && log "  Module: ${mod} → ${rel_path}"
    done
  done

  # Copy module metadata files
  local dep_file
  for dep_file in modules.dep modules.dep.bin modules.alias modules.alias.bin modules.symbols modules.symbols.bin; do
    [[ -f "${modules_src}/${dep_file}" ]] && cp "${modules_src}/${dep_file}" "${modules_dest}/"
  done

  # Generate a minimal modules.dep if not present
  if [[ ! -f "${modules_dest}/modules.dep" ]]; then
    (cd "$modules_dest" && find . -name '*.ko*' -print0 | xargs -0 -I{} echo "{}: " | sort > modules.dep)
  fi

  log "Copied ${mod_count} kernel modules"
}

# ============================================================================
# Generate /init for the initrd
# ============================================================================
generate_init() {
  local initrd_root="$1"
  local init_file="${initrd_root}/init"

  log "Generating /init for ${INIT_STYLE} style..."

  case "$INIT_STYLE" in
    systemd)
      # systemd init: mount virtual filesystems, then exec systemd which
      # will start kata-agent.service via the kata-containers.target
      cat > "$init_file" << 'INITEOF'
#!/bin/sh
# Kata Containers initrd — systemd init style
# Mounts virtual filesystems, then hands off to systemd which starts kata-agent.

mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mkdir -p /dev/pts /dev/shm /run
mount -t devpts devpts /dev/pts
mount -t tmpfs tmpfs /dev/shm
mount -t tmpfs tmpfs /run

# Load kernel modules from initrd
[ -d /lib/modules ] && {
  KVER=$(ls /lib/modules/ 2>/dev/null | head -1)
  [ -n "$KVER" ] && depmod -a "$KVER" 2>/dev/null
}

exec /sbin/init
INITEOF
      ;;
    openrc)
      # OpenRC init: mount virtual filesystems, then exec openrc-init
      cat > "$init_file" << 'INITEOF'
#!/bin/sh
# Kata Containers initrd — OpenRC init style
# Mounts virtual filesystems, then hands off to openrc-init which starts kata-agent.

mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mkdir -p /dev/pts /dev/shm /run
mount -t devpts devpts /dev/pts
mount -t tmpfs tmpfs /dev/shm
mount -t tmpfs tmpfs /run

# Load kernel modules from initrd
[ -d /lib/modules ] && {
  KVER=$(ls /lib/modules/ 2>/dev/null | head -1)
  [ -n "$KVER" ] && depmod -a "$KVER" 2>/dev/null
}

exec /sbin/openrc-init
INITEOF
      ;;
    busybox|*)
      # Busybox /init: mount virtual filesystems and exec kata-agent as PID 1.
      # This is the minimal, fastest boot path — no service manager overhead.
      cat > "$init_file" << 'INITEOF'
#!/bin/sh
# Kata Containers initrd — busybox init style
# Mounts virtual filesystems and execs kata-agent as PID 1.
# This is the minimal boot path for fastest VM startup.

mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mkdir -p /dev/pts /dev/shm /run
mount -t devpts devpts /dev/pts
mount -t tmpfs tmpfs /dev/shm
mount -t tmpfs tmpfs /run

# Load kernel modules from initrd
[ -d /lib/modules ] && {
  KVER=$(ls /lib/modules/ 2>/dev/null | head -1)
  [ -n "$KVER" ] && depmod -a "$KVER" 2>/dev/null
}

# Execute kata-agent as the final process (PID 1)
exec /usr/bin/kata-agent
INITEOF
      ;;
  esac

  chmod 755 "$init_file"
  log "Generated /init (${INIT_STYLE} style, $(wc -l < "$init_file") lines)"
}

# ============================================================================
# Helper: Find first existing file in rootfs
# ============================================================================
first_in_rootfs() {
  local name="$1"
  local rootfs="$2"
  local -a paths=("bin/${name}" "sbin/${name}" "usr/bin/${name}" "usr/sbin/${name}" "usr/local/bin/${name}")
  local p
  for p in "${paths[@]}"; do
    if [[ -f "${rootfs}/${p}" ]]; then
      echo "${rootfs}/${p}"
      return 0
    fi
  done
  return 1
}

# ============================================================================
# Helper: Copy if source exists
# ============================================================================
copy_if_exists() {
  local src="$1"
  local dest="$2"
  if [[ -e "$src" ]]; then
    if [[ -d "$src" ]]; then
      mkdir -p "$dest"
      cp -r "${src}/"* "$dest/" 2>/dev/null || true
    else
      mkdir -p "$(dirname "$dest")"
      cp "$src" "$dest"
    fi
  fi
}

# ============================================================================
# Copy required shared libraries
# ============================================================================
copy_required_libs() {
  local initrd_root="$1"
  local rootfs_dir="$2"

  log "Resolving shared library dependencies..."

  # Collect all ELF binaries in the initrd
  local -a elf_files=()
  local elf
  while IFS= read -r -d '' elf; do
    elf_files+=("$elf")
  done < <(find "${initrd_root}" -type f -executable -print0 2>/dev/null | head -50)

  # Use ldd to resolve dependencies (if available)
  if command -v ldd &>/dev/null; then
    local -a lib_set=()
    local file
    for file in "${elf_files[@]}"; do
      local libs
      libs="$(ldd "$file" 2>/dev/null | grep -oP '/[^ ]+\.so[^ ]*' || true)"
      local lib
      for lib in $libs; do
        lib_set+=("$lib")
      done
    done

    # Deduplicate and copy
    local -a unique_libs=()
    local -A seen=()
    local lib
    for lib in "${lib_set[@]}"; do
      [[ -z "${seen[$lib]:-}" ]] && { seen[$lib]=1; unique_libs+=("$lib"); }
    done

    local copied=0
    for lib in "${unique_libs[@]}"; do
      if [[ -f "$lib" ]]; then
        local dest_dir="${initrd_root}/$(dirname "$lib")"
        mkdir -p "$dest_dir"
        cp "$lib" "${initrd_root}/${lib}" 2>/dev/null && copied=$((copied + 1)) || true
      fi
    done
    log "Copied ${copied} shared libraries"

    # Copy the dynamic linker
    local -a linker_candidates=(
      "${rootfs_dir}/lib64/ld-linux-x86-64.so.2"
      "${rootfs_dir}/lib/ld-linux-x86-64.so.2"
      "${rootfs_dir}/lib64/ld-musl-x86_64.so.1"
      "${rootfs_dir}/lib/ld-musl-x86_64.so.1"
    )
    local linker
    for linker in "${linker_candidates[@]}"; do
      if [[ -f "$linker" ]]; then
        mkdir -p "${initrd_root}/$(dirname "${linker#${rootfs_dir}}")"
        cp "$linker" "${initrd_root}/${linker#${rootfs_dir}}"
        log "Copied dynamic linker: ${linker#${rootfs_dir}}"
        break
      fi
    done
  else
    # No ldd — copy common library paths from rootfs
    warn "ldd not available — copying common library paths from rootfs"
    local -a lib_dirs=("lib" "lib64" "usr/lib" "usr/lib64")
    local ldir
    for ldir in "${lib_dirs[@]}"; do
      if [[ -d "${rootfs_dir}/${ldir}" ]]; then
        mkdir -p "${initrd_root}/${ldir}"
        cp "${rootfs_dir}/${ldir}/"*.so* "${initrd_root}/${ldir}/" 2>/dev/null || true
      fi
    done
  fi
}

# ============================================================================
# Strip unnecessary files from initrd
# ============================================================================
strip_initrd() {
  local initrd_root="$1"

  log "Stripping unnecessary files from initrd..."

  # Stripping targets driven by flag-state table
  local -a strip_targets=()
  [[ "$STRIP_LOCALES" == "true" ]]  && strip_targets+=("usr/share/locale")
  [[ "$STRIP_DOCS" == "true" ]]     && strip_targets+=("usr/share/doc")
  [[ "$STRIP_MAN" == "true" ]]      && strip_targets+=("usr/share/man")
  # Always strip these
  strip_targets+=("usr/share/info" "usr/share/groff" "usr/share/emacs")
  strip_targets+=("var/cache" "var/lib/apt" "var/lib/dpkg" "var/lib/pacman")

  local path_pattern
  for path_pattern in "${strip_targets[@]}"; do
    rm -rf "${initrd_root}/${path_pattern}" 2>/dev/null || true
  done

  # Remove Python cache files
  find "$initrd_root" -name '*.pyc' -delete 2>/dev/null || true
  find "$initrd_root" -name '__pycache__' -type d -exec rm -rf {} + 2>/dev/null || true

  # Remove .a and .la files (static libraries not needed at runtime)
  find "$initrd_root" -name '*.a' -delete 2>/dev/null || true
  find "$initrd_root" -name '*.la' -delete 2>/dev/null || true

  # Strip debug symbols from ELF binaries (reduces size significantly)
  if command -v strip &>/dev/null; then
    find "$initrd_root" -type f -executable -print0 2>/dev/null \
      | xargs -0 strip --strip-unneeded 2>/dev/null || true
  fi

  log "Stripping complete: $(du -sh "$initrd_root" | cut -f1)"
}

# ============================================================================
# Build the initrd root filesystem
# ============================================================================
build_initrd_root() {
  local initrd_root="${TMPDIR}/initrd-root"
  local rootfs_dir="${EXTRACT_DIR}/rootfs-extracted"

  log "Building initrd root filesystem..."
  mkdir -p "$initrd_root"

  # Essential directory structure
  local -a essential_dirs=(
    "bin" "sbin" "usr/bin" "usr/sbin" "usr/lib" "usr/lib64"
    "lib" "lib64" "etc" "proc" "sys" "dev" "run"
    "dev/pts" "dev/shm" "tmp" "root" "var/run"
    "var/log" "usr/share" "usr/libexec"
  )
  local dir
  for dir in "${essential_dirs[@]}"; do
    mkdir -p "${initrd_root}/${dir}"
  done

  # ── Copy kata-agent ───────────────────────────────────────────────────
  if [[ -n "$AGENT_PATH" && -f "$AGENT_PATH" ]]; then
    cp "$AGENT_PATH" "${initrd_root}/usr/bin/kata-agent"
    chmod 755 "${initrd_root}/usr/bin/kata-agent"
    log "Installed kata-agent from: ${AGENT_PATH}"
  else
    # Copy from rootfs
    local -a agent_src_candidates=(
      "${rootfs_dir}/usr/bin/kata-agent"
      "${rootfs_dir}/usr/local/bin/kata-agent"
      "${rootfs_dir}/usr/sbin/kata-agent"
      "${rootfs_dir}/sbin/kata-agent"
    )
    local src
    for src in "${agent_src_candidates[@]}"; do
      if [[ -f "$src" ]]; then
        cp "$src" "${initrd_root}/usr/bin/kata-agent"
        chmod 755 "${initrd_root}/usr/bin/kata-agent"
        log "Installed kata-agent from rootfs: ${src#${rootfs_dir}}"
        break
      fi
    done
  fi

  [[ ! -f "${initrd_root}/usr/bin/kata-agent" ]] && die "kata-agent not found — cannot build initrd"

  # ── Copy essential binaries and libraries ──────────────────────────────
  # Minimal set needed for init and kata-agent operation
  local -a essential_bins=(
    "sh" "mount" "umount" "mkdir" "mknod" "ln" "cp" "mv"
    "rm" "ls" "cat" "echo" "sleep" "depmod" "modprobe"
  )

  local bin
  for bin in "${essential_bins[@]}"; do
    local bin_src
    bin_src="$(first_in_rootfs "$bin" "$rootfs_dir")" || true
    if [[ -n "$bin_src" ]]; then
      cp "$bin_src" "${initrd_root}/bin/${bin}" 2>/dev/null || true
    fi
  done

  # Copy busybox if available (provides most essential binaries via symlinks)
  if [[ -f "${rootfs_dir}/bin/busybox" ]]; then
    cp "${rootfs_dir}/bin/busybox" "${initrd_root}/bin/busybox"
    chmod 755 "${initrd_root}/bin/busybox"
    # Create busybox symlinks for our essential bins that aren't already present
    local sym_bin
    for sym_bin in "${essential_bins[@]}"; do
      [[ ! -f "${initrd_root}/bin/${sym_bin}" ]] && \
        ln -sf busybox "${initrd_root}/bin/${sym_bin}" 2>/dev/null || true
    done
    log "Installed busybox with symlinks"
  fi

  # ── Copy shared libraries ─────────────────────────────────────────────
  copy_required_libs "$initrd_root" "$rootfs_dir"

  # ── Copy kernel modules ───────────────────────────────────────────────
  copy_kernel_modules "$initrd_root"

  # ── Copy init system files ────────────────────────────────────────────
  case "$INIT_STYLE" in
    systemd)
      copy_if_exists "${rootfs_dir}/sbin/init" "${initrd_root}/sbin/init"
      copy_if_exists "${rootfs_dir}/lib/systemd" "${initrd_root}/lib/systemd"
      copy_if_exists "${rootfs_dir}/usr/lib/systemd" "${initrd_root}/usr/lib/systemd"
      copy_if_exists "${rootfs_dir}/etc/systemd" "${initrd_root}/etc/systemd"

      # Ensure kata-agent service unit exists
      mkdir -p "${initrd_root}/etc/systemd/system"
      mkdir -p "${initrd_root}/etc/systemd/target"
      if [[ -f "${rootfs_dir}/etc/systemd/system/kata-agent.service" ]]; then
        cp "${rootfs_dir}/etc/systemd/system/kata-agent.service" "${initrd_root}/etc/systemd/system/"
      else
        # Generate minimal service unit
        cat > "${initrd_root}/etc/systemd/system/kata-agent.service" << 'EOF'
[Unit]
Description=kata-agent
Requires=kata-containers.target
After=kata-containers.target

[Service]
ExecStart=/usr/bin/kata-agent
Restart=always
RestartSec=1

[Install]
WantedBy=kata-containers.target
EOF
      fi

      # Generate kata-containers.target
      cat > "${initrd_root}/etc/systemd/system/kata-containers.target" << 'EOF'
[Unit]
Description=kata containers target
Requires=multi-user.target
After=multi-user.target
EOF
      ;;
    openrc)
      copy_if_exists "${rootfs_dir}/sbin/openrc-init" "${initrd_root}/sbin/openrc-init"
      copy_if_exists "${rootfs_dir}/etc/init.d" "${initrd_root}/etc/init.d"
      # Generate OpenRC init script for kata-agent
      mkdir -p "${initrd_root}/etc/init.d"
      cat > "${initrd_root}/etc/init.d/kata-agent" << 'EOF'
#!/sbin/openrc-run
description="kata-agent"
command="/usr/bin/kata-agent"
command_background=true
pidfile="/run/kata-agent.pid"

depend() {
  after localmount
}
EOF
      chmod 755 "${initrd_root}/etc/init.d/kata-agent"
      ;;
  esac

  # ── Generate /init ────────────────────────────────────────────────────
  generate_init "$initrd_root"

  # ── Copy DNS config ───────────────────────────────────────────────────
  if [[ -f "${rootfs_dir}/etc/resolv.conf" ]]; then
    cp "${rootfs_dir}/etc/resolv.conf" "${initrd_root}/etc/resolv.conf"
  else
    echo "nameserver 8.8.8.8" > "${initrd_root}/etc/resolv.conf"
  fi

  # ── Copy /etc/passwd and /etc/group (minimal) ────────────────────────
  copy_if_exists "${rootfs_dir}/etc/passwd" "${initrd_root}/etc/passwd"
  copy_if_exists "${rootfs_dir}/etc/group" "${initrd_root}/etc/group"

  # Generate minimal passwd/group if not present
  [[ ! -f "${initrd_root}/etc/passwd" ]] && echo "root:x:0:0:root:/root:/bin/sh" > "${initrd_root}/etc/passwd"
  [[ ! -f "${initrd_root}/etc/group" ]] && echo "root:x:0:" > "${initrd_root}/etc/group"

  # ── Strip unnecessary files ───────────────────────────────────────────
  strip_initrd "$initrd_root"

  log "Initrd root built: $(du -sh "$initrd_root" | cut -f1)"
}

# ============================================================================
# Create cpio archive
# ============================================================================
create_cpio_archive() {
  local initrd_root="${TMPDIR}/initrd-root"
  local output_file="${OUTPUT:-kata-initrd${COMPRESS_EXT[$FORMAT]:-.img}}"

  log "Creating ${FORMAT} initrd image..."

  # Verify /init exists
  [[ -f "${initrd_root}/init" ]] || die "Missing /init in initrd root — cannot create archive"

  local compress_cmd="${COMPRESS_CMD[$FORMAT]:-cat}"

  # Create cpio archive with compression
  (cd "$initrd_root" && find . | cpio -o -H newc 2>/dev/null) | $compress_cmd > "$output_file"

  local size
  size="$(du -sh "$output_file" | cut -f1)"
  log "Initrd created: ${output_file} (${size})"

  # Validate the output
  validate_initrd "$output_file"

  OUTPUT_FILE="$output_file"
}

# ============================================================================
# Create cramfs image
# ============================================================================
create_cramfs_image() {
  local initrd_root="${TMPDIR}/initrd-root"
  local output_file="${OUTPUT:-kata-initrd.cramfs}"

  log "Creating cramfs initrd image..."

  command -v mkcramfs &>/dev/null || die "mkcramfs not found — install cramfsprogs or util-linux"

  mkcramfs "$initrd_root" "$output_file"

  local size
  size="$(du -sh "$output_file" | cut -f1)"
  log "Cramfs initrd created: ${output_file} (${size})"

  OUTPUT_FILE="$output_file"
}

# ============================================================================
# Validate the generated initrd
# ============================================================================
validate_initrd() {
  local initrd_file="$1"

  log "Validating generated initrd..."

  local -a checks=()

  # Check 1: File exists and is non-empty
  [[ -f "$initrd_file" && -s "$initrd_file" ]] && checks+=("PASS:file_exists") || checks+=("FAIL:file_exists")

  # Check 2: File format matches expected
  local file_type
  file_type="$(file -b "$initrd_file" 2>/dev/null || true)"
  case "$FORMAT" in
    cpio-gzip)  echo "$file_type" | grep -qi "gzip" && checks+=("PASS:format_match") || checks+=("FAIL:format_match:${file_type}") ;;
    cpio-lz4)   echo "$file_type" | grep -qi "lz4" && checks+=("PASS:format_match") || checks+=("WARN:format_match:file(1) may not identify LZ4") ;;
    cpio-xz)    echo "$file_type" | grep -qi "xz" && checks+=("PASS:format_match") || checks+=("FAIL:format_match:${file_type}") ;;
    cpio-zstd)  echo "$file_type" | grep -qi "zstd" && checks+=("PASS:format_match") || checks+=("WARN:format_match:file(1) may not identify ZSTD") ;;
    cramfs)     echo "$file_type" | grep -qi "cramfs" && checks+=("PASS:format_match") || checks+=("FAIL:format_match:${file_type}") ;;
  esac

  # Check 3: Can list archive contents (cpio formats only)
  if [[ "$FORMAT" != "cramfs" ]]; then
    local decompress_cmd
    case "$FORMAT" in
      cpio-gzip) decompress_cmd="zcat" ;;
      cpio-lz4)  decompress_cmd="lz4 -d -c" ;;
      cpio-xz)   decompress_cmd="xzcat" ;;
      cpio-zstd) decompress_cmd="zstdcat" ;;
    esac
    local listing
    listing="$($decompress_cmd "$initrd_file" 2>/dev/null | cpio -t 2>/dev/null | head -50 || true)"
    if echo "$listing" | grep -q "init"; then
      checks+=("PASS:init_exists")
    else
      checks+=("FAIL:init_exists")
    fi
    if echo "$listing" | grep -q "kata-agent"; then
      checks+=("PASS:agent_exists")
    else
      checks+=("FAIL:agent_exists")
    fi
  fi

  # Check 4: Size is reasonable (< 512MB for initrd)
  local size_bytes
  size_bytes="$(stat --format="%s" "$initrd_file" 2>/dev/null || stat -f "%z" "$initrd_file" 2>/dev/null || echo 0)"
  local size_mb=$(( size_bytes / (1024 * 1024) ))
  if [[ $size_mb -lt 512 ]]; then
    checks+=("PASS:size_reasonable:${size_mb}MB")
  else
    checks+=("WARN:size_large:${size_mb}MB — consider stripping more content")
  fi

  # Report results
  local pass=0 fail=0 warn=0
  local check
  for check in "${checks[@]}"; do
    local status="${check%%:*}"
    case "$status" in
      PASS) pass=$((pass + 1)) ;;
      FAIL) fail=$((fail + 1)); error "  ${check}" ;;
      WARN) warn=$((warn + 1)); warn "  ${check}" ;;
    esac
  done

  log "Validation: ${pass} passed, ${fail} failed, ${warn} warnings"
  [[ $fail -gt 0 ]] && die "Initrd validation failed — see errors above"
}

# ============================================================================
# Update QCrows bundle in-place
# ============================================================================
update_qcrows_inplace() {
  [[ "$IN_PLACE" != "true" ]] && return

  local qcrows_file="$INPUT"
  local new_initrd="$OUTPUT_FILE"

  log "Updating QCrows bundle in-place with new initrd..."

  # Create a new bundle with the updated initrd
  local staging="${TMPDIR}/inplace-staging"
  mkdir -p "$staging"

  # Copy all existing bundle contents
  cp -r "${EXTRACT_DIR}/"* "$staging/" 2>/dev/null || true

  # Replace the initrd
  local initrd_name
  initrd_name="$(ls "${staging}"/initrd* 2>/dev/null | head -1 || true)"
  if [[ -n "$initrd_name" ]]; then
    rm -f "${staging}"/initrd.*
    cp "$new_initrd" "${staging}/initrd.img"
  else
    cp "$new_initrd" "${staging}/initrd.img"
  fi

  # Recompute hashes
  (cd "$staging" && find . -type f ! -name "hashes.sha256" -print0 | sort -z | xargs -0 sha256sum) \
    > "${staging}/hashes.sha256"

  # Rebuild the archive
  local backup="${qcrows_file}.bak"
  cp "$qcrows_file" "$backup"
  log "Backup created: ${backup}"

  (cd "$staging" && tar czf - ./*) > "$qcrows_file"

  log "QCrows bundle updated: ${qcrows_file}"
}

# ============================================================================
# Print summary
# ============================================================================
print_summary() {
  local initrd_file="${OUTPUT_FILE:-}"
  local size=""
  [[ -f "$initrd_file" ]] && size="$(du -sh "$initrd_file" | cut -f1)"

  echo ""
  echo -e "${GREEN}╔══════════════════════════════════════════════════════╗${NC}"
  echo -e "${GREEN}║   Initrd regenerated successfully                   ║${NC}"
  echo -e "${GREEN}╚══════════════════════════════════════════════════════╝${NC}"
  echo ""
  echo -e "  Output:      ${CYAN}${initrd_file}${NC}"
  echo -e "  Size:        ${CYAN}${size}${NC}"
  echo -e "  Format:      ${CYAN}${FORMAT}${NC}"
  echo -e "  VMM:         ${CYAN}${VMM}${NC}"
  echo -e "  Init style:  ${CYAN}${INIT_STYLE}${NC}"
  echo -e "  Kernel:      ${CYAN}${KERNEL_VERSION:-unknown}${NC}"
  echo -e "  Modules:     ${CYAN}${#FINAL_MODULES[@]} included${NC}"
  echo -e "  In-place:    ${CYAN}${IN_PLACE}${NC}"
  echo ""
  echo -e "  Next steps:"
  echo -e "    ${CYAN}qcrows-verify ${INPUT}${NC}"
  echo -e "    ${CYAN}qcrows-export --format qcow2 ${INPUT} -o disk.qcow2${NC}"
  echo -e "    ${CYAN}qcrows-export --format pxe ${INPUT} --server-ip 192.168.1.1${NC}"
  echo ""
}

# ============================================================================
# Entry Point
# ============================================================================
EXTRACT_DIR=""
KERNEL_VERSION=""
OUTPUT_FILE=""
declare -a FINAL_MODULES=()

parse_args "$@"
validate_inputs

# ── 1. Setup temp directory ─────────────────────────────────────────────
TMPDIR="$(mktemp -d /tmp/qcrows-initrd-regen-XXXXXX)"
EXTRACT_DIR="${TMPDIR}/bundle"
log "Working directory: ${TMPDIR}"

# ── 2. Extract QCrows bundle ────────────────────────────────────────────
extract_qcrows "$INPUT" "$EXTRACT_DIR"

# ── 3. Detect environment ───────────────────────────────────────────────
detect_vmm
detect_init_style
detect_compression

# ── 4. Extract rootfs for content access ────────────────────────────────
extract_rootfs

# ── 5. Discover kata-agent ──────────────────────────────────────────────
discover_agent

# ── 6. Resolve kernel modules ───────────────────────────────────────────
resolve_kernel_modules

# ── 7. Build initrd root filesystem ─────────────────────────────────────
build_initrd_root

# ── 8. Create the initrd image ──────────────────────────────────────────
if [[ "$FORMAT" == "cramfs" ]]; then
  create_cramfs_image
else
  create_cpio_archive
fi

# ── 9. Update bundle in-place if requested ──────────────────────────────
update_qcrows_inplace

# ── 10. Print summary ───────────────────────────────────────────────────
print_summary