#!/bin/bash
# Initramfs assembly and optimization

# Source common functions
LIB_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${LIB_SCRIPT_DIR}/common.sh"

# Initramfs configuration
INITRAMFS_COMPRESSION="${INITRAMFS_COMPRESSION:-xz}"
XZ_COMPRESSION_LEVEL="${XZ_COMPRESSION_LEVEL:-9}"

# Setup zinit as init system (replaces OpenRC completely)
function initramfs_setup_zinit() {
    local initramfs_dir="$1"
    local zinit_config_dir="$2"
    
    section_header "Setting up zinit as init system"
    
    if [[ ! -d "$zinit_config_dir" ]]; then
        log_error "zinit configuration directory not found: ${zinit_config_dir}"
        return 1
    fi
    
    # Verify zinit binary exists
    if [[ ! -x "${initramfs_dir}/sbin/zinit" ]]; then
        log_error "zinit binary not found or not executable: ${initramfs_dir}/sbin/zinit"
        return 1
    fi
    
    # Copy config/init as /sbin/init (not a symlink to zinit)
    log_info "Installing config/init as /sbin/init"
    safe_execute rm -f "${initramfs_dir}/sbin/init"
    local config_init="${PROJECT_ROOT}/config/init"
    if [[ -f "$config_init" ]]; then
        safe_execute cp "$config_init" "${initramfs_dir}/sbin/init"
        safe_execute chmod 755 "${initramfs_dir}/sbin/init"
        log_info "✓ Installed config/init as /sbin/init"
    else
        log_error "config/init not found: $config_init"
        return 1
    fi
    
    # Copy zinit configuration (all YAML and scripts)
    log_info "Installing zinit configuration"
    safe_mkdir "${initramfs_dir}/etc/zinit"
    safe_execute cp -r "${zinit_config_dir}"/* "${initramfs_dir}/etc/zinit/"
    
    # Ensure proper permissions
    safe_execute chmod 755 "${initramfs_dir}/sbin/zinit"
    safe_execute chmod -R 644 "${initramfs_dir}/etc/zinit"
    safe_execute find "${initramfs_dir}/etc/zinit" -name "*.sh" -exec chmod 755 {} \;
    
    # Create zinit working directories
    safe_mkdir "${initramfs_dir}/var/log/zinit"
    safe_mkdir "${initramfs_dir}/run/zinit"
    
    # Remove any OpenRC remnants (ensure complete replacement)
    local openrc_paths=(
        "/etc/init.d"
        "/etc/runlevels"
        "/etc/conf.d"
        "/sbin/openrc"
        "/sbin/rc-service"
        "/sbin/rc-status"
        "/sbin/rc-update"
    )
    
    for path in "${openrc_paths[@]}"; do
        if [[ -e "${initramfs_dir}${path}" ]]; then
            log_info "Removing OpenRC remnant: ${path}"
            safe_execute rm -rf "${initramfs_dir}${path}"
        fi
    done
    
    log_info "zinit setup complete - OpenRC completely replaced"
}

# Install the critical /init script for initramfs boot
function initramfs_install_init_script() {
    local initramfs_dir="$1"
    local init_script_path="$2"
    
    section_header "Installing initramfs /init script"
    
    if [[ ! -f "$init_script_path" ]]; then
        log_error "Init script not found: ${init_script_path}"
        return 1
    fi
    
    # Install the init script as /init in initramfs root
    log_info "Installing init script from: ${init_script_path}"
    safe_execute cp "$init_script_path" "${initramfs_dir}/init"
    safe_execute chmod 755 "${initramfs_dir}/init"
    
    # Verify installation
    if [[ ! -x "${initramfs_dir}/init" ]]; then
        log_error "Failed to install init script"
        return 1
    fi
    
    log_info "✓ Initramfs /init script installed successfully"
    log_info "  Boot flow: /init -> setup environment -> switch_root -> /sbin/zinit init"
}

# Copy built components to initramfs (separate from building)
function initramfs_copy_components() {
    local initramfs_dir="$1"
    local components_dir="${2:-${PROJECT_ROOT}/components}"
    
    section_header "Copying Built Components to Initramfs"
    
    if [[ ! -d "$components_dir" ]]; then
        log_error "Components directory not found: ${components_dir}"
        return 1
    fi
    
    local copied_count=0
    local missing_count=0
    
    # Copy zinit to /sbin
    local zinit_binary="${components_dir}/zinit/target/x86_64-unknown-linux-musl/release/zinit"
    if [[ -f "$zinit_binary" ]]; then
        safe_mkdir "${initramfs_dir}/sbin"
        safe_execute cp "$zinit_binary" "${initramfs_dir}/sbin/zinit"
        safe_execute chmod +x "${initramfs_dir}/sbin/zinit"
        
        # Strip and UPX compress zinit
        local original_size=$(get_file_size "${initramfs_dir}/sbin/zinit")
        if strip "${initramfs_dir}/sbin/zinit" 2>/dev/null || true; then
            log_debug "Stripped zinit"
        else
            log_debug "zinit already stripped or strip failed"
        fi
        
        if command_exists "upx" && upx --best --force "${initramfs_dir}/sbin/zinit" >/dev/null 2>&1 || true; then
            log_debug "UPX compressed zinit"
        else
            log_debug "UPX failed or already compressed"
        fi
        
        local final_size=$(get_file_size "${initramfs_dir}/sbin/zinit")
        log_info "✓ Copied zinit ${original_size} → ${final_size} to /sbin/zinit"
        ((copied_count++))
    else
        log_error "✗ zinit binary not found: ${zinit_binary}"
        ((missing_count++))
    fi
    
    # Copy rfs to /usr/bin
    local rfs_binary="${components_dir}/rfs/rfs"
    if [[ -f "$rfs_binary" ]]; then
        safe_mkdir "${initramfs_dir}/usr/bin"
        safe_execute cp "$rfs_binary" "${initramfs_dir}/usr/bin/rfs"
        safe_execute chmod +x "${initramfs_dir}/usr/bin/rfs"
        
        # Strip and UPX compress rfs
        local original_size=$(get_file_size "${initramfs_dir}/usr/bin/rfs")
        if strip "${initramfs_dir}/usr/bin/rfs" 2>/dev/null || true; then
            log_debug "Stripped rfs"
        else
            log_debug "rfs already stripped or strip failed"
        fi
        
        if command_exists "upx" && upx --best --force "${initramfs_dir}/usr/bin/rfs" >/dev/null 2>&1 || true; then
            log_debug "UPX compressed rfs"
        else
            log_debug "UPX failed or already compressed"
        fi
        
        local final_size=$(get_file_size "${initramfs_dir}/usr/bin/rfs")
        log_info "✓ Copied rfs ${original_size} → ${final_size} to /usr/bin/rfs"
        ((copied_count++))
    else
        log_error "✗ rfs binary not found: ${rfs_binary}"
        ((missing_count++))
    fi
    
    # Copy mycelium to /usr/bin
    local mycelium_binary="${components_dir}/mycelium/myceliumd/target/x86_64-unknown-linux-musl/release/mycelium"
    if [[ -f "$mycelium_binary" ]]; then
        safe_mkdir "${initramfs_dir}/usr/bin"
        safe_execute cp "$mycelium_binary" "${initramfs_dir}/usr/bin/mycelium"
        safe_execute chmod +x "${initramfs_dir}/usr/bin/mycelium"
        
        # Strip and UPX compress mycelium
        local original_size=$(get_file_size "${initramfs_dir}/usr/bin/mycelium")
        if strip "${initramfs_dir}/usr/bin/mycelium" 2>/dev/null || true; then
            log_debug "Stripped mycelium"
        else
            log_debug "mycelium already stripped or strip failed"
        fi
        
        if command_exists "upx" && upx --best --force "${initramfs_dir}/usr/bin/mycelium" >/dev/null 2>&1 || true; then
            log_debug "UPX compressed mycelium"
        else
            log_debug "UPX failed or already compressed"
        fi
        
        local final_size=$(get_file_size "${initramfs_dir}/usr/bin/mycelium")
        log_info "✓ Copied mycelium ${original_size} → ${final_size} to /usr/bin/mycelium"
        ((copied_count++))
    else
        log_error "✗ mycelium binary not found: ${mycelium_binary}"
        ((missing_count++))
    fi
    
    # Copy corex to /usr/bin
    local corex_binary="${components_dir}/corex/corex"
    if [[ -f "$corex_binary" ]]; then
        safe_mkdir "${initramfs_dir}/usr/bin"
        safe_execute cp "$corex_binary" "${initramfs_dir}/usr/bin/corex"
        safe_execute chmod +x "${initramfs_dir}/usr/bin/corex"
        
        # Strip and UPX compress corex
        local original_size=$(get_file_size "${initramfs_dir}/usr/bin/corex")
        if strip "${initramfs_dir}/usr/bin/corex" 2>/dev/null || true; then
            log_debug "Stripped corex"
        else
            log_debug "corex already stripped or strip failed"
        fi
        
        if command_exists "upx" && upx --best --force "${initramfs_dir}/usr/bin/corex" >/dev/null 2>&1 || true; then
            log_debug "UPX compressed corex"
        else
            log_debug "UPX failed or already compressed"
        fi
        
        local final_size=$(get_file_size "${initramfs_dir}/usr/bin/corex")
        log_info "✓ Copied corex ${original_size} → ${final_size} to /usr/bin/corex"
        ((copied_count++))
    else
        log_error "✗ corex binary not found: ${corex_binary}"
        ((missing_count++))
    fi
    
    log_info "Component copy complete: ${copied_count} copied, ${missing_count} missing"
    
    if [[ $missing_count -gt 0 ]]; then
        log_error "Some components missing - build may have failed"
        return 1
    fi
    
    return 0
}

# Setup 2-stage module loading system with dependency resolution and firmware correlation
function initramfs_setup_modules() {
    local initramfs_dir="$1"
    local modules_conf="$2"
    local kernel_version="${3:-$(uname -r)}"
    
    section_header "Setting up 2-stage module loading with firmware correlation"
    
    if [[ ! -f "$modules_conf" ]]; then
        log_error "Modules configuration file not found: ${modules_conf}"
        return 1
    fi
    
    local modules_dir="${initramfs_dir}/lib/modules/${kernel_version}"
    safe_mkdir "$modules_dir"
    
    # Track required firmware packages
    local required_firmware=()
    
    # Create stage1 module list with dependencies and firmware
    log_info "Resolving stage1 module dependencies (critical boot modules)"
    local stage1_modules=()
    local stage1_firmware=()
    
    while IFS=: read -r stage module firmware_line; do
        if [[ "$stage" == "stage1" && -n "$module" ]]; then
            stage1_modules+=("$module")
            # Extract firmware package name (before any comment)
            local firmware=$(echo "$firmware_line" | sed 's/[[:space:]]*#.*//' | tr -d ' ')
            if [[ -n "$firmware" && "$firmware" != "none" ]]; then
                stage1_firmware+=("$firmware")
                required_firmware+=("$firmware")
                log_debug "Module $module requires firmware: $firmware"
            fi
        fi
    done < <(grep "^stage1:" "$modules_conf")
    
    local stage1_with_deps=()
    if [[ ${#stage1_modules[@]} -gt 0 ]]; then
        log_info "Resolving dependencies for ${#stage1_modules[@]} stage1 modules: ${stage1_modules[*]}"
        stage1_with_deps=($(initramfs_resolve_module_dependencies "${stage1_modules[@]}"))
        log_info "Stage1 expanded to ${#stage1_with_deps[@]} modules (including dependencies)"
    fi
    
    # Write stage1 list with dependencies
    printf '%s\n' "${stage1_with_deps[@]}" > "${modules_dir}/stage1.list"
    
    # Debug: show what's in stage1 list
    if [[ ${#stage1_with_deps[@]} -gt 50 ]]; then
        log_warn "Stage1 module list unexpectedly large: ${#stage1_with_deps[@]} modules"
        log_warn "This may indicate excessive dependency resolution"
        log_debug "First 10 stage1 modules: ${stage1_with_deps[@]:0:10}"
    fi
    
    # Create empty stage2 list (stage2 not used, only stage1)
    log_info "Creating empty stage2 list (only using stage1 modules)"
    echo > "${modules_dir}/stage2.list"
    local stage2_with_deps=()
    
    # Create firmware requirements list (remove duplicates)
    local unique_firmware=($(printf '%s\n' "${required_firmware[@]}" | sort -u))
    if [[ ${#unique_firmware[@]} -gt 0 ]]; then
        printf '%s\n' "${unique_firmware[@]}" > "${modules_dir}/required-firmware.list"
        log_info "Created firmware requirements list: ${#unique_firmware[@]} packages"
        for fw in "${unique_firmware[@]}"; do
            log_info "  Required firmware: $fw"
        done
        
        # Export for use by firmware installation
        export REQUIRED_FIRMWARE_PACKAGES="${unique_firmware[*]}"
    else
        log_info "No firmware packages required"
        export REQUIRED_FIRMWARE_PACKAGES=""
    fi
    
    # Create module loading scripts
    initramfs_create_module_scripts "$initramfs_dir" "$kernel_version"
    
    # Report final counts
    local stage1_count=${#stage1_with_deps[@]}
    local firmware_count=${#unique_firmware[@]}
    
    log_info "Module configuration complete:"
    log_info "  Stage1 (critical + deps): ${stage1_count} modules"
    log_info "  Stage2: disabled (only using stage1)"
    log_info "  Required firmware packages: ${firmware_count}"
    log_info "  Total modules: ${stage1_count}"
}

# Resolve module dependencies using proper depmod + modinfo -k approach
function initramfs_resolve_module_dependencies() {
    local modules=("$@")
    local resolved_modules=()
    local processed_modules=()
    
    log_debug "Resolving dependencies recursively using modinfo -k: ${modules[*]}"
    
    # Use full kernel version from build process
    local kernel_version="${KERNEL_FULL_VERSION:-${FULL_KERNEL_VERSION:-}}"
    if [[ -z "$kernel_version" ]]; then
        source "${LIB_SCRIPT_DIR}/kernel.sh"
        kernel_version=$(kernel_get_full_version "${KERNEL_VERSION:-6.12.44}" "${PROJECT_ROOT}/config/kernel.config")
    fi
    
    # Set up container modules directory for proper modinfo -k usage
    local container_modules_path="/lib/modules/${kernel_version}"
    local source_modules_path="${PROJECT_ROOT}/initramfs/lib/modules/${kernel_version}"
    
    log_debug "Using kernel version: ${kernel_version}"
    log_debug "Source modules: ${source_modules_path}"
    log_debug "Container modules: ${container_modules_path}"
    
    # Verify source modules exist
    if [[ ! -d "$source_modules_path" ]]; then
        log_warn "Built modules not found at: $source_modules_path"
        log_warn "Falling back to simple module list"
        printf '%s\n' "${modules[@]}"
        return 0
    fi
    
    # Copy modules to container /lib/modules for proper modinfo -k usage
    log_info "Setting up container modules for dependency resolution"
    safe_mkdir "$container_modules_path"
    safe_execute cp -r "${source_modules_path}"/* "$container_modules_path/"
    
    # Run depmod to build dependency database
    log_info "Running depmod -av for dependency resolution"
    safe_execute depmod -av "$kernel_version" >/dev/null
    
    # Recursive dependency resolution function
    function resolve_single_module() {
        local module="$1"
        local depth="${2:-0}"
        local max_depth=5  # Allow deeper recursion for proper resolution
        
        # Skip if already processed
        if printf '%s\n' "${processed_modules[@]}" | grep -q "^${module}$"; then
            return 0
        fi
        
        # Skip if too deep (prevent infinite loops)
        if [[ $depth -gt $max_depth ]]; then
            log_debug "Skipping $module - depth limit ($max_depth) reached"
            return 0
        fi
        
        processed_modules+=("$module")
        
        # Use modinfo -k for proper dependency resolution
        local deps=()
        if command_exists "modinfo"; then
            local depends_line
            # Use modinfo -k to query the properly indexed modules
            if depends_line=$(modinfo -k "$kernel_version" "$module" 2>/dev/null | grep '^depends:' | head -1 | cut -d: -f2- | tr -d ' '); then
                if [[ -n "$depends_line" && "$depends_line" != "-" ]]; then
                    IFS=',' read -ra deps <<< "$depends_line"
                    if [[ ${#deps[@]} -gt 0 ]]; then
                        log_debug "Module $module (depth $depth) depends on: ${deps[*]}"
                    fi
                fi
            else
                # Check if module is built-in
                if modinfo -k "$kernel_version" "$module" >/dev/null 2>&1; then
                    log_debug "Module $module found but no dependencies"
                else
                    log_debug "Module $module not found - likely built-in"
                fi
            fi
        fi
        
        # Recursively resolve dependencies first (depth-first)
        for dep in "${deps[@]}"; do
            if [[ -n "$dep" ]]; then
                resolve_single_module "$dep" $((depth + 1))
            fi
        done
        
        # Add current module to resolved list (after dependencies for correct load order)
        resolved_modules+=("$module")
    }
    
    # Process all requested modules
    for module in "${modules[@]}"; do
        resolve_single_module "$module"
    done
    
    # Remove duplicates while preserving dependency order
    local unique_modules=()
    local seen_modules=()
    
    for module in "${resolved_modules[@]}"; do
        if ! printf '%s\n' "${seen_modules[@]}" | grep -q "^${module}$"; then
            unique_modules+=("$module")
            seen_modules+=("$module")
        fi
    done
    
    log_debug "Resolved ${#unique_modules[@]} unique modules with recursive dependencies"
    printf '%s\n' "${unique_modules[@]}"
}


# Create module loading scripts for zinit
function initramfs_create_module_scripts() {
    local initramfs_dir="$1"
    local kernel_version="$2"
    
    log_info "Creating module loading scripts"
    
    safe_mkdir "${initramfs_dir}/etc/zinit/init"
    
    # Stage1 module loading script (critical modules)
    cat > "${initramfs_dir}/etc/zinit/init/stage1-modules.sh" << 'EOF'
#!/bin/sh
# Stage1 module loading - critical boot modules

KERNEL_VERSION=$(uname -r)
STAGE1_LIST="/lib/modules/${KERNEL_VERSION}/stage1.list"

echo "Loading stage1 modules (critical boot)"

if [ -f "$STAGE1_LIST" ]; then
    while read -r module; do
        if [ -n "$module" ] && [ "$module" != "#"* ]; then
            echo "Loading critical module: $module"
            modprobe "$module" 2>/dev/null || echo "Warning: Failed to load $module"
        fi
    done < "$STAGE1_LIST"
else
    echo "Warning: Stage1 module list not found: $STAGE1_LIST"
fi

echo "Stage1 module loading complete"
EOF

    # Stage2 module loading script (extended hardware)
    cat > "${initramfs_dir}/etc/zinit/init/stage2-modules.sh" << 'EOF'
#!/bin/sh
# Stage2 module loading - extended hardware support

KERNEL_VERSION=$(uname -r)
STAGE2_LIST="/lib/modules/${KERNEL_VERSION}/stage2.list"

echo "Loading stage2 modules (extended hardware)"

if [ -f "$STAGE2_LIST" ]; then
    while read -r module; do
        if [ -n "$module" ] && [ "$module" != "#"* ]; then
            echo "Loading hardware module: $module"
            modprobe "$module" 2>/dev/null || echo "Info: Module $module not available"
        fi
    done < "$STAGE2_LIST"
else
    echo "Warning: Stage2 module list not found: $STAGE2_LIST"
fi

echo "Stage2 module loading complete"
EOF

    # Make scripts executable
    safe_execute chmod 755 "${initramfs_dir}/etc/zinit/init/stage1-modules.sh"
    safe_execute chmod 755 "${initramfs_dir}/etc/zinit/init/stage2-modules.sh"
    
    log_info "Module loading scripts created"
}

# Strip and UPX compress all binaries for maximum size optimization
function initramfs_strip_and_upx() {
    local initramfs_dir="$1"
    
    section_header "Stripping and UPX compressing binaries"
    
    local stripped_count=0
    local upx_count=0
    local failed_strip=0
    local failed_upx=0
    
    # Find and process all executable files
    log_info "Processing executable files..."
    
    while IFS= read -r -d '' file; do
        # Check if it's a valid ELF executable
        if file "$file" | grep -q "ELF.*executable"; then
            log_debug "Processing executable: $file"
            
            # Strip debug symbols (ignore already stripped files)
            if strip "$file" 2>/dev/null || true; then
                if [[ $? -eq 0 ]]; then
                    ((stripped_count++))
                    log_debug "Stripped: $file"
                else
                    ((failed_strip++))
                    log_debug "Already stripped or failed: $file"
                fi
            fi
            
            # UPX compress (best compression) - skip if already compressed or fails
            if command_exists "upx"; then
                if upx --best --force "$file" >/dev/null 2>&1 || true; then
                    if [[ $? -eq 0 ]]; then
                        ((upx_count++))
                        log_debug "UPX compressed: $file"
                    else
                        ((failed_upx++))
                        log_debug "UPX failed or already compressed: $file"
                    fi
                fi
            else
                log_debug "UPX not available, skipping compression for: $file"
                ((failed_upx++))
            fi
        fi
    done < <(find "$initramfs_dir" -type f -executable -print0)
    
    # Process shared libraries
    log_info "Processing shared libraries..."
    
    local lib_stripped=0
    local lib_failed=0
    
    while IFS= read -r -d '' file; do
        if file "$file" | grep -q "ELF.*shared object"; then
            log_debug "Processing library: $file"
            
            # Strip libraries (more conservative - keep function symbols)
            if strip --strip-unneeded "$file" 2>/dev/null || true; then
                if [[ $? -eq 0 ]]; then
                    ((lib_stripped++))
                    log_debug "Stripped library: $file"
                else
                    ((lib_failed++))
                    log_debug "Library already stripped or failed: $file"
                fi
            fi
        fi
    done < <(find "$initramfs_dir" -name "*.so*" -type f -print0)
    
    # Summary
    log_info "Binary optimization complete:"
    log_info "  Executables stripped: ${stripped_count} (${failed_strip} failed)"
    log_info "  Executables UPX compressed: ${upx_count} (${failed_upx} failed)"
    log_info "  Libraries stripped: ${lib_stripped} (${lib_failed} failed)"
    
    # Calculate space savings
    local total_size=$(du -sb "$initramfs_dir" 2>/dev/null | cut -f1 || echo "0")
    local total_mb=$((total_size / 1024 / 1024))
    log_info "Total initramfs size after optimization: ${total_mb}MB"
}

# Create final initramfs.cpio.xz archive
function initramfs_create_cpio() {
    local initramfs_dir="$1"
    local output_file="$2"
    local compression="${3:-$INITRAMFS_COMPRESSION}"
    
    section_header "Creating initramfs.cpio.${compression}"
    
    if [[ ! -d "$initramfs_dir" ]]; then
        log_error "Initramfs directory not found: ${initramfs_dir}"
        return 1
    fi
    
    # Ensure output directory exists
    local output_dir=$(dirname "$output_file")
    safe_mkdir "$output_dir"
    
    # Remove any existing output file
    safe_execute rm -f "$output_file"
    
    log_info "Source directory: ${initramfs_dir}"
    log_info "Output file: ${output_file}"
    log_info "Compression: ${compression}"
    
    # Change to initramfs directory for relative paths
    safe_execute cd "$initramfs_dir"
    
    case "$compression" in
        "xz")
            log_info "Creating XZ compressed CPIO archive"
            safe_execute find . -print0 | cpio -o -H newc -0 | xz -${XZ_COMPRESSION_LEVEL} --check=crc32 > "$output_file"
            ;;
        "gzip"|"gz")
            log_info "Creating gzip compressed CPIO archive"
            safe_execute find . -print0 | cpio -o -H newc -0 | gzip -9 > "$output_file"
            ;;
        "zstd")
            log_info "Creating zstd compressed CPIO archive"
            safe_execute find . -print0 | cpio -o -H newc -0 | zstd -19 > "$output_file"
            ;;
        "none"|"uncompressed")
            log_info "Creating uncompressed CPIO archive"
            safe_execute find . -print0 | cpio -o -H newc -0 > "$output_file"
            ;;
        *)
            log_error "Unsupported compression format: ${compression}"
            return 1
            ;;
    esac
    
    # Verify output file was created
    if [[ ! -f "$output_file" ]]; then
        log_error "Failed to create initramfs archive: ${output_file}"
        return 1
    fi
    
    # Report final size
    local final_size=$(get_file_size "$output_file")
    local uncompressed_size=$(du -sh "$initramfs_dir" | cut -f1)
    
    log_info "Initramfs creation complete:"
    log_info "  Uncompressed size: ${uncompressed_size}"
    log_info "  Final archive size: ${final_size}"
    log_info "  Archive: ${output_file}"
}

# Validate initramfs contents
function initramfs_validate() {
    local initramfs_dir="$1"
    
    section_header "Validating initramfs contents"
    
    local errors=0
    
    # Check essential files and directories
    local essential_items=(
        "init"
        "sbin/init"
        "sbin/zinit"
        "bin/busybox"
        "etc/zinit"
        "lib"
        "usr/bin"
        "var"
        "tmp"
        "proc"
        "sys"
        "dev"
    )
    
    for item in "${essential_items[@]}"; do
        if [[ ! -e "${initramfs_dir}/${item}" ]]; then
            log_error "Missing essential item: ${item}"
            ((errors++))
        else
            log_debug "Found: ${item}"
        fi
    done
    
    # Check that initramfs /init script exists
    if [[ -x "${initramfs_dir}/init" ]]; then
        log_info "✓ Initramfs /init script found"
    else
        log_error "✗ Initramfs /init script missing"
        ((errors++))
    fi
    
    # Check that /sbin/init is a script
    if [[ -f "${initramfs_dir}/sbin/init" && -x "${initramfs_dir}/sbin/init" ]]; then
        log_info "✓ /sbin/init script found"
    else
        log_error "✗ /sbin/init is missing or not executable"
        ((errors++))
    fi
    
    # Check zinit configuration
    if [[ -f "${initramfs_dir}/etc/zinit/zinit.conf" ]]; then
        log_info "✓ zinit configuration found"
    else
        log_error "✗ zinit configuration missing"
        ((errors++))
    fi
    
    # Verify no OpenRC remnants
    local openrc_check=(
        "etc/init.d"
        "etc/runlevels"
        "sbin/openrc"
    )
    
    for path in "${openrc_check[@]}"; do
        if [[ -e "${initramfs_dir}/${path}" ]]; then
            log_warn "OpenRC remnant found: ${path}"
        fi
    done
    
    # Check component binaries
    local component_binaries=(
        "usr/bin/rfs"
        "usr/bin/mycelium"
        "usr/bin/corex"
    )
    
    for binary in "${component_binaries[@]}"; do
        if [[ -x "${initramfs_dir}/${binary}" ]]; then
            log_info "✓ Component binary: ${binary}"
        else
            log_warn "Component binary missing or not executable: ${binary}"
        fi
    done
    
    if [[ $errors -eq 0 ]]; then
        log_info "Initramfs validation passed"
        return 0
    else
        log_error "Initramfs validation failed with ${errors} errors"
        return 1
    fi
}

# Test initramfs archive integrity
function initramfs_test_archive() {
    local archive_file="$1"
    
    section_header "Testing initramfs archive integrity"
    
    if [[ ! -f "$archive_file" ]]; then
        log_error "Archive file not found: ${archive_file}"
        return 1
    fi
    
    # Test based on file extension
    case "$archive_file" in
        *.xz)
            log_info "Testing XZ archive integrity"
            safe_execute xz -t "$archive_file"
            ;;
        *.gz)
            log_info "Testing gzip archive integrity"
            safe_execute gzip -t "$archive_file"
            ;;
        *.zst)
            log_info "Testing zstd archive integrity"
            safe_execute zstd -t "$archive_file"
            ;;
        *.cpio)
            log_info "Testing CPIO archive integrity"
            safe_execute cpio -t < "$archive_file" >/dev/null
            ;;
        *)
            log_warn "Unknown archive format, skipping integrity test"
            return 0
            ;;
    esac
    
    log_info "Archive integrity test passed"
}

# Copy resolved modules from container to initramfs
function initramfs_copy_resolved_modules() {
    local initramfs_dir="$1"
    local kernel_version="${2:-$(uname -r)}"
    
    log_info "ENTRY: initramfs_copy_resolved_modules called with:"
    log_info "  initramfs_dir: ${initramfs_dir}"
    log_info "  kernel_version: ${kernel_version}"
    log_info "  CONTAINER_MODULES_PATH: ${CONTAINER_MODULES_PATH:-not-set}"
    
    section_header "Copying Resolved Modules to Initramfs"
    
    local container_modules_path="${CONTAINER_MODULES_PATH:-/lib/modules/${kernel_version}}"
    local initramfs_modules_dir="${initramfs_dir}/lib/modules/${kernel_version}"
    
    log_info "Container modules path: ${container_modules_path}"
    log_info "Initramfs modules dir: ${initramfs_modules_dir}"
    
    if [[ ! -d "$container_modules_path" ]]; then
        log_error "Container modules not found: $container_modules_path"
        log_info "Available module directories:"
        ls -la /lib/modules/ 2>/dev/null || log_error "No /lib/modules directory"
        return 1
    fi
    
    log_info "Container modules directory exists, proceeding..."
    
    # Create initramfs modules directory
    safe_mkdir "$initramfs_modules_dir"
    
    # Read the stage1 module list (only using stage1, no stage2)
    local stage1_modules=()
    
    if [[ -f "${initramfs_modules_dir}/stage1.list" ]]; then
        while IFS= read -r module; do
            [[ -n "$module" ]] && stage1_modules+=("$module")
        done < "${initramfs_modules_dir}/stage1.list"
    fi
    
    local all_modules=("${stage1_modules[@]}")
    
    if [[ ${#all_modules[@]} -eq 0 ]]; then
        log_warn "No modules to copy (empty stage1 module list)"
        return 0
    fi
    
    log_info "Copying ${#all_modules[@]} stage1 modules to initramfs"
    log_info "Stage1 modules: ${#stage1_modules[@]} (stage2 disabled)"
    
    # Debug: show what we're copying if the list is reasonable
    if [[ ${#all_modules[@]} -lt 200 ]]; then
        log_debug "Modules to copy: ${all_modules[*]}"
    else
        log_warn "Module list is very large (${#all_modules[@]}), showing sample:"
        log_debug "First 10: ${all_modules[@]:0:10}"
        log_debug "Last 10: ${all_modules[@]:-10}"
    fi
    
    # Copy essential module metadata first
    local essential_files=(
        "modules.order"
        "modules.builtin"
        "modules.builtin.modinfo"
        "modules.dep"
        "modules.dep.bin"
        "modules.symbols"
        "modules.symbols.bin"
        "modules.alias"
        "modules.alias.bin"
        "modules.devname"
    )
    
    for file in "${essential_files[@]}"; do
        if [[ -f "${container_modules_path}/${file}" ]]; then
            safe_execute cp "${container_modules_path}/${file}" "${initramfs_modules_dir}/"
            log_debug "Copied module metadata: $file"
        fi
    done
    
    # Copy actual module files
    local copied_count=0
    local failed_count=0
    
    for module in "${all_modules[@]}"; do
        local module_found=false
        
        # Try to find module file
        local module_file=$(find "$container_modules_path" -name "${module}.ko*" -type f | head -1)
        
        if [[ -n "$module_file" && -f "$module_file" ]]; then
            # Preserve directory structure
            local rel_path="${module_file#${container_modules_path}/}"
            local target_dir="${initramfs_modules_dir}/$(dirname "$rel_path")"
            
            safe_mkdir "$target_dir"
            safe_execute cp "$module_file" "${initramfs_modules_dir}/${rel_path}"
            
            log_debug "Copied module: $module from $rel_path"
            ((copied_count++))
            module_found=true
        fi
        
        if [[ "$module_found" == "false" ]]; then
            log_warn "Module not found in container: $module"
            ((failed_count++))
        fi
    done
    
    log_info "Module copy complete: ${copied_count} copied, ${failed_count} failed"
    
    # Run depmod in initramfs context (chroot-like) for final dependency resolution
    if command_exists "depmod" && [[ $copied_count -gt 0 ]]; then
        log_info "Running depmod in initramfs context for final module database"
        safe_execute depmod -av -b "$initramfs_dir" "$kernel_version"
        
        # Verify modules.dep was created properly
        if [[ -f "${initramfs_modules_dir}/modules.dep" ]]; then
            local dep_count=$(wc -l < "${initramfs_modules_dir}/modules.dep")
            log_info "✓ Initramfs modules.dep created with ${dep_count} entries"
        else
            log_warn "modules.dep not created in initramfs"
        fi
    fi
    
    return 0
}

# Export all functions at the end after they're all defined
export -f initramfs_setup_zinit initramfs_setup_modules initramfs_resolve_module_dependencies
export -f initramfs_install_init_script initramfs_copy_components initramfs_create_module_scripts
export -f initramfs_strip_and_upx initramfs_create_cpio initramfs_validate initramfs_test_archive
export -f initramfs_copy_resolved_modules