#!/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
    
    # Remove existing init (if any) and replace with zinit
    log_info "Replacing system init with zinit"
    safe_execute rm -f "${initramfs_dir}/sbin/init"
    safe_execute ln -sf zinit "${initramfs_dir}/sbin/init"
    
    # Copy zinit configuration
    log_info "Installing zinit configuration"
    safe_mkdir "${initramfs_dir}/etc/zinit"
    safe_copy "${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"
}

# Setup 2-stage module loading system
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"
    
    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"
    
    # Create stage1 module list (critical boot modules)
    log_info "Creating stage1 module list (critical boot modules)"
    grep "^stage1:" "$modules_conf" | cut -d: -f2 > "${modules_dir}/stage1.list"
    
    # Create stage2 module list (extended hardware support)
    log_info "Creating stage2 module list (extended hardware support)"
    grep "^stage2:" "$modules_conf" | cut -d: -f2 > "${modules_dir}/stage2.list"
    
    # Create module loading scripts
    initramfs_create_module_scripts "$initramfs_dir" "$kernel_version"
    
    # Count modules
    local stage1_count=$(wc -l < "${modules_dir}/stage1.list" 2>/dev/null || echo 0)
    local stage2_count=$(wc -l < "${modules_dir}/stage2.list" 2>/dev/null || echo 0)
    
    log_info "Module configuration complete:"
    log_info "  Stage1 (critical): ${stage1_count} modules"
    log_info "  Stage2 (extended): ${stage2_count} 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
            if strip "$file" 2>/dev/null; then
                ((stripped_count++))
                log_debug "Stripped: $file"
            else
                ((failed_strip++))
                log_debug "Failed to strip: $file"
            fi
            
            # UPX compress (best compression)
            if upx --best --force "$file" 2>/dev/null; then
                ((upx_count++))
                log_debug "UPX compressed: $file"
            else
                ((failed_upx++))
                log_debug "Failed to UPX: $file"
            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; then
                ((lib_stripped++))
                log_debug "Stripped library: $file"
            else
                ((lib_failed++))
                log_debug "Failed to strip library: $file"
            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=(
        "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 init is properly linked to zinit
    if [[ -L "${initramfs_dir}/sbin/init" ]]; then
        local link_target=$(readlink "${initramfs_dir}/sbin/init")
        if [[ "$link_target" == "zinit" ]]; then
            log_info "✓ /sbin/init correctly linked to zinit"
        else
            log_error "✗ /sbin/init linked to wrong target: ${link_target}"
            ((errors++))
        fi
    else
        log_error "✗ /sbin/init is not a symbolic link"
        ((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"
}

# Export functions
export -f initramfs_setup_zinit initramfs_setup_modules initramfs_create_module_scripts
export -f initramfs_strip_and_upx initramfs_create_cpio
export -f initramfs_validate initramfs_test_archive