#!/bin/bash
# Build stage tracking and incremental build support

# Source common functions
LIB_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${LIB_SCRIPT_DIR}/common.sh"

# Stage tracking configuration
STAGES_DIR="${PROJECT_ROOT:-/workspace}/.build-stages"

# Initialize stage tracking
function stages_init() {
    section_header "Initializing Stage Tracking"
    
    safe_mkdir "$STAGES_DIR"
    log_info "Stage tracking directory: ${STAGES_DIR}"
    
    # Show existing completed stages
    local completed_stages=()
    if ls "$STAGES_DIR"/*.done >/dev/null 2>&1; then
        while IFS= read -r stage_file; do
            local stage_name=$(basename "$stage_file" .done)
            completed_stages+=("$stage_name")
        done < <(ls "$STAGES_DIR"/*.done 2>/dev/null)
        
        log_info "Previously completed stages: ${completed_stages[*]}"
    else
        log_info "No previously completed stages found"
    fi
}

# Check if stage is already completed
function stage_is_completed() {
    local stage_name="$1"
    local stage_file="${STAGES_DIR}/${stage_name}.done"
    
    if [[ -f "$stage_file" ]]; then
        local completion_time=$(stat -c %Y "$stage_file" 2>/dev/null || echo "unknown")
        log_debug "Stage '$stage_name' already completed at $(date -d @$completion_time 2>/dev/null || echo "unknown time")"
        return 0
    else
        log_debug "Stage '$stage_name' not completed yet"
        return 1
    fi
}

# Mark stage as completed
function stage_mark_completed() {
    local stage_name="$1"
    local stage_details="${2:-}"
    local stage_file="${STAGES_DIR}/${stage_name}.done"
    
    # Create completion marker with metadata
    cat > "$stage_file" << EOF
# Stage completion marker
STAGE_NAME="$stage_name"
COMPLETED_AT="$(date -Iseconds)"
COMPLETED_BY="$(whoami)"
DETAILS="$stage_details"
EOF
    
    log_info "Stage completed: $stage_name"
    if [[ -n "$stage_details" ]]; then
        log_debug "  Details: $stage_details"
    fi
}

# Remove stage completion marker (force rebuild)
function stage_force_rebuild() {
    local stage_name="$1"
    local stage_file="${STAGES_DIR}/${stage_name}.done"
    
    if [[ -f "$stage_file" ]]; then
        safe_execute rm "$stage_file"
        log_info "Stage marked for rebuild: $stage_name"
    else
        log_debug "Stage not completed, no need to force rebuild: $stage_name"
    fi
}

# Clear all stage markers (full rebuild)
function stages_clear_all() {
    section_header "Clearing All Stage Markers"
    
    if [[ -d "$STAGES_DIR" ]]; then
        local marker_count=$(ls "$STAGES_DIR"/*.done 2>/dev/null | wc -l || echo "0")
        if [[ $marker_count -gt 0 ]]; then
            safe_execute rm -f "$STAGES_DIR"/*.done
            log_info "Cleared ${marker_count} stage completion markers"
        else
            log_info "No stage markers to clear"
        fi
    fi
    
    log_info "Next build will rebuild all stages"
}

# Run stage with completion tracking
function stage_run() {
    local stage_name="$1"
    local stage_function="$2"
    shift 2
    local stage_args=("$@")
    
    log_info "=== STAGE: $stage_name ==="
    
    # Check if stage is already completed (unless forced)
    if [[ "${FORCE_REBUILD:-false}" != "true" ]] && stage_is_completed "$stage_name"; then
        log_info "Skipping completed stage: $stage_name"
        log_info "   (Use FORCE_REBUILD=true or remove ${STAGES_DIR}/${stage_name}.done to rebuild)"
        return 0
    fi
    
    log_info "Running stage: $stage_name"
    local start_time=$(date +%s)
    
    # Run the stage function with its arguments
    if "$stage_function" "${stage_args[@]}"; then
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        
        # Mark as completed with timing info
        stage_mark_completed "$stage_name" "Duration: ${duration}s, Function: $stage_function"
        log_info "Stage '$stage_name' completed successfully (${duration}s)"
        return 0
    else
        local exit_code=$?
        log_error "Stage '$stage_name' failed with exit code: $exit_code"
        return $exit_code
    fi
}

# Show stage status
function stages_status() {
    section_header "Build Stages Status"
    
    local all_stages=(
        "alpine_extract"
        "alpine_configure" 
        "alpine_packages"
        "alpine_firmware"
        "components_build"
        "components_verify"
        "kernel_modules"
        "zinit_setup"
        "init_script"
        "modules_setup"
        "modules_copy"
        "cleanup"
        "validation"
        "initramfs_create"
        "initramfs_test"
        "kernel_build"
    )
    
    log_info "Stage completion status:"
    
    for stage in "${all_stages[@]}"; do
        if stage_is_completed "$stage"; then
            local completion_time=$(stat -c %Y "${STAGES_DIR}/${stage}.done" 2>/dev/null || echo "0")
            local time_str=$(date -d @$completion_time '+%H:%M:%S' 2>/dev/null || echo "unknown")
            log_info "  [DONE] $stage (completed at $time_str)"
        else
            log_info "  [TODO] $stage (pending)"
        fi
    done
}

# Export functions
export -f stages_init stage_is_completed stage_mark_completed stage_force_rebuild
export -f stages_clear_all stage_run stages_status