#!/bin/bash
set -e

INITRAMFS_ROOT="/build/initramfs"
CONFIG_DIR="/build/configs"
OUTPUT_DIR="/build/output"

echo "[+] Setting up initramfs structure..."

# Copy init script
echo "  Installing init script..."
cp "$CONFIG_DIR/init" "$INITRAMFS_ROOT/init"
chmod +x "$INITRAMFS_ROOT/init"

# Copy zinit configuration (mounted from existing config)
echo "  Installing zinit configuration..."
echo "[DEBUG] Checking zinit config sources..."
echo "[DEBUG] /mnt/zinit exists: $([ -d "/mnt/zinit" ] && echo 'YES' || echo 'NO')"
echo "[DEBUG] /mnt/zinit contents: $(ls -la /mnt/zinit 2>/dev/null | wc -l) files"
echo "[DEBUG] $CONFIG_DIR/zinit exists: $([ -d "$CONFIG_DIR/zinit" ] && echo 'YES' || echo 'NO')"
echo "[DEBUG] $CONFIG_DIR/zinit contents: $(ls -la $CONFIG_DIR/zinit 2>/dev/null | wc -l) files"

mkdir -p "$INITRAMFS_ROOT/etc/zinit"
if [ -d "/mnt/zinit" ]; then
    cp -r "/mnt/zinit/"* "$INITRAMFS_ROOT/etc/zinit/"
    echo "    Copied zinit configs from /mnt/zinit"
    echo "[DEBUG] Copied $(ls -la $INITRAMFS_ROOT/etc/zinit 2>/dev/null | wc -l) zinit config files"
elif [ -d "$CONFIG_DIR/zinit" ]; then
    cp -r "$CONFIG_DIR/zinit/"* "$INITRAMFS_ROOT/etc/zinit/"
    echo "    Copied zinit configs from $CONFIG_DIR/zinit"
    echo "[DEBUG] Copied $(ls -la $INITRAMFS_ROOT/etc/zinit 2>/dev/null | wc -l) zinit config files"
else
    echo "    Warning: zinit config directory not found"
    echo "[DEBUG] ERROR: No zinit config source available - this will likely cause runtime failures"
fi

# Create debug script for compatibility
echo "  Creating debug script..."
cat > "$INITRAMFS_ROOT/init-debug" << 'EOF'
#!/bin/bash
# Debug file injection script (Alpine version)
echo "[+] debug mode enabled"

# Check for debug device (same as original)
if [ -b /dev/sda1 ]; then
    echo "  debug device found: /dev/sda1"
    
    # Create temporary mount point
    mkdir -p /mnt/debug
    
    # Try to mount as vfat (same as original)
    if mount -t vfat -o ro /dev/sda1 /mnt/debug 2>/dev/null; then
        echo "  mounted debug device"
        
        # Check for debug marker file
        if [ -f /mnt/debug/.zero-os-debug ]; then
            echo "  debug marker found, copying files..."
            
            # Copy all files from debug device to tmpfs root
            target="/mnt/root"
            if [ -d "$target" ]; then
                cp -rv /mnt/debug/* "$target/" 2>/dev/null || true
                echo "  debug files copied to tmpfs"
            else
                echo "  warning: tmpfs root not ready yet"
            fi
        else
            echo "  no debug marker (.zero-os-debug) found"
        fi
        
        # Unmount debug device
        umount /mnt/debug 2>/dev/null || true
    else
        echo "  failed to mount debug device"
    fi
    
    # Cleanup
    rmdir /mnt/debug 2>/dev/null || true
else
    echo "  no debug device found"
fi

echo "[+] debug initialization complete"
EOF

chmod +x "$INITRAMFS_ROOT/init-debug"

# Ensure compiled Zero-OS binaries are installed
echo "  Checking Zero-OS binaries..."
echo "[DEBUG] Checking for compiled binaries..."

# Check if binaries exist and copy them if needed
BINARIES_FOUND=true

# Check zinit (critical for boot)
if [ ! -f "$INITRAMFS_ROOT/sbin/zinit" ]; then
    if [ -f "/build/output/zinit" ]; then
        echo "    Installing zinit from output directory..."
        mkdir -p "$INITRAMFS_ROOT/sbin"
        cp "/build/output/zinit" "$INITRAMFS_ROOT/sbin/"
        chmod +x "$INITRAMFS_ROOT/sbin/zinit"
    else
        echo "    ERROR: zinit binary not found! Build may have failed."
        BINARIES_FOUND=false
    fi
else
    echo "    ✓ zinit found at $INITRAMFS_ROOT/sbin/zinit"
fi

# Check other binaries
for binary in "usr/bin/rfs" "usr/bin/mycelium" "usr/bin/corex"; do
    if [ ! -f "$INITRAMFS_ROOT/$binary" ]; then
        basename_binary=$(basename "$binary")
        if [ -f "/build/output/$basename_binary" ]; then
            echo "    Installing $basename_binary from output directory..."
            mkdir -p "$INITRAMFS_ROOT/$(dirname $binary)"
            cp "/build/output/$basename_binary" "$INITRAMFS_ROOT/$binary"
            chmod +x "$INITRAMFS_ROOT/$binary"
        else
            echo "    Warning: $basename_binary not found in output directory"
        fi
    else
        echo "    ✓ $basename_binary found"
    fi
done

if [ "$BINARIES_FOUND" = "false" ]; then
    echo "[DEBUG] ERROR: Critical binaries missing - initramfs will not boot properly"
    echo "[DEBUG] Make sure to run compile-components.sh before setup-initramfs.sh"
fi

# Install kernel modules for network interfaces
echo "  Installing essential network kernel modules..."
KERNEL_VERSION=$(cat /build/configs/kernel-version)
MODULES_DIR="$INITRAMFS_ROOT/lib/modules/${KERNEL_VERSION}-Zero-OS"
mkdir -p "$MODULES_DIR/kernel/drivers/net/ethernet"
mkdir -p "$MODULES_DIR/kernel/drivers/net"
mkdir -p "$MODULES_DIR/kernel/drivers"

KERNEL_BUILD_DIR="/build/kernel/linux-${KERNEL_VERSION}"
if [ -d "$KERNEL_BUILD_DIR" ]; then
    echo "    Copying essential ethernet drivers from kernel build..."
    
    # Copy essential network driver modules and their dependencies
    find "$KERNEL_BUILD_DIR" -name "*.ko" | grep -E "(e1000|r8169|bnx2|tg3|virtio|igb|ixgbe|i40e|ice|atl1|alx|8139|b44|mii|mdio|libphy|net_failover|failover)" | while read ko; do
        if [ -f "$ko" ]; then
            cp "$ko" "$MODULES_DIR/kernel/drivers/net/ethernet/" 2>/dev/null || \
            cp "$ko" "$MODULES_DIR/kernel/drivers/net/" 2>/dev/null || \
            cp "$ko" "$MODULES_DIR/kernel/drivers/" 2>/dev/null || \
            cp "$ko" "$MODULES_DIR/kernel/net/" 2>/dev/null
        fi
    done
    
    # Also copy any net_failover modules from net/ directory
    find "$KERNEL_BUILD_DIR" -path "*/net/*" -name "*failover*.ko" | while read ko; do
        if [ -f "$ko" ]; then
            mkdir -p "$MODULES_DIR/kernel/net/"
            cp "$ko" "$MODULES_DIR/kernel/net/"
        fi
    done
    
    # Count actual installed modules
    actual_modules=$(find "$MODULES_DIR" -name "*.ko" 2>/dev/null | wc -l)
    echo "    Installed $actual_modules network driver modules"
    
    # Copy essential kernel module metadata files from kernel build
    echo "    Copying kernel module metadata files..."
    if [ -f "$KERNEL_BUILD_DIR/modules.order" ]; then
        cp "$KERNEL_BUILD_DIR/modules.order" "$MODULES_DIR/"
        echo "      ✓ modules.order copied"
    else
        echo "      Warning: modules.order not found in kernel build"
        touch "$MODULES_DIR/modules.order"
    fi
    
    if [ -f "$KERNEL_BUILD_DIR/modules.builtin" ]; then
        cp "$KERNEL_BUILD_DIR/modules.builtin" "$MODULES_DIR/"
        echo "      ✓ modules.builtin copied"
    else
        echo "      Warning: modules.builtin not found in kernel build"
        touch "$MODULES_DIR/modules.builtin"
    fi
    
    if [ -f "$KERNEL_BUILD_DIR/modules.builtin.modinfo" ]; then
        cp "$KERNEL_BUILD_DIR/modules.builtin.modinfo" "$MODULES_DIR/"
        echo "      ✓ modules.builtin.modinfo copied"
    else
        echo "      Warning: modules.builtin.modinfo not found in kernel build"
        touch "$MODULES_DIR/modules.builtin.modinfo"
    fi
    
    # Create basic module dependency files (these will be updated by depmod)
    touch "$MODULES_DIR/modules.dep"
    touch "$MODULES_DIR/modules.alias"
    touch "$MODULES_DIR/modules.symbols"
    
    echo "    Module metadata files installed"
else
    echo "    Warning: Kernel build directory not found at $KERNEL_BUILD_DIR"
fi

# Copy other system configurations
echo "  Setting up system configuration..."

# Ensure proper hostname
echo "zero-os" > "$INITRAMFS_ROOT/etc/hostname"

# Basic hosts file
cat > "$INITRAMFS_ROOT/etc/hosts" << EOF
127.0.0.1	localhost zero-os
::1		localhost zero-os
EOF

# Minimal passwd file (compatible with existing)
if [ ! -f "$INITRAMFS_ROOT/etc/passwd" ]; then
    cat > "$INITRAMFS_ROOT/etc/passwd" << EOF
root:x:0:0:root:/root:/bin/sh
EOF
fi

# Minimal group file
if [ ! -f "$INITRAMFS_ROOT/etc/group" ]; then
    cat > "$INITRAMFS_ROOT/etc/group" << EOF
root:x:0:
EOF
fi

# Set up profile for environment
cat > "$INITRAMFS_ROOT/etc/profile" << EOF
export PATH=/sbin:/bin:/usr/sbin:/usr/bin
export HOME=/root
export TERM=linux
EOF

# Set up environment file for PATH
cat > "$INITRAMFS_ROOT/etc/environment" << EOF
PATH=/sbin:/bin:/usr/sbin:/usr/bin
EOF

# Ensure SSH directory exists
mkdir -p "$INITRAMFS_ROOT/etc/ssh"
mkdir -p "$INITRAMFS_ROOT/root/.ssh"

# Set proper permissions
chmod 700 "$INITRAMFS_ROOT/root"
chmod 700 "$INITRAMFS_ROOT/root/.ssh" 2>/dev/null || true

# Create necessary run directories
mkdir -p "$INITRAMFS_ROOT/run"
mkdir -p "$INITRAMFS_ROOT/var"

# Ensure terminfo is available for better terminal support
if [ -d "$INITRAMFS_ROOT/usr/share/terminfo" ]; then
    echo "  terminfo available"
else
    echo "  warning: terminfo not found"
fi

# Set up symlinks (similar to current system)
echo "  Setting up symlinks..."
cd "$INITRAMFS_ROOT"

# Ensure lib symlinks
ln -sf usr/lib lib || true
ln -sf usr/lib lib64 || true

# Ensure var/run symlink (standard Linux practice: /var/run -> /run)
cd "$INITRAMFS_ROOT/var"
rm -rf run
ln -sf ../run run || true
cd "$INITRAMFS_ROOT"

# Create initramfs archive
echo "[+] Creating initramfs archive..."
cd "$INITRAMFS_ROOT"

# Create the cpio archive
find . | cpio -H newc -o | xz --check=crc32 --x86 --lzma2=dict=1MiB > "$OUTPUT_DIR/initramfs.cpio.xz"

if [ -f "$OUTPUT_DIR/initramfs.cpio.xz" ]; then
    size=$(stat -c%s "$OUTPUT_DIR/initramfs.cpio.xz")
    echo "[+] Initramfs created: $OUTPUT_DIR/initramfs.cpio.xz ($size bytes)"
else
    echo "[-] Error: Failed to create initramfs archive"
    exit 1
fi

# Also create uncompressed version for debugging
find . | cpio -H newc -o > "$OUTPUT_DIR/initramfs.cpio"

echo "[+] Initramfs setup completed successfully"