#!/bin/bash

set -e

# Disk space management script for particle-os builds
# This script helps manage disk space and clean up build artifacts

WORK_DIRS=(
    "/tmp/particle-os-build"
    "/tmp/test-*"
    "/tmp/particle-os-*"
    "/tmp/deb-bootc-*"
    "/tmp/test-sudo-fix*"
    "/tmp/test-improvements*"
)

BUILD_DIRS=(
    "work/*"
    "output/*"
    "*.img"
    "*.qcow2"
    "*.vmdk"
    "*.vdi"
)

# 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 "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

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

# Check disk space
check_disk_space() {
    local path="$1"
    local available_gb=$(df -BG "$path" | tail -1 | awk '{print $4}' | sed 's/G//')
    local used_percent=$(df -h "$path" | tail -1 | awk '{print $5}' | sed 's/%//')
    
    echo "Disk space for $path:"
    echo "  Available: ${available_gb}GB"
    echo "  Used: ${used_percent}%"
    
    if [ "$available_gb" -lt 5 ]; then
        log_warn "Low disk space: ${available_gb}GB available, 5GB recommended for builds"
        return 1
    else
        log_info "Sufficient disk space: ${available_gb}GB available"
        return 0
    fi
}

# Clean up work directories
cleanup_work_dirs() {
    log_info "Cleaning up work directories..."
    
    local cleaned=0
    for pattern in "${WORK_DIRS[@]}"; do
        for dir in $pattern; do
            if [ -d "$dir" ]; then
                log_info "Removing: $dir"
                sudo rm -rf "$dir"
                cleaned=$((cleaned + 1))
            fi
        done
    done
    
    if [ $cleaned -eq 0 ]; then
        log_info "No work directories to clean up"
    else
        log_info "Cleaned up $cleaned work directories"
    fi
}

# Clean up build artifacts
cleanup_build_artifacts() {
    log_info "Cleaning up build artifacts..."
    
    local cleaned=0
    for pattern in "${BUILD_DIRS[@]}"; do
        for file in $pattern; do
            if [ -e "$file" ]; then
                log_info "Removing: $file"
                rm -rf "$file"
                cleaned=$((cleaned + 1))
            fi
        done
    done
    
    if [ $cleaned -eq 0 ]; then
        log_info "No build artifacts to clean up"
    else
        log_info "Cleaned up $cleaned build artifacts"
    fi
}

# Clean up package caches
cleanup_package_caches() {
    log_info "Cleaning up package caches..."
    
    # Clean apt cache
    if command -v apt >/dev/null 2>&1; then
        log_info "Cleaning apt cache..."
        sudo apt clean
        sudo apt autoremove -y
    fi
    
    # Clean podman cache
    if command -v podman >/dev/null 2>&1; then
        log_info "Cleaning podman cache..."
        podman system prune -f
    fi
    
    # Clean docker cache
    if command -v docker >/dev/null 2>&1; then
        log_info "Cleaning docker cache..."
        docker system prune -f
    fi
}

# Find large files and directories
find_large_files() {
    local path="$1"
    local size="${2:-100M}"
    
    log_info "Finding files larger than $size in $path..."
    
    if [ -d "$path" ]; then
        find "$path" -type f -size "+$size" -exec ls -lh {} \; 2>/dev/null | head -20
    else
        log_warn "Path $path does not exist"
    fi
}

# Create custom work directory with more space
create_custom_work_dir() {
    local custom_dir="$1"
    
    if [ -z "$custom_dir" ]; then
        custom_dir="/tmp/particle-os-custom-$(date +%s)"
    fi
    
    log_info "Creating custom work directory: $custom_dir"
    
    if mkdir -p "$custom_dir"; then
        log_info "Custom work directory created successfully"
        echo "Use this directory for builds:"
        echo "  ./bib/particle-os build --work-dir $custom_dir recipes/minimal-debug.yml"
        echo ""
        echo "Directory: $custom_dir"
    else
        log_error "Failed to create custom work directory"
        return 1
    fi
}

# Show help
show_help() {
    cat << EOF
Disk Space Management Script for particle-os

Usage: $0 [COMMAND] [OPTIONS]

Commands:
    check [PATH]           Check disk space for PATH (default: /tmp)
    cleanup                Clean up all work directories and build artifacts
    cleanup-work           Clean up only work directories
    cleanup-builds         Clean up only build artifacts
    cleanup-caches         Clean up package and container caches
    find-large [PATH]      Find large files in PATH (default: current directory)
    create-work-dir [DIR]  Create custom work directory with more space
    status                 Show current disk space status

Options:
    -h, --help            Show this help message

Examples:
    $0 check /tmp
    $0 cleanup
    $0 create-work-dir /home/joe/particle-os-builds
    $0 find-large /tmp

EOF
}

# Main execution
case "${1:-status}" in
    "check")
        path="${2:-/tmp}"
        check_disk_space "$path"
        ;;
    "cleanup")
        cleanup_work_dirs
        cleanup_build_artifacts
        cleanup_package_caches
        check_disk_space "/tmp"
        ;;
    "cleanup-work")
        cleanup_work_dirs
        ;;
    "cleanup-builds")
        cleanup_build_artifacts
        ;;
    "cleanup-caches")
        cleanup_package_caches
        ;;
    "find-large")
        path="${2:-.}"
        size="${3:-100M}"
        find_large_files "$path" "$size"
        ;;
    "create-work-dir")
        create_custom_work_dir "$2"
        ;;
    "status")
        echo "=== Disk Space Status ==="
        check_disk_space "/tmp"
        echo ""
        check_disk_space "/opt/Projects/deb-bootc-image-builder"
        echo ""
        echo "=== Work Directory Status ==="
        for pattern in "${WORK_DIRS[@]}"; do
            for dir in $pattern; do
                if [ -d "$dir" ]; then
                    size=$(du -sh "$dir" 2>/dev/null | cut -f1)
                    echo "  $dir: $size"
                fi
            done
        done
        ;;
    "-h"|"--help"|"help")
        show_help
        ;;
    *)
        log_error "Unknown command: $1"
        show_help
        exit 1
        ;;
esac