#!/bin/bash
# Main orchestrator script for Zero OS Alpine Initramfs Builder

set -euo pipefail

# Script directory and project root detection
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Source all libraries
source "${SCRIPT_DIR}/lib/common.sh"
source "${SCRIPT_DIR}/lib/stages.sh"
source "${SCRIPT_DIR}/lib/docker.sh"
source "${SCRIPT_DIR}/lib/alpine.sh"
source "${SCRIPT_DIR}/lib/components.sh"
source "${SCRIPT_DIR}/lib/initramfs.sh"
source "${SCRIPT_DIR}/lib/kernel.sh"

# Build configuration loaded from config/build.conf via common.sh
# Environment variables can override config file values
ALPINE_VERSION="${ALPINE_VERSION:-3.22}"
KERNEL_VERSION="${KERNEL_VERSION:-6.12.44}"
RUST_TARGET="${RUST_TARGET:-x86_64-unknown-linux-musl}"
OPTIMIZATION_LEVEL="${OPTIMIZATION_LEVEL:-max}"

# Directory configuration
export INSTALL_DIR="${PROJECT_ROOT}/initramfs"
export COMPONENTS_DIR="${PROJECT_ROOT}/components"
export KERNEL_DIR="${PROJECT_ROOT}/kernel"
export DIST_DIR="${PROJECT_ROOT}/dist"

# Configuration files
CONFIG_DIR="${PROJECT_ROOT}/config"
PACKAGES_LIST="${CONFIG_DIR}/packages.list"
SOURCES_CONF="${CONFIG_DIR}/sources.conf"
MODULES_CONF="${CONFIG_DIR}/modules.conf"
KERNEL_CONFIG="${CONFIG_DIR}/kernel.config"
FIRMWARE_CONF="${CONFIG_DIR}/firmware.conf"
ZINIT_CONFIG_DIR="${CONFIG_DIR}/zinit"

# Build options
USE_CONTAINER="${USE_CONTAINER:-auto}"
CLEAN_BUILD="${CLEAN_BUILD:-false}"
KEEP_ARTIFACTS="${KEEP_ARTIFACTS:-false}"

# Display usage information
function show_usage() {
    cat << EOF
Zero OS Alpine Initramfs Builder (Incremental)

Usage: $0 [OPTIONS]

Options:
  --clean              Clean build (remove all artifacts first)
  --keep-artifacts     Keep build artifacts after completion
  --force-rebuild      Force rebuild all stages (ignore completion markers)
  --rebuild-from=STAGE Force rebuild from specific stage onward
  --show-stages        Show stage completion status
  --help               Show this help message

Environment Variables:
  ALPINE_VERSION       Alpine Linux version (default: 3.22)
  KERNEL_VERSION       Linux kernel version (default: 6.12.44)
  RUST_TARGET          Rust compilation target (default: x86_64-unknown-linux-musl)
  OPTIMIZATION_LEVEL   Optimization level: max|size|speed (default: max)
  DEBUG                Enable debug output (default: 0)
  FORCE_REBUILD        Force rebuild all stages (default: false)

Examples:
  $0                          # Incremental build (skip completed stages)
  $0 --clean                  # Clean build (remove artifacts + stage markers)
  $0 --force-rebuild          # Force rebuild all stages
  $0 --rebuild-from=zinit_setup  # Rebuild from zinit setup onward
  $0 --show-stages           # Show which stages are completed
  DEBUG=1 $0                 # Build with debug output

Development Workflow:
  ./scripts/dev-container.sh start   # Start persistent container
  ./scripts/dev-container.sh shell   # Enter container for debugging
  DEBUG=1 ./scripts/build.sh         # Run incremental build inside container
EOF
}

# Parse command line arguments
function parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --clean)
                CLEAN_BUILD="true"
                shift
                ;;
            --keep-artifacts)
                KEEP_ARTIFACTS="true"
                shift
                ;;
            --force-rebuild)
                export FORCE_REBUILD="true"
                shift
                ;;
            --rebuild-from=*)
                export REBUILD_FROM_STAGE="${1#*=}"
                export FORCE_REBUILD="true"
                shift
                ;;
            --show-stages)
                # Initialize stage tracking and show status
                source "${SCRIPT_DIR}/lib/stages.sh"
                stages_init
                stages_status
                exit 0
                ;;
            --help|-h)
                show_usage
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_usage
                exit 1
                ;;
        esac
    done
}

# Setup build environment
function setup_build_environment() {
    section_header "Setting up build environment"
    
    log_info "Project root: ${PROJECT_ROOT}"
    log_info "Alpine version: ${ALPINE_VERSION}"
    log_info "Kernel version: ${KERNEL_VERSION}"
    log_info "Rust target: ${RUST_TARGET}"
    log_info "Optimization level: ${OPTIMIZATION_LEVEL}"

    # Ensure a stable CWD inside the container (prefer /workspace)
    if in_container; then
        if [[ "$(pwd)" != "${PROJECT_ROOT}" ]]; then
            log_info "Ensuring container CWD=${PROJECT_ROOT}"
            safe_execute cd "${PROJECT_ROOT}"
        fi
    fi
    
    # Create build directories only if we're in container
    # Host will let container create them to avoid permission issues
    if in_container; then
        safe_mkdir "$INSTALL_DIR"
        safe_mkdir "$COMPONENTS_DIR"
        safe_mkdir "$KERNEL_DIR"
        safe_mkdir "$DIST_DIR"
    else
        log_info "Skipping directory creation on host (container will create them)"
    fi
    
    # Check dependencies
    if ! check_dependencies; then
        log_error "Dependency check failed"
        return 1
    fi
    
    # Verify configuration files exist
    verify_configuration_files
    
    log_info "Build environment setup complete"
}

# Verify all required configuration files exist
function verify_configuration_files() {
    section_header "Verifying Configuration Files"
    
    local required_configs=(
        "$PACKAGES_LIST"
        "$SOURCES_CONF"
        "$MODULES_CONF"
        "$KERNEL_CONFIG"
        "$FIRMWARE_CONF"
    )
    
    local missing_configs=()
    
    for config in "${required_configs[@]}"; do
        if [[ ! -f "$config" ]]; then
            missing_configs+=("$config")
        else
            log_info "✓ Configuration found: $(basename "$config")"
        fi
    done
    
    if [[ ${#missing_configs[@]} -gt 0 ]]; then
        log_error "Missing configuration files:"
        for config in "${missing_configs[@]}"; do
            log_error "  - $config"
        done
        log_error "Run the setup script or create configuration files manually"
        return 1
    fi
    
    # Check zinit configuration directory
    if [[ ! -d "$ZINIT_CONFIG_DIR" ]]; then
        log_error "zinit configuration directory not found: ${ZINIT_CONFIG_DIR}"
        return 1
    fi
    
    log_info "All configuration files verified"
}

# Incremental build process with stage tracking
function main_build_process() {
    section_header "Starting Zero OS Alpine Initramfs Build (Incremental)"
    
    local start_time=$(date +%s)
    
    # Initialize stage tracking
    stages_init
    
    # Define stage wrapper functions
    function stage_alpine_extract() {
        alpine_extract_miniroot "$INSTALL_DIR" "$ALPINE_VERSION"
    }
    
    function stage_alpine_configure() {
        alpine_configure_repos "$INSTALL_DIR" "$ALPINE_VERSION"
        alpine_configure_system "$INSTALL_DIR"
    }
    
    function stage_alpine_packages() {
        alpine_install_packages "$INSTALL_DIR" "$PACKAGES_LIST"
    }
    
    function stage_alpine_firmware() {
        alpine_install_firmware "$INSTALL_DIR" "$FIRMWARE_CONF"
    }
    
    function stage_components_build() {
        components_parse_sources_conf "$SOURCES_CONF" "$COMPONENTS_DIR" "$INSTALL_DIR"
    }
    
    function stage_components_verify() {
        components_verify_installation
    }
    
    function stage_kernel_modules() {
        # Create placeholder for kernel build
        local initramfs_archive="${DIST_DIR}/initramfs.cpio.xz"
        safe_mkdir "$DIST_DIR"
        safe_execute touch "$initramfs_archive"
        
        # Download and build kernel modules
        kernel_download_source "$KERNEL_DIR" "$KERNEL_VERSION"
        kernel_apply_config "$KERNEL_DIR" "$initramfs_archive" "$KERNEL_CONFIG"
        kernel_build_modules "$KERNEL_DIR" "$INSTALL_DIR" "$KERNEL_VERSION"
        
        # Set full kernel version for later stages
        FULL_KERNEL_VERSION=$(kernel_get_full_version "$KERNEL_VERSION" "$KERNEL_CONFIG")
        export FULL_KERNEL_VERSION
        log_info "Full kernel version: ${FULL_KERNEL_VERSION}"
    }
    
    function stage_zinit_setup() {
        initramfs_setup_zinit "$INSTALL_DIR" "$ZINIT_CONFIG_DIR"
    }
    
    function stage_init_script() {
        initramfs_install_init_script "$INSTALL_DIR" "${CONFIG_DIR}/init"
    }
    
    function stage_components_copy() {
        initramfs_copy_components "$INSTALL_DIR" "$COMPONENTS_DIR"
    }
    
    function stage_modules_setup() {
        # Calculate full kernel version (needed for incremental builds when kernel_modules stage is skipped)
        local full_kernel_version=$(kernel_get_full_version "$KERNEL_VERSION" "$KERNEL_CONFIG")
        export FULL_KERNEL_VERSION="$full_kernel_version"
        log_info "Using full kernel version: ${full_kernel_version}"
        
        initramfs_setup_modules "$INSTALL_DIR" "$MODULES_CONF" "$full_kernel_version"
    }
    
    function stage_modules_copy() {
        # Ensure FULL_KERNEL_VERSION is set (for incremental builds)
        if [[ -z "${FULL_KERNEL_VERSION:-}" ]]; then
            FULL_KERNEL_VERSION=$(kernel_get_full_version "$KERNEL_VERSION" "$KERNEL_CONFIG")
            export FULL_KERNEL_VERSION
        fi
        
        log_info "About to copy modules - FULL_KERNEL_VERSION: ${FULL_KERNEL_VERSION}"
        log_info "CONTAINER_MODULES_PATH: ${CONTAINER_MODULES_PATH:-not-set}"
        initramfs_copy_resolved_modules "$INSTALL_DIR" "$FULL_KERNEL_VERSION"
    }
    
    # Create RFS flists and embed them into initramfs prior to CPIO
    function stage_rfs_flists() {
        section_header "Creating RFS flists and embedding into initramfs"

        # Ensure FULL_KERNEL_VERSION is available
        if [[ -z "${FULL_KERNEL_VERSION:-}" ]]; then
            FULL_KERNEL_VERSION=$(kernel_get_full_version "$KERNEL_VERSION" "$KERNEL_CONFIG")
            export FULL_KERNEL_VERSION
            log_info "Resolved FULL_KERNEL_VERSION: ${FULL_KERNEL_VERSION}"
        fi

        # Normalize working directory to the project root to avoid relative path issues
        local _oldpwd
        _oldpwd="$(pwd)"
        safe_execute cd "${PROJECT_ROOT}"
        log_debug "stage_rfs_flists CWD (normalized): $(pwd)"

        # Ensure rfs scripts are executable when present (be robust if directory is missing)
        if [[ -d "${PROJECT_ROOT}/scripts/rfs" ]]; then
            safe_execute find "${PROJECT_ROOT}/scripts/rfs" -type f -name "*.sh" -exec chmod +x {} \;
        else
            log_warn "scripts/rfs directory not found under PROJECT_ROOT=${PROJECT_ROOT}; invoking packers via bash with absolute paths"
        fi

        # Build modules flist (writes to dist/flists/modules-${FULL_KERNEL_VERSION}.fl)
        safe_execute bash "${PROJECT_ROOT}/scripts/rfs/pack-modules.sh"

        # Build firmware flist with a reproducible tag:
        # Priority: env FIRMWARE_TAG > config/build.conf: FIRMWARE_TAG > "latest"
        local fw_tag
        if [[ -n "${FIRMWARE_TAG:-}" ]]; then
            fw_tag="${FIRMWARE_TAG}"
        else
            if [[ -f "${CONFIG_DIR}/build.conf" ]]; then
                # shellcheck source=/dev/null
                source "${CONFIG_DIR}/build.conf"
            fi
            fw_tag="${FIRMWARE_TAG:-latest}"
        fi
        log_info "Using firmware tag: ${fw_tag}"
        safe_execute env FIRMWARE_TAG="${fw_tag}" bash "${PROJECT_ROOT}/scripts/rfs/pack-firmware.sh"

        # Embed flists inside initramfs at /etc/rfs for zinit init scripts
        local etc_rfs_dir="${INSTALL_DIR}/etc/rfs"
        safe_mkdir "${etc_rfs_dir}"

        local modules_fl="${PROJECT_ROOT}/dist/flists/modules-${FULL_KERNEL_VERSION}.fl"
        if [[ -f "${modules_fl}" ]]; then
            safe_execute cp "${modules_fl}" "${etc_rfs_dir}/"
            log_info "Embedded modules flist: ${modules_fl} -> ${etc_rfs_dir}/"
        else
            log_warn "Modules flist not found: ${modules_fl}"
        fi

        local firmware_fl="${PROJECT_ROOT}/dist/flists/firmware-${fw_tag}.fl"
        if [[ -f "${firmware_fl}" ]]; then
            # Provide canonical name firmware-latest.fl expected by firmware.sh
            safe_execute cp "${firmware_fl}" "${etc_rfs_dir}/firmware-latest.fl"
            log_info "Embedded firmware flist: ${firmware_fl} -> ${etc_rfs_dir}/firmware-latest.fl"
        else
            log_warn "Firmware flist not found: ${firmware_fl}"
        fi

        log_info "RFS flists embedded into initramfs"
        # Restore previous working directory
        safe_execute cd "${_oldpwd}"
    }

    function stage_cleanup() {
        alpine_aggressive_cleanup "$INSTALL_DIR"
    }
    
    function stage_validation() {
        initramfs_validate "$INSTALL_DIR"
    }
    
    function stage_initramfs_create() {
        local initramfs_archive="${DIST_DIR}/initramfs.cpio.xz"
        # Normalize to absolute path to avoid CWD-related issues in later stages
        if [[ "${initramfs_archive}" != /* ]]; then
            initramfs_archive="${PROJECT_ROOT}/${initramfs_archive#./}"
        fi
        initramfs_create_cpio "$INSTALL_DIR" "$initramfs_archive"
        export INITRAMFS_ARCHIVE="$initramfs_archive"
        log_debug "stage_initramfs_create: INITRAMFS_ARCHIVE=${INITRAMFS_ARCHIVE}"
    }
    
    function stage_initramfs_test() {
        # Ensure INITRAMFS_ARCHIVE is set when skipping directly to this stage
        if [[ -z "${INITRAMFS_ARCHIVE:-}" ]]; then
            local archive_path="${DIST_DIR}/initramfs.cpio.xz"
            if [[ "${archive_path}" != /* ]]; then
                archive_path="${PROJECT_ROOT}/${archive_path#./}"
            fi
            export INITRAMFS_ARCHIVE="${archive_path}"
            log_debug "stage_initramfs_test: defaulting INITRAMFS_ARCHIVE=${INITRAMFS_ARCHIVE}"
        fi
        initramfs_test_archive "$INITRAMFS_ARCHIVE"
    }
    
    function stage_kernel_build() {
        local kernel_output="${DIST_DIR}/vmlinuz.efi"

        # Ensure INITRAMFS_ARCHIVE is set even if initramfs_create/test were skipped previously
        if [[ -z "${INITRAMFS_ARCHIVE:-}" ]]; then
            local archive_path="${DIST_DIR}/initramfs.cpio.xz"
            if [[ "${archive_path}" != /* ]]; then
                archive_path="${PROJECT_ROOT}/${archive_path#./}"
            fi
            export INITRAMFS_ARCHIVE="${archive_path}"
            log_debug "stage_kernel_build: defaulting INITRAMFS_ARCHIVE=${INITRAMFS_ARCHIVE}"
        fi

        # Ensure FULL_KERNEL_VERSION is set for versioned output filename
        if [[ -z "${FULL_KERNEL_VERSION:-}" ]]; then
            FULL_KERNEL_VERSION=$(kernel_get_full_version "$KERNEL_VERSION" "$KERNEL_CONFIG")
            export FULL_KERNEL_VERSION
            log_debug "stage_kernel_build: resolved FULL_KERNEL_VERSION=${FULL_KERNEL_VERSION}"
        fi

        kernel_build_with_initramfs "$KERNEL_CONFIG" "$INITRAMFS_ARCHIVE" "$kernel_output"
        export KERNEL_OUTPUT="$kernel_output"
    }

    # Boot tests removed - use runit.sh for testing instead

    # Run all stages with incremental tracking
    stage_run "alpine_extract" stage_alpine_extract
    stage_run "alpine_configure" stage_alpine_configure
    stage_run "alpine_packages" stage_alpine_packages
    stage_run "alpine_firmware" stage_alpine_firmware
    stage_run "components_build" stage_components_build
    stage_run "components_verify" stage_components_verify
    stage_run "kernel_modules" stage_kernel_modules
    stage_run "init_script" stage_init_script
    stage_run "components_copy" stage_components_copy
    stage_run "zinit_setup" stage_zinit_setup
    stage_run "modules_setup" stage_modules_setup
    stage_run "modules_copy" stage_modules_copy
    stage_run "cleanup" stage_cleanup
    stage_run "rfs_flists" stage_rfs_flists
    stage_run "validation" stage_validation
    stage_run "initramfs_create" stage_initramfs_create
    stage_run "initramfs_test" stage_initramfs_test
    stage_run "kernel_build" stage_kernel_build

    # Calculate build time
    local end_time=$(date +%s)
    local build_time=$((end_time - start_time))
    local build_minutes=$((build_time / 60))
    local build_seconds=$((build_time % 60))
    
    section_header "Build Complete"
    log_info "Build time: ${build_minutes}m ${build_seconds}s"
    log_info "Output files:"
    if [[ -n "${KERNEL_OUTPUT:-}" ]]; then
        log_info "  Kernel: ${KERNEL_OUTPUT} ($(get_file_size "$KERNEL_OUTPUT"))"
    else
        log_info "  Kernel: ${DIST_DIR}/vmlinuz.efi ($(get_file_size "${DIST_DIR}/vmlinuz.efi" 2>/dev/null || echo "not found"))"
    fi
    if [[ -n "${INITRAMFS_ARCHIVE:-}" ]]; then
        log_info "  Initramfs: ${INITRAMFS_ARCHIVE} ($(get_file_size "$INITRAMFS_ARCHIVE"))"
    else
        log_info "  Initramfs: ${DIST_DIR}/initramfs.cpio.xz ($(get_file_size "${DIST_DIR}/initramfs.cpio.xz" 2>/dev/null || echo "not found"))"
    fi
}


# Main function
function main() {
    # Parse command line arguments
    parse_arguments "$@"
    
    # Show banner
    echo ""
    echo "=================================================="
    echo "==     ZERO-OS ALPINE INITRAMFS BUILDER        =="
    echo "==           ThreeFold Edition                  =="
    echo "=================================================="
    echo ""
    
    # Clean build if requested
    if [[ "$CLEAN_BUILD" == "true" ]]; then
        section_header "Clean Build Requested"
        "$SCRIPT_DIR/clean.sh"
    fi
    
    # Setup environment
    setup_build_environment
    
    # Always use container builds for consistency
    if in_container; then
        log_info "Already in container, proceeding with build"
        # Enable debug mode in container for better output visibility
        if [[ "${DEBUG:-0}" != "1" ]]; then
            log_info "Enabling debug mode for container build visibility"
            export DEBUG=1
        fi
        main_build_process
    elif command_exists "podman" || command_exists "docker"; then
        log_info "Starting container build"
        docker_detect_runtime
        docker_build_container

        # Pass through relevant arguments to container
        local container_args=""
        if [[ "$KEEP_ARTIFACTS" == "true" ]]; then
            container_args="$container_args --keep-artifacts"
        fi

        docker_run_build "./scripts/build.sh${container_args}"
    else
        log_error "Container runtime required (podman or docker)"
        log_error "Install with: apt-get install podman"
        return 1
    fi
    
    
    section_header "Zero OS Build Complete"
    log_info "Ready to deploy Zero OS with Alpine Linux and zinit"
}

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