#!/bin/bash

# System Cleanup and Maintenance Script
# 
# Description: Comprehensive system cleanup for Docker containers and Linux systems
# Features:
#   - Docker resource cleanup (images, containers, volumes, networks)
#   - Package manager cache cleanup (APK/APT)
#   - System cache and temporary file cleanup
#   - Log rotation and cleanup
#   - Memory cache optimization
#   - Journal cleanup (systemd)
#   - Thumbnail and user cache cleanup
# Author: ivanch
# Version: 2.0

set -euo pipefail  # Exit on error, undefined vars, and pipe failures

#==============================================================================
# CONFIGURATION
#==============================================================================

# Color definitions for output formatting
readonly NC='\033[0m'
readonly RED='\033[1;31m'
readonly GREEN='\033[1;32m'
readonly LIGHT_GREEN='\033[1;32m'
readonly LIGHT_BLUE='\033[1;34m'
readonly LIGHT_GREY='\033[0;37m'
readonly YELLOW='\033[1;33m'

# Cleanup configuration
readonly LOG_RETENTION_DAYS=30
readonly JOURNAL_RETENTION_DAYS=7
readonly TEMP_DIRS=("/tmp" "/var/tmp")
readonly CACHE_DIRS=("/var/cache" "/root/.cache")

# Auto-update configuration
readonly AUTO_UPDATE_ENABLED=true

#==============================================================================
# UTILITY FUNCTIONS
#==============================================================================

# Print formatted log messages
log_info() { echo -e "${LIGHT_GREY}[i] $1${NC}"; }
log_success() { echo -e "${LIGHT_GREEN}[✓] $1${NC}"; }
log_warning() { echo -e "${YELLOW}[!] $1${NC}"; }
log_error() { echo -e "${RED}[x] $1${NC}" >&2; }
log_step() { echo -e "${LIGHT_BLUE}[i] $1${NC}"; }

# Exit with error message
die() {
    log_error "$1"
    exit 1
}

# Check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Get directory size in human readable format
get_dir_size() {
    local dir="$1"
    if [[ -d "$dir" ]]; then
        du -sh "$dir" 2>/dev/null | cut -f1 || echo "0B"
    else
        echo "0B"
    fi
}

# Safe directory cleanup with size reporting
clean_directory() {
    local dir="$1"
    local description="$2"
    
    if [[ ! -d "$dir" ]]; then
        return 0
    fi
    
    local size_before
    size_before=$(get_dir_size "$dir")
    
    if [[ "$size_before" == "0B" ]]; then
        log_info "$description: already clean"
        return 0
    fi
    
    log_step "$description (was $size_before)..."
    
    # Use find with -delete for safer cleanup
    if find "$dir" -mindepth 1 -delete 2>/dev/null; then
        log_success "$description: freed $size_before"
    else
        # Fallback to rm if find fails
        if rm -rf "$dir"/* 2>/dev/null; then
            log_success "$description: freed $size_before"
        else
            log_warning "$description: partial cleanup completed"
        fi
    fi
}

# Get system information for reporting
get_system_info() {
    local info=""
    
    # Memory info
    if [[ -f /proc/meminfo ]]; then
        local mem_total mem_available
        mem_total=$(grep MemTotal /proc/meminfo | awk '{print $2}')
        mem_available=$(grep MemAvailable /proc/meminfo | awk '{print $2}')
        if [[ -n "$mem_total" && -n "$mem_available" ]]; then
            info+="Memory: $((mem_available/1024))MB available of $((mem_total/1024))MB total"
        fi
    fi
    
    # Disk space info
    if command_exists df; then
        local disk_info
        disk_info=$(df -h / 2>/dev/null | tail -1 | awk '{print $4 " available of " $2 " total"}')
        if [[ -n "$disk_info" ]]; then
            info+="${info:+, }Disk: $disk_info"
        fi
    fi
    
    echo "$info"
}

#==============================================================================
# DOCKER CLEANUP FUNCTIONS
#==============================================================================

# Clean Docker resources
cleanup_docker() {
    if ! command_exists docker; then
        log_info "Docker not found, skipping Docker cleanup"
        return 0
    fi
    
    log_step "Starting Docker cleanup..."
    
    # Check if Docker daemon is running
    if ! docker info >/dev/null 2>&1; then
        log_warning "Docker daemon not running, skipping Docker cleanup"
        return 0
    fi
    
    # Get initial Docker disk usage
    local docker_usage_before=""
    if docker system df >/dev/null 2>&1; then
        docker_usage_before=$(docker system df 2>/dev/null || echo "")
    fi
    
    # Remove unused images
    log_info "Removing unused Docker images..."
    if docker image prune -af >/dev/null 2>&1; then
        log_success "Docker images cleaned"
    else
        log_warning "Docker image cleanup failed"
    fi
    
    # Remove stopped containers
    log_info "Removing stopped Docker containers..."
    if docker container prune -f >/dev/null 2>&1; then
        log_success "Docker containers cleaned"
    else
        log_warning "Docker container cleanup failed"
    fi
    
    # Remove unused volumes
    log_info "Removing unused Docker volumes..."
    if docker volume prune -f >/dev/null 2>&1; then
        log_success "Docker volumes cleaned"
    else
        log_warning "Docker volume cleanup failed"
    fi
    
    # Remove unused networks
    log_info "Removing unused Docker networks..."
    if docker network prune -f >/dev/null 2>&1; then
        log_success "Docker networks cleaned"
    else
        log_warning "Docker network cleanup failed"
    fi
    
    # Complete system cleanup
    log_info "Running Docker system cleanup..."
    if docker system prune -af >/dev/null 2>&1; then
        log_success "Docker system cleanup completed"
    else
        log_warning "Docker system cleanup failed"
    fi
    
    # Show space freed if possible
    if [[ -n "$docker_usage_before" ]] && docker system df >/dev/null 2>&1; then
        log_info "Docker cleanup completed"
    fi
}

#==============================================================================
# PACKAGE MANAGER CLEANUP FUNCTIONS
#==============================================================================

# Clean APK cache (Alpine Linux)
cleanup_apk() {
    if ! command_exists apk; then
        return 0
    fi
    
    log_step "Cleaning APK cache..."
    
    # Clean APK cache
    if [[ -d /var/cache/apk ]]; then
        clean_directory "/var/cache/apk" "APK cache directory"
    fi
    
    # Clean APK cache using apk command
    if apk cache clean >/dev/null 2>&1; then
        log_success "APK cache cleaned"
    fi
    
    # Update package index
    log_info "Updating APK package index..."
    if apk update >/dev/null 2>&1; then
        log_success "APK index updated"
    else
        log_warning "APK index update failed"
    fi
}

# Clean APT cache (Debian/Ubuntu)
cleanup_apt() {
    if ! command_exists apt-get; then
        return 0
    fi
    
    log_step "Cleaning APT cache..."
    
    # Clean downloaded packages
    if apt-get clean >/dev/null 2>&1; then
        log_success "APT cache cleaned"
    else
        log_warning "APT clean failed"
    fi
    
    # Remove orphaned packages
    if apt-get autoclean >/dev/null 2>&1; then
        log_success "APT autocleaned"
    else
        log_warning "APT autoclean failed"
    fi
    
    # Remove unnecessary packages
    if apt-get autoremove -y >/dev/null 2>&1; then
        log_success "Unnecessary packages removed"
    else
        log_warning "APT autoremove failed"
    fi
    
    # Update package index
    log_info "Updating APT package index..."
    if apt-get update >/dev/null 2>&1; then
        log_success "APT index updated"
    else
        log_warning "APT index update failed"
    fi
}

#==============================================================================
# SYSTEM CLEANUP FUNCTIONS
#==============================================================================

# Clean system temporary directories
cleanup_temp_dirs() {
    log_step "Cleaning temporary directories..."
    
    for temp_dir in "${TEMP_DIRS[@]}"; do
        if [[ -d "$temp_dir" ]]; then
            # Clean contents but preserve the directory
            find "$temp_dir" -mindepth 1 -maxdepth 1 -mtime +1 -exec rm -rf {} + 2>/dev/null || true
            log_success "Cleaned old files in $temp_dir"
        fi
    done
}

# Clean system cache directories
cleanup_cache_dirs() {
    log_step "Cleaning cache directories..."
    
    for cache_dir in "${CACHE_DIRS[@]}"; do
        if [[ -d "$cache_dir" ]]; then
            clean_directory "$cache_dir" "Cache directory $cache_dir"
        fi
    done
    
    # Clean additional cache locations
    local additional_caches=(
        "/var/lib/apt/lists"
        "/var/cache/debconf"
        "/root/.npm"
        "/root/.pip"
        "/home/*/.cache"
        "/home/*/.npm"
        "/home/*/.pip"
    )
    
    for cache_pattern in "${additional_caches[@]}"; do
        # Use shell expansion for patterns
        for cache_path in $cache_pattern; do
            if [[ -d "$cache_path" ]]; then
                clean_directory "$cache_path" "Additional cache $cache_path"
            fi
        done 2>/dev/null || true
    done
}

# Clean old log files
cleanup_logs() {
    log_step "Cleaning old log files..."
    
    # Clean logs older than retention period
    if [[ -d /var/log ]]; then
        local cleaned_count=0
        
        # Find and remove old log files
        while IFS= read -r -d '' logfile; do
            rm -f "$logfile" 2>/dev/null && ((cleaned_count++))
        done < <(find /var/log -type f -name "*.log" -mtime +"$LOG_RETENTION_DAYS" -print0 2>/dev/null || true)
        
        # Clean compressed logs
        while IFS= read -r -d '' logfile; do
            rm -f "$logfile" 2>/dev/null && ((cleaned_count++))
        done < <(find /var/log -type f \( -name "*.log.gz" -o -name "*.log.bz2" -o -name "*.log.xz" \) -mtime +"$LOG_RETENTION_DAYS" -print0 2>/dev/null || true)
        
        if [[ $cleaned_count -gt 0 ]]; then
            log_success "Removed $cleaned_count old log files"
        else
            log_info "No old log files to remove"
        fi
    fi
    
    # Truncate large active log files
    local large_logs
    while IFS= read -r -d '' logfile; do
        if [[ -f "$logfile" && -w "$logfile" ]]; then
            truncate -s 0 "$logfile" 2>/dev/null || true
        fi
    done < <(find /var/log -type f -name "*.log" -size +100M -print0 2>/dev/null || true)
}

# Clean systemd journal
cleanup_journal() {
    if ! command_exists journalctl; then
        return 0
    fi
    
    log_step "Cleaning systemd journal..."
    
    # Clean journal older than retention period
    if journalctl --vacuum-time="${JOURNAL_RETENTION_DAYS}d" >/dev/null 2>&1; then
        log_success "Journal cleaned (older than $JOURNAL_RETENTION_DAYS days)"
    else
        log_warning "Journal cleanup failed"
    fi
    
    # Limit journal size
    if journalctl --vacuum-size=100M >/dev/null 2>&1; then
        log_success "Journal size limited to 100MB"
    fi
}

# Clean thumbnail caches
cleanup_thumbnails() {
    log_step "Cleaning thumbnail caches..."
    
    local thumbnail_dirs=(
        "/root/.thumbnails"
        "/root/.cache/thumbnails"
        "/home/*/.thumbnails"
        "/home/*/.cache/thumbnails"
    )
    
    for thumb_pattern in "${thumbnail_dirs[@]}"; do
        for thumb_dir in $thumb_pattern; do
            if [[ -d "$thumb_dir" ]]; then
                clean_directory "$thumb_dir" "Thumbnail cache $thumb_dir"
            fi
        done 2>/dev/null || true
    done
}

# Optimize memory caches
optimize_memory() {
    log_step "Optimizing memory caches..."
    
    # Sync filesystem
    if sync; then
        log_info "Filesystem synced"
    fi
    
    # Drop caches (page cache, dentries and inodes)
    if [[ -w /proc/sys/vm/drop_caches ]]; then
        echo 3 > /proc/sys/vm/drop_caches 2>/dev/null && log_success "Memory caches dropped" || log_warning "Failed to drop memory caches"
    fi
}

#==============================================================================
# REPORTING FUNCTIONS
#==============================================================================

# Generate cleanup summary
generate_summary() {
    log_step "Generating cleanup summary..."
    
    local system_info
    system_info=$(get_system_info)
    
    if [[ -n "$system_info" ]]; then
        log_info "System status: $system_info"
    fi
    
    # Show disk usage of important directories
    local important_dirs=("/" "/var" "/tmp" "/var/log" "/var/cache")
    for dir in "${important_dirs[@]}"; do
        if [[ -d "$dir" ]]; then
            local usage
            usage=$(df -h "$dir" 2>/dev/null | tail -1 | awk '{print $5 " used (" $4 " available)"}' || echo "unknown")
            log_info "$dir: $usage"
        fi
    done
}

#==============================================================================
# MAIN EXECUTION
#==============================================================================

main() {
    log_step "Starting System Cleanup and Maintenance"
    echo
    
    # Show initial system status
    local initial_info
    initial_info=$(get_system_info)
    if [[ -n "$initial_info" ]]; then
        log_info "Initial system status: $initial_info"
        echo
    fi
    
    # Docker cleanup
    cleanup_docker
    
    # Package manager cleanup
    cleanup_apk
    cleanup_apt
    
    # System cleanup
    cleanup_temp_dirs
    cleanup_cache_dirs
    cleanup_logs
    cleanup_journal
    cleanup_thumbnails
    
    # Memory optimization
    optimize_memory
    
    # Generate summary
    echo
    generate_summary
    
    echo
    log_success "System cleanup and maintenance completed!"
}

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