#!/bin/bash
set -e

# Detect if running in container or locally
if [ -d "/build" ]; then
    # Running in Docker container
    ALPINE_ROOT="/build/initramfs"
    OUTPUT_DIR="/build/output"
    WORKSPACE_ROOT="/build"
else
    # Running locally
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    WORKSPACE_ROOT="$(dirname "$SCRIPT_DIR")"
    ALPINE_ROOT="$WORKSPACE_ROOT/initramfs"
    OUTPUT_DIR="$WORKSPACE_ROOT/output"
fi

echo "[+] Compiling Zero-OS components individually..."

# Check if components exist
if [ ! -d "$WORKSPACE_ROOT/components" ]; then
    echo "Error: Components directory not found."
    exit 1
fi

echo "  Components found:"
ls -1 "$WORKSPACE_ROOT/components/" | sed 's/^/    - /'
echo ""

# Function to build a single component
build_component() {
    local component="$1"
    local component_dir="$WORKSPACE_ROOT/components/$component"
    
    echo "  Building $component..."
    
    if [ ! -d "$component_dir" ]; then
        echo "    ✗ Component directory not found: $component_dir"
        return 1
    fi
    
    if [ ! -f "$component_dir/Cargo.toml" ]; then
        echo "    ✗ Cargo.toml not found in: $component_dir"
        return 1
    fi
    
    cd "$component_dir"
    
    # Source cargo environment
    . ~/.cargo/env
    
    # Handle special case for myceliumd (requires subdirectory build)
    if [ "$component" = "mycelium" ]; then
        echo "    Building myceliumd from mycelium/myceliumd subdirectory..."
        cd "$component_dir/myceliumd"
        if [ ! -f "Cargo.toml" ]; then
            echo "    ✗ myceliumd/Cargo.toml not found"
            return 1
        fi
        
        # Set OpenSSL environment variables for musl cross-compilation
        export OPENSSL_STATIC=1
        export OPENSSL_DIR=/usr
        
        echo "    Running: cargo build --release --target x86_64-unknown-linux-musl"
        CARGO_TARGET_DIR="$WORKSPACE_ROOT/target" cargo build --release --target x86_64-unknown-linux-musl || {
            echo "    ⚠ myceliumd build failed (OpenSSL cross-compilation issue)"
            echo "    This is a known issue with OpenSSL and musl cross-compilation"
            return 1
        }
    elif [ "$component" = "rfs" ]; then
        # Build RFS with required features
        echo "    Removing "
        echo "    Running: cargo build --release --target x86_64-unknown-linux-musl --features build-binary"
        CARGO_TARGET_DIR="$WORKSPACE_ROOT/target" cargo build --release --target x86_64-unknown-linux-musl --features build-binary
    else
        # Build component with musl target for static linking
        echo "    Running: cargo build --release --target x86_64-unknown-linux-musl"
        CARGO_TARGET_DIR="$WORKSPACE_ROOT/target" cargo build --release --target x86_64-unknown-linux-musl
    fi
    
    echo "    ✓ $component build completed"
    return 0
}

# Function to optimize binaries (strip + upx)
optimize_binary() {
    local binary_path="$1"
    local binary_name=$(basename "$binary_path")
    
    if [ ! -f "$binary_path" ]; then
        echo "    ⚠ Binary not found for optimization: $binary_path"
        return 1
    fi
    
    local original_size=$(stat -c%s "$binary_path")
    
    # Strip debug symbols and unnecessary sections
    echo "    Stripping $binary_name..."
    strip --strip-all "$binary_path" 2>/dev/null || {
        echo "    ⚠ Failed to strip $binary_name"
    }
    
    local stripped_size=$(stat -c%s "$binary_path")
    
    # Compress with UPX (with fallback if UPX fails)
    echo "    Compressing $binary_name with UPX..."
    if upx --best --lzma "$binary_path" 2>/dev/null; then
        local final_size=$(stat -c%s "$binary_path")
        echo "    ✓ $binary_name optimized: $original_size → $stripped_size → $final_size bytes"
    else
        echo "    ⚠ UPX compression failed for $binary_name, keeping stripped version"
        echo "    ✓ $binary_name stripped: $original_size → $stripped_size bytes"
    fi
}

# Function to install workspace binaries
install_component_binaries() {
    local target_dir="$WORKSPACE_ROOT/target/x86_64-unknown-linux-musl/release"
    
    echo "  Installing component binaries..."
    echo "    Target directory: $target_dir"
    
    # Debug: Show what was built
    echo "    Built binaries:"
    ls -la "$target_dir" 2>/dev/null | grep -E '^-.*x.*' | awk '{print "      " $9 " (" $5 " bytes)"}' || echo "      No binaries found"
    
    # Create output and alpine root directories
    mkdir -p "$OUTPUT_DIR"
    mkdir -p "$ALPINE_ROOT/sbin"
    mkdir -p "$ALPINE_ROOT/usr/bin"
    
    # Install specific binaries
    local installed=0
    
    # Install zinit
    if [ -f "$target_dir/zinit" ]; then
        cp "$target_dir/zinit" "$ALPINE_ROOT/sbin/"
        cp "$target_dir/zinit" "$OUTPUT_DIR/"
        chmod +x "$ALPINE_ROOT/sbin/zinit"
        
        # Optimize binaries
        optimize_binary "$ALPINE_ROOT/sbin/zinit"
        optimize_binary "$OUTPUT_DIR/zinit"
        
        echo "    ✓ zinit -> /sbin/zinit (optimized)"
        installed=$((installed + 1))
    else
        echo "    ✗ zinit binary not found"
    fi
    
    # Install rfs (check for different possible names)
    local rfs_binary=""
    for name in "rfs" "rfs-server" "rfs-bin" "fl-server"; do
        if [ -f "$target_dir/$name" ]; then
            rfs_binary="$name"
            break
        fi
    done
    
    if [ -n "$rfs_binary" ]; then
        cp "$target_dir/$rfs_binary" "$ALPINE_ROOT/usr/bin/rfs"
        cp "$target_dir/$rfs_binary" "$OUTPUT_DIR/rfs"
        chmod +x "$ALPINE_ROOT/usr/bin/rfs"
        
        # Optimize binaries
        optimize_binary "$ALPINE_ROOT/usr/bin/rfs"
        optimize_binary "$OUTPUT_DIR/rfs"
        
        echo "    ✓ $rfs_binary -> /usr/bin/rfs (optimized)"
        installed=$((installed + 1))
    else
        echo "    ✗ rfs binary not found (checked: rfs, rfs-server, rfs-bin)"
    fi
    
    # Install mycelium (check for different possible names)
    local mycelium_binary=""
    for name in "mycelium" "myceliumd" "mycelium-cli"; do
        if [ -f "$target_dir/$name" ]; then
            mycelium_binary="$name"
            break
        fi
    done
    
    if [ -n "$mycelium_binary" ]; then
        cp "$target_dir/$mycelium_binary" "$ALPINE_ROOT/usr/bin/mycelium"
        cp "$target_dir/$mycelium_binary" "$OUTPUT_DIR/mycelium"
        chmod +x "$ALPINE_ROOT/usr/bin/mycelium"
        
        # Optimize binaries
        optimize_binary "$ALPINE_ROOT/usr/bin/mycelium"
        optimize_binary "$OUTPUT_DIR/mycelium"
        
        echo "    ✓ $mycelium_binary -> /usr/bin/mycelium (optimized)"
        installed=$((installed + 1))
    else
        echo "    ✗ mycelium binary not found (checked: mycelium, myceliumd, mycelium-cli)"
    fi
    
    echo "    Installed $installed out of 3 expected binaries"
    return 0
}

# Ensure Alpine root directory exists
if [ ! -d "$ALPINE_ROOT" ]; then
    echo "  Creating Alpine initramfs root at $ALPINE_ROOT"
    mkdir -p "$ALPINE_ROOT"
fi

# Build each component individually
failed_components=()

for component in zinit mycelium rfs; do
    if build_component "$component"; then
        echo "  ✓ $component build completed"
    else
        echo "  ✗ $component build failed"
        failed_components+=("$component")
    fi
done

# Report build results
if [ ${#failed_components[@]} -eq 0 ]; then
    echo "  ✓ All components built successfully"
else
    echo "  ✗ Failed to build: ${failed_components[*]}"
    echo "  Note: Continuing with available binaries..."
fi

# Install binaries
install_component_binaries

echo ""
echo "[+] Zero-OS component compilation completed"

# Show installed binaries info
echo ""
echo "[+] Compiled Zero-OS components in Alpine root:"
for binary in "/sbin/zinit" "/usr/bin/rfs" "/usr/bin/mycelium"; do
    if [ -x "$ALPINE_ROOT$binary" ]; then
        size=$(stat -c%s "$ALPINE_ROOT$binary" 2>/dev/null || echo "unknown")
        echo "  $binary (${size} bytes)"
    else
        echo "  $binary (NOT BUILT)"
    fi
done

echo ""
echo "[+] Compiled binaries also saved to: $OUTPUT_DIR"
echo "[+] Zero-OS Rust component compilation complete!"