#!/bin/bash

# Comprehensive Benchmark Cleanup Script
# Cleans up all benchmark-related artifacts and temporary files

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
LOCAL_ARTIFACTS_DIR="$PROJECT_ROOT/.local-artifacts"

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

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

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

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

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

# Cleanup functions
cleanup_benchmark_artifacts() {
    print_status "Cleaning benchmark artifacts..."
    
    if [ -d "$LOCAL_ARTIFACTS_DIR" ]; then
        local count_before=$(ls -1d "$LOCAL_ARTIFACTS_DIR"/run_* 2>/dev/null | wc -l)
        local size_before=$(du -sh "$LOCAL_ARTIFACTS_DIR" 2>/dev/null | cut -f1 || echo "0B")
        
        case "${1:-keep-10}" in
            "all")
                print_status "Removing ALL benchmark artifacts..."
                rm -rf "$LOCAL_ARTIFACTS_DIR"
                print_success "Removed all artifacts (was $size_before)"
                ;;
            "keep-5")
                print_status "Keeping last 5 runs, removing older ones..."
                cd "$LOCAL_ARTIFACTS_DIR"
                ls -1t run_* 2>/dev/null | tail -n +6 | xargs rm -rf 2>/dev/null || true
                local count_after=$(ls -1d run_* 2>/dev/null | wc -l)
                local size_after=$(du -sh . 2>/dev/null | cut -f1 || echo "0B")
                print_success "Cleaned old runs: $count_before → $count_after runs ($size_before → $size_after)"
                ;;
            "keep-10")
                print_status "Keeping last 10 runs, removing older ones..."
                cd "$LOCAL_ARTIFACTS_DIR"
                ls -1t run_* 2>/dev/null | tail -n +11 | xargs rm -rf 2>/dev/null || true
                local count_after=$(ls -1d run_* 2>/dev/null | wc -l)
                local size_after=$(du -sh . 2>/dev/null | cut -f1 || echo "0B")
                print_success "Cleaned old runs: $count_before → $count_after runs ($size_before → $size_after)"
                ;;
            *)
                print_error "Invalid cleanup level: ${1}"
                print_status "Valid options: all, keep-5, keep-10"
                return 1
                ;;
        esac
    else
        print_warning "No benchmark artifacts directory found"
    fi
}

cleanup_temp_files() {
    print_status "Cleaning temporary files..."
    
    # Clean temp directories
    local temp_cleaned=0
    
    # /tmp cleanup
    if [ -d "/tmp" ]; then
        local tmp_files=$(find /tmp -name "benchmark_*" -type f 2>/dev/null | wc -l)
        if [ "$tmp_files" -gt 0 ]; then
            find /tmp -name "benchmark_*" -type f -mmin +60 -delete 2>/dev/null || true
            print_success "Cleaned $tmp_files temporary files from /tmp"
            temp_cleaned=$((temp_cleaned + tmp_files))
        fi
    fi
    
    # /var/tmp cleanup
    if [ -d "/var/tmp" ]; then
        local vartmp_files=$(find /var/tmp -name "benchmark_*" -type f 2>/dev/null | wc -l)
        if [ "$vartmp_files" -gt 0 ]; then
            find /var/tmp -name "benchmark_*" -type f -mmin +60 -delete 2>/dev/null || true
            print_success "Cleaned $vartmp_files temporary files from /var/tmp"
            temp_cleaned=$((temp_cleaned + vartmp_files))
        fi
    fi
    
    # User temp cleanup
    if [ -d "$HOME/tmp" ]; then
        local user_tmp_files=$(find "$HOME/tmp" -name "benchmark_*" -type f 2>/dev/null | wc -l)
        if [ "$user_tmp_files" -gt 0 ]; then
            find "$HOME/tmp" -name "benchmark_*" -type f -mmin +60 -delete 2>/dev/null || true
            print_success "Cleaned $user_tmp_files temporary files from ~/tmp"
            temp_cleaned=$((temp_cleaned + user_tmp_files))
        fi
    fi
    
    if [ "$temp_cleaned" -eq 0 ]; then
        print_status "No temporary files to clean"
    fi
}

cleanup_go_cache() {
    print_status "Cleaning Go build cache..."
    
    # Clean test cache
    if command -v go >/dev/null 2>&1; then
        local cache_before=$(go env GOCACHE 2>/dev/null || echo "")
        if [ -n "$cache_before" ] && [ -d "$cache_before" ]; then
            local cache_size_before=$(du -sh "$cache_before" 2>/dev/null | cut -f1 || echo "0B")
            go clean -testcache 2>/dev/null || true
            local cache_size_after=$(du -sh "$cache_before" 2>/dev/null | cut -f1 || echo "0B")
            print_success "Cleaned Go test cache: $cache_size_before → $cache_size_after"
        fi
        
        # Clean build cache (optional, more aggressive)
        if [ "${1:-}" = "aggressive" ]; then
            go clean -cache 2>/dev/null || true
            print_success "Cleaned Go build cache (aggressive)"
        fi
    else
        print_warning "Go not found, skipping cache cleanup"
    fi
}

cleanup_docker() {
    print_status "Cleaning Docker artifacts..."
    
    if command -v docker >/dev/null 2>&1; then
        # Clean stopped containers with benchmark labels
        local containers_removed=$(docker container prune -f --filter "label=benchmark" 2>/dev/null | grep "Total reclaimed space" | cut -d: -f2 | tr -d ' ' || echo "0B")
        if [ "$containers_removed" != "0B" ]; then
            print_success "Cleaned Docker containers: $containers_removed"
        fi
        
        # Clean unused images (aggressive mode only)
        if [ "${1:-}" = "aggressive" ]; then
            local images_removed=$(docker image prune -f 2>/dev/null | grep "Total reclaimed space" | cut -d: -f2 | tr -d ' ' || echo "0B")
            if [ "$images_removed" != "0B" ]; then
                print_success "Cleaned Docker images: $images_removed"
            fi
        fi
        
        # Clean unused volumes (aggressive mode only)
        if [ "${1:-}" = "aggressive" ]; then
            local volumes_removed=$(docker volume prune -f 2>/dev/null | grep "Total reclaimed space" | cut -d: -f2 | tr -d ' ' || echo "0B")
            if [ "$volumes_removed" != "0B" ]; then
                print_success "Cleaned Docker volumes: $volumes_removed"
            fi
        fi
    else
        print_warning "Docker not found, skipping Docker cleanup"
    fi
}

cleanup_logs() {
    print_status "Cleaning log files..."
    
    # Clean old log files
    local log_dirs=("$PROJECT_ROOT/logs" "$HOME/.local/share/fetch_ml/logs")
    
    for log_dir in "${log_dirs[@]}"; do
        if [ -d "$log_dir" ]; then
            local log_size_before=$(du -sh "$log_dir" 2>/dev/null | cut -f1 || echo "0B")
            # Remove log files older than 7 days
            find "$log_dir" -name "*.log" -type f -mtime +7 -delete 2>/dev/null || true
            find "$log_dir" -name "*.log.*" -type f -mtime +7 -delete 2>/dev/null || true
            local log_size_after=$(du -sh "$log_dir" 2>/dev/null | cut -f1 || echo "0B")
            if [ "$log_size_before" != "$log_size_after" ]; then
                print_success "Cleaned old logs in $log_dir: $log_size_before → $log_size_after"
            fi
        fi
    done
}

show_disk_usage() {
    print_status "Current disk usage:"
    
    # Project root
    if [ -d "$PROJECT_ROOT" ]; then
        local project_size=$(du -sh "$PROJECT_ROOT" 2>/dev/null | cut -f1 || echo "N/A")
        echo "  Project: $project_size"
    fi
    
    # Artifacts
    if [ -d "$LOCAL_ARTIFACTS_DIR" ]; then
        local artifacts_size=$(du -sh "$LOCAL_ARTIFACTS_DIR" 2>/dev/null | cut -f1 || echo "0B")
        local artifacts_count=$(ls -1d "$LOCAL_ARTIFACTS_DIR"/run_* 2>/dev/null | wc -l)
        echo "  Artifacts: $artifacts_size ($artifacts_count runs)"
    fi
    
    # Go cache
    if command -v go >/dev/null 2>&1; then
        local gocache=$(go env GOCACHE 2>/dev/null || echo "")
        if [ -n "$gocache" ] && [ -d "$gocache" ]; then
            local cache_size=$(du -sh "$gocache" 2>/dev/null | cut -f1 || echo "0B")
            echo "  Go cache: $cache_size"
        fi
    fi
}

# Main execution
main() {
    echo "=== Benchmark Cleanup Script ==="
    echo ""
    
    case "${1:-help}" in
        "artifacts")
            cleanup_benchmark_artifacts "${2:-keep-10}"
            ;;
        "temp")
            cleanup_temp_files
            ;;
        "go")
            cleanup_go_cache "${2:-}"
            ;;
        "docker")
            cleanup_docker "${2:-}"
            ;;
        "logs")
            cleanup_logs
            ;;
        "benchmarks")
            print_status "Running standard benchmark cleanup..."
            cleanup_benchmark_artifacts "keep-10"
            cleanup_temp_files
            cleanup_go_cache
            ;;
        "all")
            print_status "Running comprehensive cleanup..."
            cleanup_benchmark_artifacts "keep-5"
            cleanup_temp_files
            cleanup_go_cache
            cleanup_docker
            cleanup_logs
            ;;
        "aggressive")
            print_warning "Running AGGRESSIVE cleanup - this will remove more data!"
            read -p "Are you sure? [y/N] " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                cleanup_benchmark_artifacts "keep-5"
                cleanup_temp_files
                cleanup_go_cache "aggressive"
                cleanup_docker "aggressive"
                cleanup_logs
            else
                print_status "Aggressive cleanup cancelled"
                exit 0
            fi
            ;;
        "status")
            show_disk_usage
            ;;
        "help"|*)
            echo "Benchmark Cleanup Script"
            echo ""
            echo "Usage: $0 <command> [options]"
            echo ""
            echo "Commands:"
            echo "  artifacts [all|keep-5|keep-10]  Clean benchmark artifacts"
            echo "  temp                          Clean temporary files"
            echo "  go [aggressive]               Clean Go build cache"
            echo "  docker [aggressive]           Clean Docker artifacts"
            echo "  logs                          Clean old log files"
            echo "  benchmarks                    Standard benchmark cleanup"
            echo "  all                           Comprehensive cleanup"
            echo "  aggressive                    Aggressive cleanup (more data removed)"
            echo "  status                        Show current disk usage"
            echo "  help                          Show this help"
            echo ""
            echo "Examples:"
            echo "  $0 benchmarks                 # Standard cleanup"
            echo "  $0 artifacts keep-5            # Keep last 5 runs"
            echo "  $0 all                        # Comprehensive cleanup"
            echo "  $0 aggressive                 # Aggressive cleanup"
            ;;
    esac
    
    echo ""
    print_success "Cleanup completed!"
    
    # Show final status
    if [ "${1:-}" != "status" ] && [ "${1:-}" != "help" ]; then
        echo ""
        show_disk_usage
    fi
}

main "$@"