#!/bin/bash
set -e

KERNEL_DIR="/build/kernel"
OUTPUT_DIR="/build/output"
CONFIG_DIR="/build/configs"

# Parse command line arguments
KERNEL_ONLY=false
EMBED_ONLY=false

while [ $# -gt 0 ]; do
    case $1 in
        --kernel-only)
            KERNEL_ONLY=true
            shift
            ;;
        --embed-only)
            EMBED_ONLY=true
            shift
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

echo "[+] Building custom Zero-OS kernel..."
echo "    Mode: $([ "$KERNEL_ONLY" = "true" ] && echo "kernel-only" || [ "$EMBED_ONLY" = "true" ] && echo "embed-only" || echo "full-build")"

# DIAGNOSTIC: Check prerequisites
echo "[DEBUG] Checking kernel build prerequisites..."
echo "[DEBUG] Initramfs file exists: $([ -f "$OUTPUT_DIR/initramfs.cpio.xz" ] && echo 'YES' || echo 'NO')"
echo "[DEBUG] Kernel config exists: $([ -f "$CONFIG_DIR/kernel-config-generic" ] && echo 'YES' || echo 'NO')"

# Only check for initramfs if we're doing embedding
if [ "$KERNEL_ONLY" != "true" ] && [ ! -f "$OUTPUT_DIR/initramfs.cpio.xz" ]; then
    echo "[DEBUG] ERROR: Initramfs not found - kernel build will fail"
    exit 1
fi

if [ ! -f "$CONFIG_DIR/kernel-config-generic" ]; then
    echo "[DEBUG] ERROR: Kernel config not found - kernel build will fail" 
    exit 1
fi

echo "[DEBUG] Initramfs size: $(stat -c%s "$OUTPUT_DIR/initramfs.cpio.xz") bytes"

# Read kernel version from config
if [ -f "$CONFIG_DIR/kernel-version" ]; then
    KERNEL_VERSION=$(cat "$CONFIG_DIR/kernel-version" | tr -d '\n')
    echo "  Using kernel version from config: $KERNEL_VERSION"
else
    echo "  Error: kernel-version config file not found" >&2
    exit 1
fi

KERNEL_MAJOR=$(echo $KERNEL_VERSION | cut -d. -f1)
KERNEL_URL="https://cdn.kernel.org/pub/linux/kernel/v${KERNEL_MAJOR}.x/linux-${KERNEL_VERSION}.tar.xz"

echo "  Kernel version: $KERNEL_VERSION"
echo "  Download URL: $KERNEL_URL"

# Create kernel build directory
echo "  Setting up kernel build environment..."
mkdir -p "$KERNEL_DIR"
cd "$KERNEL_DIR"

# Install kernel build dependencies
echo "  Installing kernel build dependencies..."
apk add --no-cache build-base bc bison flex openssl-dev elfutils-dev \
    perl linux-headers cpio xz wget

# Download kernel source if not already present
if [ ! -f "linux-${KERNEL_VERSION}.tar.xz" ]; then
    echo "  Downloading Linux kernel ${KERNEL_VERSION}..."
    wget -q --show-progress "$KERNEL_URL"
fi

# Extract kernel source if not already extracted
if [ ! -d "linux-${KERNEL_VERSION}" ]; then
    echo "  Extracting kernel source..."
    tar -xf "linux-${KERNEL_VERSION}.tar.xz"
fi

cd "linux-${KERNEL_VERSION}"

# Copy kernel config
echo "  Setting up kernel configuration..."
cp "$CONFIG_DIR/kernel-config-generic" .config

# Handle different build modes
if [ "$KERNEL_ONLY" = "true" ]; then
    echo "  Building kernel only (no initramfs embedding)..."
    
    # Disable initramfs in config for kernel-only build
    sed -i 's/CONFIG_INITRAMFS_SOURCE="..\/..\/root"/CONFIG_INITRAMFS_SOURCE=""/' .config
    
    # Update config and build kernel
    make olddefconfig
    make -j$(nproc) bzImage
    make -j$(nproc) modules
    
    echo "  Kernel built successfully (cached for later embedding)"
    exit 0
    
elif [ "$EMBED_ONLY" = "true" ]; then
    echo "  Embedding initramfs into existing kernel..."
    
    # The kernel config expects initramfs at "../../root" relative to kernel build directory
    echo "    Extracting initramfs for embedding..."
    mkdir -p ../../root
    cd ../../root
    rm -rf ./*
    xz -dc "$OUTPUT_DIR/initramfs.cpio.xz" | cpio -id
    cd "$KERNEL_DIR/linux-${KERNEL_VERSION}"
    
    # Restore initramfs path in config
    sed -i 's/CONFIG_INITRAMFS_SOURCE=""/CONFIG_INITRAMFS_SOURCE="..\/..\/root"/' .config
    
    # Rebuild with embedded initramfs
    make olddefconfig
    make -j$(nproc) bzImage
    
    # Copy final kernel
    echo "    Copying embedded kernel..."
    cp arch/x86/boot/bzImage "$OUTPUT_DIR/vmlinuz.efi"
    
    kernel_size=$(stat -c%s "$OUTPUT_DIR/vmlinuz.efi")
    echo "  Embedded kernel created: $(echo "$kernel_size / 1024 / 1024" | bc)MB"
    exit 0
    
else
    # Full build mode (original behavior)
    echo "  Full build: kernel + embedded initramfs..."
    
    # The kernel config expects initramfs at "../../root" relative to kernel build directory
    echo "  Setting up initramfs for kernel embedding..."
    mkdir -p ../../root
    cd ../../root
    
    # Extract our initramfs to the location expected by kernel config
    echo "  Extracting initramfs to expected location..."
    # Remove any existing contents first
    rm -rf ./*
    # Extract the initramfs cpio archive
    xz -dc "$OUTPUT_DIR/initramfs.cpio.xz" | cpio -id
    
    echo "  Initramfs extracted to: $(pwd)"
    echo "  Initramfs contents:"
    find . -maxdepth 2 -type f -o -type d | head -20
    
    # Return to kernel source directory
    cd "$KERNEL_DIR/linux-${KERNEL_VERSION}"
    
    # Verify kernel config is properly set up
    echo "  Verifying kernel configuration..."
    if grep -q 'CONFIG_INITRAMFS_SOURCE="../../root"' .config; then
        echo "    ✓ Initramfs source path configured correctly"
    else
        echo "    ✗ Warning: Initramfs source path not found in config"
    fi
    
    if grep -q 'CONFIG_INITRAMFS_COMPRESSION_XZ=y' .config; then
        echo "    ✓ XZ compression enabled"
    else
        echo "    ✗ Warning: XZ compression not enabled"
    fi
    
    # Update config to handle any new options
    echo "  Updating kernel configuration..."
    make olddefconfig
    
    # Build the kernel
    echo "  Building kernel with embedded initramfs..."
    echo "    This may take several minutes..."
    make -j$(nproc) bzImage
fi

# Verify kernel was built
if [ ! -f "arch/x86/boot/bzImage" ]; then
    echo "    Error: Kernel build failed - bzImage not found" >&2
    exit 1
fi

# Copy built kernel to output location
echo "  Installing built kernel..."
cp arch/x86/boot/bzImage "$OUTPUT_DIR/vmlinuz.efi"

# Get kernel size info
kernel_size=$(stat -c%s "$OUTPUT_DIR/vmlinuz.efi")
initramfs_size=$(stat -c%s "$OUTPUT_DIR/initramfs.cpio.xz")

echo "[DEBUG] Built kernel size: $kernel_size bytes"
echo "[DEBUG] Original initramfs size: $initramfs_size bytes"

# Create version info file
echo "  Creating version information..."
cat > "$OUTPUT_DIR/version.txt" << EOF
Zero-OS Custom Kernel (Linux $KERNEL_VERSION)
Built: $(date)
Source: $KERNEL_URL
Config: kernel-config-generic
Initramfs: Embedded during compilation
Final size: $kernel_size bytes
Original initramfs: $initramfs_size bytes

Boot Options:
  Method: Custom kernel with embedded initramfs
  Usage: Boot vmlinuz.efi directly
  GRUB: linux vmlinuz.efi
  No separate initramfs needed

Build Environment:
  Alpine Linux $(cat /etc/alpine-release 2>/dev/null || echo 'unknown')
  GCC: $(gcc --version | head -1)
  Make: $(make --version | head -1)
EOF

# Build kernel modules
echo "  Building kernel modules..."
make -j$(nproc) modules

# Install modules to output directory
echo "  Installing kernel modules..."
mkdir -p "$OUTPUT_DIR/modules"
make INSTALL_MOD_PATH="$OUTPUT_DIR/modules" modules_install

# Find the installed module directory
MODULE_DIR=$(find "$OUTPUT_DIR/modules/lib/modules" -maxdepth 1 -type d -name "*" | head -1)
if [ -n "$MODULE_DIR" ]; then
    echo "    Modules installed to: $MODULE_DIR"
    echo "Modules: $MODULE_DIR" >> "$OUTPUT_DIR/version.txt"
else
    echo "    Warning: No modules installed"
fi

# Clean up intermediate files
echo "  Cleaning up build artifacts..."
cd "$KERNEL_DIR"
rm -rf "linux-${KERNEL_VERSION}"
rm -f "linux-${KERNEL_VERSION}.tar.xz"

echo "[+] Custom kernel build completed successfully"
echo "    Kernel: $OUTPUT_DIR/vmlinuz.efi ($(du -h $OUTPUT_DIR/vmlinuz.efi | cut -f1))"
echo "    Version: Linux $KERNEL_VERSION Zero-OS"
if [ -d "$OUTPUT_DIR/modules" ]; then
    echo "    Modules: $OUTPUT_DIR/modules/ ($(du -sh $OUTPUT_DIR/modules | cut -f1))"
fi
echo "    Method: Custom compilation with embedded initramfs"

# Final verification
if [ ! -f "$OUTPUT_DIR/vmlinuz.efi" ] || [ "$(stat -c%s "$OUTPUT_DIR/vmlinuz.efi")" -eq 0 ]; then
    echo "    Error: Final kernel verification failed" >&2
    exit 1
fi

echo "    ✓ Kernel ready for deployment"