#!/bin/bash
# dracut-module.sh - Custom dracut module for Ubuntu uBlue immutable system
# This script creates a dracut module that mounts squashfs layers via overlayfs at boot time
# to achieve true immutability for the root filesystem.

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODULE_NAME="90-ublue-immutable"
DRACUT_MODULES_DIR="/usr/lib/dracut/modules.d"
MODULE_DIR="${DRACUT_MODULES_DIR}/${MODULE_NAME}"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Create the dracut module structure
create_module_structure() {
    log_info "Creating dracut module structure..."
    
    # Create module directory
    mkdir -p "${MODULE_DIR}"
    
    # Create module files
    cat > "${MODULE_DIR}/module-setup.sh" << 'EOF'
#!/bin/bash
# module-setup.sh - Dracut module setup script for uBlue immutable system

check() {
    # Check if composefs layers exist (more generic than uBlue-specific)
    if [[ -d /var/lib/composefs-alternative/layers ]]; then
        return 0
    fi
    
    # Check if we're on a uBlue system as fallback
    if [[ -f /etc/os-release ]] && grep -q "ublue" /etc/os-release; then
        return 0
    fi
    
    return 255
}

depends() {
    echo "base"
}

install() {
    # Install the module script
    inst_hook cmdline 95 "$moddir/parse-ublue-cmdline.sh"
    inst_hook mount 95 "$moddir/mount-ublue-layers.sh"
    inst_hook pre-pivot 95 "$moddir/setup-ublue-root.sh"
    
    # Install required binaries (removed mksquashfs/unsquashfs - build tools, not runtime)
    dracut_install mount umount losetup find sort jq
    
    # Install configuration files
    inst /etc/os-release
    
    # Install kernel modules for overlayfs and squashfs
    instmods overlay squashfs loop
    
    # Create secure state directory for inter-hook communication
    mkdir -p "$initdir/run/initramfs/ublue-state"
    chmod 700 "$initdir/run/initramfs/ublue-state"
}
EOF

    cat > "${MODULE_DIR}/parse-ublue-cmdline.sh" << 'EOF'
#!/bin/bash
# parse-ublue-cmdline.sh - Parse kernel command line for uBlue options

# Parse kernel command line for uBlue-specific options
for param in $(cat /proc/cmdline); do
    case $param in
        ublue.immutable=*)
            UBLUE_IMMUTABLE="${param#*=}"
            ;;
        ublue.layers=*)
            UBLUE_LAYERS="${param#*=}"
            ;;
        ublue.upper=*)
            UBLUE_UPPER="${param#*=}"
            ;;
        ublue.manifest=*)
            UBLUE_MANIFEST="${param#*=}"
            ;;
    esac
done

# Set defaults if not specified
: ${UBLUE_IMMUTABLE:=1}
: ${UBLUE_LAYERS:=/var/lib/composefs-alternative/layers}
: ${UBLUE_UPPER:=tmpfs}
: ${UBLUE_MANIFEST:=/var/lib/composefs-alternative/layers/manifest.json}

# Export variables for other scripts (using secure state directory)
STATE_DIR="/run/initramfs/ublue-state"
mkdir -p "$STATE_DIR"
chmod 700 "$STATE_DIR"

cat > "$STATE_DIR/cmdline.conf" << INNEREOF
UBLUE_IMMUTABLE="$UBLUE_IMMUTABLE"
UBLUE_LAYERS="$UBLUE_LAYERS"
UBLUE_UPPER="$UBLUE_UPPER"
UBLUE_MANIFEST="$UBLUE_MANIFEST"
INNEREOF
EOF

    cat > "${MODULE_DIR}/mount-ublue-layers.sh" << 'EOF'
#!/bin/bash
# mount-ublue-layers.sh - Mount squashfs layers for uBlue immutable system

# Source the parsed command line variables from secure state directory
STATE_DIR="/run/initramfs/ublue-state"
if [[ -f "$STATE_DIR/cmdline.conf" ]]; then
    source "$STATE_DIR/cmdline.conf"
else
    echo "uBlue command line configuration not found"
    exit 1
fi

# Skip if immutable mode is disabled
if [[ "${UBLUE_IMMUTABLE:-1}" != "1" ]]; then
    echo "uBlue immutable mode disabled, skipping layer mounting"
    exit 0
fi

echo "Setting up uBlue immutable layers..."

# Create mount points
mkdir -p /run/ublue/layers
mkdir -p /run/ublue/overlay

# Find and mount squashfs layers with deterministic ordering
LAYER_COUNT=0
LOWER_DIRS=""

# Function to read layer manifest for deterministic ordering
read_layer_manifest() {
    local manifest_file="$1"
    local layers_dir="$2"
    
    if [[ -f "$manifest_file" ]]; then
        echo "Reading layer manifest: $manifest_file"
        # Use jq if available, otherwise fallback to simple parsing
        if command -v jq >/dev/null 2>&1; then
            jq -r '.layers[] | .name + ":" + .file' "$manifest_file" 2>/dev/null | while IFS=: read -r name file; do
                if [[ -n "$name" && -n "$file" ]]; then
                    echo "$name:$file"
                fi
            done
        else
            # Improved fallback: extract both name and file fields
            # This handles the case where jq is not available but we still want manifest ordering
            local temp_file=$(mktemp)
            local name_file_pairs=()
            
            # Extract name and file pairs using grep and sed
            while IFS= read -r line; do
                if [[ "$line" =~ "name"[[:space:]]*:[[:space:]]*"([^"]*)" ]]; then
                    local name="${BASH_REMATCH[1]}"
                    # Look for the corresponding file field in the same object
                    local file=""
                    # Simple heuristic: assume file follows name in the same object
                    # This is not perfect but better than just extracting names
                    if [[ -n "$name" ]]; then
                        # Try to find the file by looking for the name.squashfs in the layers directory
                        if [[ -f "$layers_dir/$name.squashfs" ]]; then
                            file="$layers_dir/$name.squashfs"
                        else
                            # Fallback: just use the name as a hint for alphabetical ordering
                            file="$name"
                        fi
                        echo "$name:$file"
                    fi
                fi
            done < "$manifest_file" > "$temp_file"
            
            # Output the results and clean up
            if [[ -s "$temp_file" ]]; then
                cat "$temp_file"
            fi
            rm -f "$temp_file"
        fi
    fi
}

# Function to mount layers in specified order
mount_layers_ordered() {
    local layers_dir="$1"
    local manifest_file="$2"
    
    # Try to read manifest for ordering
    local ordered_layers=()
    if [[ -f "$manifest_file" ]]; then
        while IFS= read -r layer_info; do
            if [[ -n "$layer_info" ]]; then
                ordered_layers+=("$layer_info")
            fi
        done < <(read_layer_manifest "$manifest_file" "$layers_dir")
    fi
    
    # If no manifest or empty, fall back to alphabetical ordering
    if [[ ${#ordered_layers[@]} -eq 0 ]]; then
        echo "No manifest found, using alphabetical ordering"
        while IFS= read -r -d '' layer_file; do
            ordered_layers+=("$(basename "$layer_file" .squashfs):$layer_file")
        done < <(find "$layers_dir" -name "*.squashfs" -print0 | sort -z)
    fi
    
    # Mount layers in order
    for layer_info in "${ordered_layers[@]}"; do
        IFS=: read -r layer_name layer_file <<< "$layer_info"
        
        if [[ -f "$layer_file" ]]; then
            mount_point="/run/ublue/layers/${layer_name}"
            
            echo "Mounting layer: $layer_name"
            mkdir -p "$mount_point"
            
            # Mount the squashfs layer
            if mount -t squashfs -o ro "$layer_file" "$mount_point"; then
                LOWER_DIRS="${LOWER_DIRS}:${mount_point}"
                ((LAYER_COUNT++))
                echo "Successfully mounted $layer_name"
            else
                echo "Failed to mount $layer_name"
            fi
        fi
    done
    
    # Remove leading colon
    LOWER_DIRS="${LOWER_DIRS#:}"
}

# Mount layers from the specified directory
if [[ -d "${UBLUE_LAYERS}" ]]; then
    echo "Scanning for squashfs layers in ${UBLUE_LAYERS}..."
    mount_layers_ordered "${UBLUE_LAYERS}" "${UBLUE_MANIFEST}"
fi

# If no layers found, try alternative locations
if [[ -z "$LOWER_DIRS" ]]; then
    echo "No squashfs layers found in ${UBLUE_LAYERS}, checking for OSTree deployment..."
    
    # Check OSTree deployment as fallback
    if [[ -d /sysroot/ostree/deploy ]]; then
        deploy_path=$(find /sysroot/ostree/deploy -maxdepth 3 -name "deploy" -type d | head -1)
        if [[ -n "$deploy_path" ]]; then
            echo "Found OSTree deployment at $deploy_path"
            echo "WARNING: Using direct OSTree deployment instead of squashfs layers"
            echo "This may indicate composefs-alternative.sh did not create layers properly"
            
            # Use the deployment as a fallback layer
            # Note: This creates a mixed overlayfs with directory + potential squashfs layers
            # This is acceptable as a fallback but not ideal for production
            LOWER_DIRS="$deploy_path"
            LAYER_COUNT=1
            
            # Store fallback information for debugging
            echo "UBLUE_FALLBACK_MODE=ostree" >> "$STATE_DIR/layers.conf"
            echo "UBLUE_FALLBACK_PATH=$deploy_path" >> "$STATE_DIR/layers.conf"
        else
            echo "ERROR: No OSTree deployment found"
            echo "Cannot proceed with immutable root setup"
            exit 1
        fi
    else
        echo "ERROR: No squashfs layers found and no OSTree deployment available"
        echo "Cannot proceed with immutable root setup"
        echo "Please ensure composefs-alternative.sh has created layers in ${UBLUE_LAYERS}"
        exit 1
    fi
fi

# Create upper layer
case "${UBLUE_UPPER}" in
    tmpfs)
        echo "Creating tmpfs upper layer..."
        mount -t tmpfs -o size=1G tmpfs /run/ublue/overlay/upper
        ;;
    var)
        echo "Using /var as upper layer..."
        mkdir -p /run/ublue/overlay/upper
        mount --bind /sysroot/var /run/ublue/overlay/upper
        ;;
    *)
        echo "Using custom upper layer: ${UBLUE_UPPER}"
        mkdir -p /run/ublue/overlay/upper
        mount --bind "${UBLUE_UPPER}" /run/ublue/overlay/upper
        ;;
esac

# Create work directory
mkdir -p /run/ublue/overlay/work

# Store layer information in secure state directory
cat > "$STATE_DIR/layers.conf" << INNEREOF
UBLUE_LOWER_DIRS="$LOWER_DIRS"
UBLUE_LAYER_COUNT=$LAYER_COUNT
UBLUE_OVERLAY_READY=1
INNEREOF

echo "uBlue layer mounting complete: $LAYER_COUNT layers mounted"
EOF

    cat > "${MODULE_DIR}/setup-ublue-root.sh" << 'EOF'
#!/bin/bash
# setup-ublue-root.sh - Set up overlayfs root for uBlue immutable system

# Source layer information from secure state directory
STATE_DIR="/run/initramfs/ublue-state"
if [[ -f "$STATE_DIR/layers.conf" ]]; then
    source "$STATE_DIR/layers.conf"
else
    echo "uBlue layer configuration not found"
    exit 1
fi

# Skip if overlay is not ready
if [[ "${UBLUE_OVERLAY_READY:-0}" != "1" ]]; then
    echo "uBlue overlay not ready, skipping root setup"
    exit 0
fi

echo "Setting up uBlue overlayfs root..."

# Create the overlayfs mount
if [[ -n "$UBLUE_LOWER_DIRS" ]]; then
    echo "Creating overlayfs with layers: $UBLUE_LOWER_DIRS"
    
    # Mount overlayfs
    mount -t overlay overlay \
        -o lowerdir="$UBLUE_LOWER_DIRS",upperdir=/run/ublue/overlay/upper,workdir=/run/ublue/overlay/work \
        /run/ublue/overlay/root
    
    # Verify the mount
    if mountpoint -q /run/ublue/overlay/root; then
        echo "Overlayfs root created successfully"
        
        # Set up pivot root
        mkdir -p /run/ublue/overlay/root/oldroot
        
        # Move essential mounts to new root
        for mount in /proc /sys /dev /run; do
            if mountpoint -q "$mount"; then
                mount --move "$mount" "/run/ublue/overlay/root$mount"
            fi
        done
        
        # Pivot to new root
        cd /run/ublue/overlay/root
        pivot_root . oldroot
        
        # Clean up old root
        umount -l /oldroot
        
        echo "uBlue immutable root setup complete"
    else
        echo "Failed to create overlayfs root"
        exit 1
    fi
else
    echo "No layers available for overlayfs"
    exit 1
fi
EOF

    # Make scripts executable
    chmod +x "${MODULE_DIR}/module-setup.sh"
    chmod +x "${MODULE_DIR}/parse-ublue-cmdline.sh"
    chmod +x "${MODULE_DIR}/mount-ublue-layers.sh"
    chmod +x "${MODULE_DIR}/setup-ublue-root.sh"
    
    log_success "Dracut module structure created in ${MODULE_DIR}"
}

# Install the module
install_module() {
    log_info "Installing dracut module..."
    
    # Check if module already exists
    if [[ -d "${MODULE_DIR}" ]]; then
        log_warning "Module already exists, backing up..."
        mv "${MODULE_DIR}" "${MODULE_DIR}.backup.$(date +%Y%m%d_%H%M%S)"
    fi
    
    create_module_structure
    
    log_success "Dracut module installed successfully"
}

# Generate initramfs with the module
generate_initramfs() {
    local kernel_version="${1:-}"
    
    if [[ -z "$kernel_version" ]]; then
        # Get current kernel version
        kernel_version=$(uname -r)
    fi
    
    log_info "Generating initramfs for kernel $kernel_version with uBlue module..."
    
    # Generate initramfs with our module
    dracut --force --add "$MODULE_NAME" --kver "$kernel_version"
    
    log_success "Initramfs generated with uBlue immutable module"
}

# Update GRUB configuration with idempotent parameter addition
update_grub() {
    log_info "Updating GRUB configuration..."
    
    if [[ -f /etc/default/grub ]]; then
        # Backup original
        cp /etc/default/grub /etc/default/grub.backup.$(date +%Y%m%d_%H%M%S)
        
        # Define uBlue parameters (including manifest for deterministic ordering)
        local ublue_params="ublue.immutable=1 ublue.layers=/var/lib/composefs-alternative/layers ublue.upper=tmpfs ublue.manifest=/var/lib/composefs-alternative/layers/manifest.json"
        
        # Check if parameters already exist
        if grep -q "ublue.immutable=" /etc/default/grub; then
            log_warning "uBlue parameters already exist in GRUB configuration"
        else
            # Add uBlue parameters to GRUB_CMDLINE_LINUX_DEFAULT
            sed -i "s/GRUB_CMDLINE_LINUX_DEFAULT=\"\([^\"]*\)\"/GRUB_CMDLINE_LINUX_DEFAULT=\"\1 $ublue_params\"/" /etc/default/grub
            
            # Update GRUB
            update-grub
            
            log_success "GRUB configuration updated with uBlue parameters"
        fi
    else
        log_warning "GRUB configuration not found, manual update may be required"
    fi
}

# Test the module
test_module() {
    log_info "Testing dracut module..."
    
    # Check if module files exist
    local required_files=(
        "module-setup.sh"
        "parse-ublue-cmdline.sh"
        "mount-ublue-layers.sh"
        "setup-ublue-root.sh"
    )
    
    for file in "${required_files[@]}"; do
        if [[ -f "${MODULE_DIR}/${file}" ]]; then
            log_success "✓ $file exists"
        else
            log_error "✗ $file missing"
            return 1
        fi
    done
    
    # Test module setup
    if (cd "${MODULE_DIR}" && bash module-setup.sh check); then
        log_success "✓ Module check passed"
    else
        log_warning "⚠ Module check failed (this may be normal on non-uBlue systems)"
    fi
    
    log_success "Module test completed"
}

# Remove the module
remove_module() {
    log_info "Removing dracut module..."
    
    if [[ -d "${MODULE_DIR}" ]]; then
        rm -rf "${MODULE_DIR}"
        log_success "Module removed"
    else
        log_warning "Module not found"
    fi
}

# Show module status
show_status() {
    log_info "uBlue dracut module status:"
    
    if [[ -d "${MODULE_DIR}" ]]; then
        echo "✓ Module installed: ${MODULE_DIR}"
        echo "  - module-setup.sh"
        echo "  - parse-ublue-cmdline.sh"
        echo "  - mount-ublue-layers.sh"
        echo "  - setup-ublue-root.sh"
    else
        echo "✗ Module not installed"
    fi
    
    # Check if module is in any initramfs
    echo ""
    echo "Initramfs files containing uBlue module:"
    find /boot -name "initrd*" -o -name "initramfs*" 2>/dev/null | while read -r initramfs; do
        if lsinitrd "$initramfs" 2>/dev/null | grep -q "ublue"; then
            echo "✓ $initramfs"
        fi
    done
}

# Show usage
show_usage() {
    cat << EOF
uBlue Dracut Module Manager

Usage: $0 [COMMAND] [OPTIONS]

Commands:
    install              Install the dracut module
    generate [KERNEL]    Generate initramfs with module (optional kernel version)
    update-grub          Update GRUB with uBlue parameters (idempotent)
    test                 Test the module installation
    remove               Remove the module
    status               Show module status
    help                 Show this help

Examples:
    $0 install                    # Install the module
    $0 generate                   # Generate initramfs for current kernel
    $0 generate 5.15.0-56-generic # Generate for specific kernel
    $0 update-grub               # Update GRUB configuration
    $0 test                      # Test the module
    $0 status                    # Show status

Kernel Parameters:
    ublue.immutable=1                                    # Enable immutable mode
    ublue.layers=/var/lib/composefs-alternative/layers   # Layer directory
    ublue.upper=tmpfs                                    # Upper layer type (tmpfs/var/custom)
    ublue.manifest=/path/to/manifest.json               # Layer ordering manifest

Security Improvements:
    - Secure state directory (/run/initramfs/ublue-state)
    - Removed initramfs bloat (no mksquashfs/unsquashfs)
    - Deterministic layer ordering via manifest
    - Idempotent GRUB parameter updates

EOF
}

# Main function
main() {
    local command="${1:-help}"
    
    case "$command" in
        install)
            check_root
            install_module
            ;;
        generate)
            check_root
            generate_initramfs "$2"
            ;;
        update-grub)
            check_root
            update_grub
            ;;
        test)
            test_module
            ;;
        remove)
            check_root
            remove_module
            ;;
        status)
            show_status
            ;;
        help|--help|-h)
            show_usage
            ;;
        *)
            log_error "Unknown command: $command"
            show_usage
            exit 1
            ;;
    esac
}

# Run main function with all arguments
main "$@"