#!/bin/bash

# Housekeeping Script
# Unified system maintenance: package index update, system cleanup, and oh-my-zsh update.
# Combines the former update.sh and clean.sh into a single script with subcommand dispatch.

set -euo pipefail

#==============================================================================
# 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 CACHE_RETENTION_DAYS=7
readonly TEMP_DIRS=("/tmp" "/var/tmp")
readonly CACHE_DIRS=("/var/cache/apt" "/root/.cache")

readonly SCRIPT_NAME="$(basename "$0")"

#==============================================================================
# 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
# Deletes files older than CACHE_RETENTION_DAYS instead of wiping everything
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 to delete files older than retention period
    # -type f -mtime +N: only files older than N days (not directories)
    local deleted=false
    if find "$dir" -type f -mtime +"${CACHE_RETENTION_DAYS:-7}" -delete 2>/dev/null; then
        deleted=true
    fi

    # Clean up empty directories left behind (but don't remove the dir itself)
    find "$dir" -mindepth 1 -type d -empty -delete 2>/dev/null || true

    if $deleted; then
        local size_after
        size_after=$(get_dir_size "$dir")
        if [[ "$size_after" == "$size_before" ]]; then
            log_info "$description: nothing old enough to clean (retention: ${CACHE_RETENTION_DAYS:-7}d)"
        else
            log_success "$description: freed (was $size_before, now $size_after)"
        fi
    else
        log_warning "$description: partial cleanup completed"
    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"
}

#==============================================================================
# OS DETECTION (for package index update)
#==============================================================================

# Detect the running OS and its ID string.
# Prints two values: the distro id (lowercase) and the family id (lowercase).
detect_os() {
    local id=""
    local family=""

    if [[ -f /etc/os-release ]]; then
        # shellcheck disable=SC1091
        . /etc/os-release
        id="${ID:-unknown}"
        family="${ID_LIKE:-}"
    elif [[ -f /etc/alpine-release ]]; then
        id="alpine"
        family="alpine"
    else
        # Fall back to package manager detection
        if command_exists apk; then
            id="alpine"; family="alpine"
        elif command_exists pacman; then
            id="arch"; family="arch"
        elif command_exists apt-get; then
            id="debian"; family="debian"
        elif command_exists dnf; then
            id="fedora"; family="rhel"
        elif command_exists yum; then
            id="centos"; family="rhel"
        elif command_exists zypper; then
            id="opensuse"; family="suse"
        elif command_exists xbps-install; then
            id="void"; family="void"
        elif command_exists emerge; then
            id="gentoo"; family="gentoo"
        else
            id="unknown"; family="unknown"
        fi
    fi

    echo "${id} ${family}"
}

#==============================================================================
# OS-SPECIFIC PACKAGE INDEX UPDATE FUNCTIONS
#==============================================================================

update_alpine() {
    if ! command_exists apk; then
        log_warning "apk not found on Alpine-type system"
        return 0
    fi

    log_info "Running apk update..."
    if apk update 2>&1; then
        log_success "Alpine package index updated"
    else
        log_warning "apk update failed"
    fi
}

update_arch() {
    if ! command_exists pacman; then
        log_warning "pacman not found on Arch-type system"
        return 0
    fi

    log_info "Refreshing pacman database (pacman -Sy)..."
    # -S = sync, -y = refresh local database only (no package installation)
    if pacman -Sy --noconfirm 2>&1; then
        log_success "Arch package database refreshed"
    else
        log_warning "pacman -Sy failed"
    fi
}

update_debian() {
    if ! command_exists apt-get; then
        log_warning "apt-get not found on Debian-type system"
        return 0
    fi

    log_info "Running apt-get update..."
    if apt-get update 2>&1; then
        log_success "Debian/Ubuntu package index updated"
    else
        log_warning "apt-get update failed"
    fi
}

update_rhel() {
    if command_exists dnf; then
        log_info "Running dnf makecache..."
        if dnf makecache 2>&1; then
            log_success "RHEL/Fedora metadata cache refreshed (dnf)"
        else
            log_warning "dnf makecache failed"
        fi
    elif command_exists yum; then
        log_info "Running yum makecache..."
        if yum makecache fast 2>&1; then
            log_success "CentOS/RHEL metadata cache refreshed (yum)"
        else
            log_warning "yum makecache failed"
        fi
    else
        log_warning "No DNF or YUM found on RHEL-type system"
    fi
}

update_suse() {
    if ! command_exists zypper; then
        log_warning "zypper not found on openSUSE-type system"
        return 0
    fi

    log_info "Running zypper refresh..."
    if zypper --non-interactive refresh 2>&1; then
        log_success "openSUSE repositories refreshed"
    else
        log_warning "zypper refresh failed"
    fi
}

update_void() {
    if ! command_exists xbps-install; then
        log_warning "xbps-install not found on Void system"
        return 0
    fi

    log_info "Running xbps-install -S (sync repository index)..."
    if xbps-install -S 2>&1; then
        log_success "Void package index synced"
    else
        log_warning "xbps-install -S failed"
    fi
}

update_gentoo() {
    if ! command_exists emerge; then
        log_warning "emerge not found on Gentoo system"
        return 0
    fi

    log_info "Running emaint sync -a..."
    if emaint sync -a 2>&1; then
        log_success "Gentoo Portage tree synced"
    else
        log_warning "emaint sync failed"
    fi
}

#==============================================================================
# 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

    # 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
}

#==============================================================================
# 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 safe cache locations only
    # NOTE: We do NOT recursively delete /var/cache or /root/.cache as a whole,
    # as that can break package managers and remove important data.
    local additional_caches=(
        "/var/cache/apt/archives"
        "/var/cache/apt/archives/partial"
        "/var/cache/debconf"
        "/var/lib/apt/lists"
        "/root/.npm"
        "/root/.cache/pip"
        "/root/.cache/npm"
        "/home/*/.cache/pip"
        "/home/*/.cache/npm"
        "/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
            if rm -f "$logfile" 2>/dev/null; then
                ((++cleaned_count)) || true
            fi
        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
            if rm -f "$logfile" 2>/dev/null; then
                ((++cleaned_count)) || true
            fi
        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
}

#==============================================================================
# SUBCOMMAND HANDLERS
#==============================================================================

# update: refresh the OS package index (no upgrades)
run_update() {
    log_step "Starting package list update (${SCRIPT_NAME})"
    echo

    # Detect the OS
    local distro_info distro_id family_id
    distro_info="$(detect_os)"
    distro_id="${distro_info%% *}"
    family_id="${distro_info##* }"

    log_info "Detected OS: ${distro_id}"
    if [[ -n "${family_id}" && "${family_id}" != "${distro_id}" ]]; then
        log_info "OS family: ${family_id}"
    fi
    echo

    # Dispatch to the correct handler based on distro family
    # Check distro_id first, then fall back to family_id for aliases
    case "${distro_id}" in
        alpine)
            update_alpine
            ;;
        arch|artix|manjaro|garuda)
            update_arch
            ;;
        debian|ubuntu|linuxmint|kali|raspbian|pop)
            update_debian
            ;;
        fedora|rhel|centos|rocky|almalinux|amazon|oracle|clear-linux-os)
            update_rhel
            ;;
        opensuse*|suse|sles)
            update_suse
            ;;
        void)
            update_void
            ;;
        gentoo|funtoo)
            update_gentoo
            ;;
        *)
            # Try matching on family ID as fallback
            case "${family_id}" in
                *alpine*)
                    update_alpine
                    ;;
                *arch*)
                    update_arch
                    ;;
                *debian*|*ubuntu*)
                    update_debian
                    ;;
                *rhel*|*fedora*|*centos*)
                    update_rhel
                    ;;
                *suse*|*opensuse*)
                    update_suse
                    ;;
                *void*)
                    update_void
                    ;;
                *gentoo*)
                    update_gentoo
                    ;;
                *)
                    die "Unsupported OS: ${distro_id} (${family_id}). Cannot update package index."
                    ;;
            esac
            ;;
    esac

    echo
    log_success "Package list update completed!"
}

# clean: full system cleanup and maintenance
run_clean() {
    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!"
}

# omz: update oh-my-zsh (safe fallback if not installed)
run_omz() {
    log_step "Starting oh-my-zsh update (${SCRIPT_NAME})"
    echo

    if ! command_exists omz; then
        log_warning "oh-my-zsh (omz) is not installed on this system, skipping omz update"
        return 0
    fi

    log_info "Running omz update..."
    if omz update 2>&1; then
        log_success "oh-my-zsh updated"
    else
        log_warning "omz update failed"
    fi
}

# all: run update, clean, then omz in order
run_all() {
    run_update
    echo
    run_clean
    echo
    run_omz
}

# Print usage information
print_usage() {
    cat <<EOF
Usage: ${SCRIPT_NAME} [subcommand]

Unified system housekeeping script.

Subcommands:
  update    Update the OS package index (detects distro, updates correct package manager)
  clean     Run system cleanup (docker, package caches, temp/cache/log dirs, journal, thumbnails, memory)
  omz       Update oh-my-zsh (skipped safely if not installed)
  all       Run update, then clean, then omz (default when no subcommand is given)

If no subcommand is provided, 'all' is assumed.
EOF
}

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

main() {
    local subcommand="${1:-all}"

    case "$subcommand" in
        update)
            run_update
            ;;
        clean)
            run_clean
            ;;
        omz)
            run_omz
            ;;
        all)
            run_all
            ;;
        -h|--help|help)
            print_usage
            ;;
        *)
            log_warning "Unknown subcommand: '${subcommand}'"
            echo
            print_usage
            exit 0
            ;;
    esac
}

main "$@"