#!/bin/bash
# Cleanup 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 common functions
source "${SCRIPT_DIR}/lib/common.sh"
source "${SCRIPT_DIR}/lib/stages.sh"

# Cleanup configuration
CLEAN_ALL="${CLEAN_ALL:-false}"
CLEAN_DOWNLOADS="${CLEAN_DOWNLOADS:-false}"
CLEAN_CONTAINER="${CLEAN_CONTAINER:-false}"
CLEAN_STAGES="${CLEAN_STAGES:-false}"

# Display usage information
function show_usage() {
    cat << EOF
Zero OS Build Cleanup Script (Incremental)

Usage: $0 [OPTIONS]

Options:
  --all                Clean everything (artifacts + downloads + containers + stages)
  --downloads          Clean downloaded sources and components
  --containers         Clean container images
  --stages             Clean stage completion markers only
  --artifacts-only     Clean only build artifacts (default)
  --help               Show this help message

Environment Variables:
  CLEAN_ALL            Clean everything (default: false)
  CLEAN_DOWNLOADS      Clean downloaded sources (default: false)
  CLEAN_CONTAINER      Clean container images (default: false)
  CLEAN_STAGES         Clean stage markers (default: false)

Examples:
  $0                   # Clean build artifacts only (preserves stages)
  $0 --all             # Complete cleanup (everything)
  $0 --stages          # Reset stage tracking (force full rebuild)
  $0 --downloads       # Clean sources and keep artifacts
EOF
}

# Parse command line arguments
function parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --all)
                CLEAN_ALL="true"
                CLEAN_DOWNLOADS="true"
                CLEAN_CONTAINER="true"
                CLEAN_STAGES="true"
                shift
                ;;
            --downloads)
                CLEAN_DOWNLOADS="true"
                shift
                ;;
            --containers)
                CLEAN_CONTAINER="true"
                shift
                ;;
            --stages)
                CLEAN_STAGES="true"
                shift
                ;;
            --artifacts-only)
                # This is the default, no action needed
                shift
                ;;
            --help|-h)
                show_usage
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_usage
                exit 1
                ;;
        esac
    done
}

# Clean build artifacts
function clean_build_artifacts() {
    section_header "Cleaning Build Artifacts"
    
    local artifacts_to_clean=(
        "${PROJECT_ROOT}/initramfs"
        "${PROJECT_ROOT}/dist"
        "${PROJECT_ROOT}/.build-stages"
    )
    
    for artifact in "${artifacts_to_clean[@]}"; do
        if [[ -d "$artifact" ]]; then
            log_info "Removing: $artifact"
            safe_rmdir "$artifact"
        else
            log_debug "Already clean: $artifact"
        fi
    done
    
    # Clean temporary files
    local temp_files=(
        "/tmp/alpine-miniroot*.tar.gz"
        "/tmp/linux-*.tar.xz"
        "/tmp/qemu-*.log"
        "/tmp/cloud-hypervisor-*.log"
    )
    
    for pattern in "${temp_files[@]}"; do
        if ls $pattern 2>/dev/null; then
            log_info "Removing temporary files: $pattern"
            safe_execute rm -f $pattern
        fi
    done
    
    log_info "Build artifacts cleaned"
}

# Clean downloaded sources and components
function clean_downloads() {
    section_header "Cleaning Downloaded Sources and Components"
    
    local download_dirs=(
        "${PROJECT_ROOT}/components"
        "${PROJECT_ROOT}/kernel"
    )
    
    for dir in "${download_dirs[@]}"; do
        if [[ -d "$dir" ]]; then
            log_info "Removing: $dir"
            safe_rmdir "$dir"
        else
            log_debug "Already clean: $dir"
        fi
    done
    
    # Clean Rust cache if it exists in project
    local rust_cache="${PROJECT_ROOT}/.cargo"
    if [[ -d "$rust_cache" ]]; then
        log_info "Removing Rust cache: $rust_cache"
        safe_rmdir "$rust_cache"
    fi
    
    log_info "Downloads and sources cleaned"
}

# Clean container images
function clean_container_images() {
    section_header "Cleaning Container Images"
    
    # Source docker functions if available
    if [[ -f "${SCRIPT_DIR}/lib/docker.sh" ]]; then
        source "${SCRIPT_DIR}/lib/docker.sh"
        
        # Detect container runtime
        if docker_detect_runtime 2>/dev/null; then
            docker_cleanup "false"
        else
            log_info "No container runtime detected"
        fi
    else
        log_warn "Docker library not found, manual container cleanup may be needed"
    fi
    
    log_info "Container images cleaned"
}

# Show disk space recovery
function show_space_recovery() {
    section_header "Disk Space Recovery"
    
    # Calculate space in current directory
    local current_usage=$(du -sh "${PROJECT_ROOT}" 2>/dev/null | cut -f1 || echo "unknown")
    log_info "Current project size: ${current_usage}"
    
    # Show what was cleaned
    if [[ "$CLEAN_ALL" == "true" ]]; then
        log_info "Complete cleanup performed:"
        log_info "  ✓ Build artifacts removed"
        log_info "  ✓ Downloaded sources removed"  
        log_info "  ✓ Container images removed"
    elif [[ "$CLEAN_DOWNLOADS" == "true" ]]; then
        log_info "Partial cleanup performed:"
        log_info "  ✓ Build artifacts removed"
        log_info "  ✓ Downloaded sources removed"
        log_info "  - Container images preserved"
    else
        log_info "Minimal cleanup performed:"
        log_info "  ✓ Build artifacts removed"
        log_info "  - Downloaded sources preserved"
        log_info "  - Container images preserved"
    fi
}

# Verify cleanup was successful
function verify_cleanup() {
    section_header "Verifying Cleanup"
    
    local remaining_artifacts=()
    
    # Check if artifacts were actually removed
    if [[ -d "${PROJECT_ROOT}/initramfs" ]]; then
        remaining_artifacts+=("initramfs/")
    fi
    
    if [[ -d "${PROJECT_ROOT}/dist" ]]; then
        remaining_artifacts+=("dist/")
    fi
    
    if [[ "$CLEAN_DOWNLOADS" == "true" ]]; then
        if [[ -d "${PROJECT_ROOT}/components" ]]; then
            remaining_artifacts+=("components/")
        fi
        
        if [[ -d "${PROJECT_ROOT}/kernel" ]]; then
            remaining_artifacts+=("kernel/")
        fi
    fi
    
    if [[ ${#remaining_artifacts[@]} -gt 0 ]]; then
        log_warn "Some artifacts may not have been cleaned:"
        for artifact in "${remaining_artifacts[@]}"; do
            log_warn "  - $artifact"
        done
        return 1
    else
        log_info "Cleanup verification passed"
        return 0
    fi
}

# Main function
function main() {
    # Parse command line arguments
    parse_arguments "$@"
    
    echo ""
    echo "=================================================="
    echo "==     ZERO-OS BUILD CLEANUP                   =="
    echo "=================================================="
    echo ""
    
    log_info "Starting cleanup process"
    log_info "Clean all: ${CLEAN_ALL}"
    log_info "Clean downloads: ${CLEAN_DOWNLOADS}"
    log_info "Clean containers: ${CLEAN_CONTAINER}"
    
    # Always clean build artifacts
    clean_build_artifacts
    
    # Clean downloads if requested
    if [[ "$CLEAN_DOWNLOADS" == "true" ]]; then
        clean_downloads
    fi
    
    # Clean containers if requested
    if [[ "$CLEAN_CONTAINER" == "true" ]]; then
        clean_container_images
    fi
    
    # Clean stage markers if requested (already handled in build artifacts for convenience)
    if [[ "$CLEAN_STAGES" == "true" ]]; then
        log_info "Stage markers already cleaned with build artifacts"
    fi
    
    # Show space recovery
    show_space_recovery
    
    # Verify cleanup
    verify_cleanup
    
    section_header "Cleanup Complete"
    log_info "Project cleaned successfully"
}

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