#!/usr/bin/env bash

#########################################################################################
#### disk-cleanup.sh — Find and clean disk space hogs on Linux servers               ####
#### Scans logs, temp files, package caches, old kernels, journal, and Docker cruft   ####
#### Dry-run by default — nothing is deleted without --force                          ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./disk-cleanup.sh --scan                                                      ####
####   ./disk-cleanup.sh --clean --force                                             ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
LOG_AGE_DAYS="${LOG_AGE_DAYS:-30}"
TMP_AGE_DAYS="${TMP_AGE_DAYS:-7}"
JOURNAL_MAX="${JOURNAL_MAX:-500M}"
LARGE_FILE_MIN="${LARGE_FILE_MIN:-100M}"
LARGE_FILE_DIRS="${LARGE_FILE_DIRS:-/var /home /opt /tmp /srv}"
DRY_RUN="${DRY_RUN:-true}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
TOTAL_RECLAIMABLE=0
TOTAL_CLEANED=0

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

human_bytes() {
    local bytes="$1"
    if [[ "$bytes" -ge 1073741824 ]]; then
        awk "BEGIN { printf \"%.1f GiB\", $bytes / 1073741824 }"
    elif [[ "$bytes" -ge 1048576 ]]; then
        awk "BEGIN { printf \"%.1f MiB\", $bytes / 1048576 }"
    elif [[ "$bytes" -ge 1024 ]]; then
        awk "BEGIN { printf \"%.1f KiB\", $bytes / 1024 }"
    else
        echo "${bytes} B"
    fi
}

add_reclaimable() {
    TOTAL_RECLAIMABLE=$((TOTAL_RECLAIMABLE + $1))
}

add_cleaned() {
    TOTAL_CLEANED=$((TOTAL_CLEANED + $1))
}

# ══════════════════════════════════════════════════════════════════════
#                       OLD LOGS
# ══════════════════════════════════════════════════════════════════════

scan_old_logs() {
    section_header "Old Log Files (> ${LOG_AGE_DAYS} days)"

    local total_size=0
    local count=0

    while IFS= read -r -d '' file; do
        local size
        size=$(stat -c%s "$file" 2>/dev/null || echo 0)
        if [[ "$size" -gt 0 ]]; then
            total_size=$((total_size + size))
            ((count++)) || true
            if [[ "$VERBOSE" == "true" ]]; then
                printf "    %10s  %s\n" "$(human_bytes "$size")" "$file"
            fi
        fi
    done < <(find /var/log -type f \( -name "*.gz" -o -name "*.xz" -o -name "*.bz2" -o -name "*.[0-9]" -o -name "*.old" \) -mtime +"$LOG_AGE_DAYS" -print0 2>/dev/null)

    # Rotated logs without compression
    while IFS= read -r -d '' file; do
        local size
        size=$(stat -c%s "$file" 2>/dev/null || echo 0)
        if [[ "$size" -gt 0 ]]; then
            total_size=$((total_size + size))
            ((count++)) || true
        fi
    done < <(find /var/log -type f -name "*.log.*" -mtime +"$LOG_AGE_DAYS" -print0 2>/dev/null)

    printf "  %-30s %s (%d files)\n" "Rotated/old logs:" "$(human_bytes "$total_size")" "$count"
    add_reclaimable "$total_size"
}

clean_old_logs() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would delete old log files in /var/log"
        return
    fi

    local cleaned=0
    find /var/log -type f \( -name "*.gz" -o -name "*.xz" -o -name "*.bz2" -o -name "*.[0-9]" -o -name "*.old" \) -mtime +"$LOG_AGE_DAYS" -print0 2>/dev/null | while IFS= read -r -d '' file; do
        local size
        size=$(stat -c%s "$file" 2>/dev/null || echo 0)
        rm -f "$file" && cleaned=$((cleaned + size))
    done

    find /var/log -type f -name "*.log.*" -mtime +"$LOG_AGE_DAYS" -delete 2>/dev/null || true
    log "Cleaned old log files"
}

# ══════════════════════════════════════════════════════════════════════
#                       JOURNAL
# ══════════════════════════════════════════════════════════════════════

scan_journal() {
    section_header "Systemd Journal"

    if ! command -v journalctl &>/dev/null; then
        printf "  %-30s %s\n" "Journal:" "N/A (no systemd)"
        return
    fi

    local journal_size
    journal_size=$(journalctl --disk-usage 2>/dev/null | grep -oP '[\d.]+[GMKT]' | head -1 || echo "0")

    # Get bytes for tracking
    local journal_bytes
    journal_bytes=$(du -sb /var/log/journal/ 2>/dev/null | awk '{print $1}' || echo "0")
    if [[ "$journal_bytes" -eq 0 ]]; then
        journal_bytes=$(du -sb /run/log/journal/ 2>/dev/null | awk '{print $1}' || echo "0")
    fi

    printf "  %-30s %s\n" "Journal size:" "${journal_size:-Unknown}"
    printf "  %-30s %s\n" "Would vacuum to:" "$JOURNAL_MAX"

    # Estimate savings
    local max_bytes=0
    local max_num="${JOURNAL_MAX%[GMKT]*}"
    local max_unit="${JOURNAL_MAX: -1}"
    case "$max_unit" in
        G) max_bytes=$((max_num * 1073741824)) ;;
        M) max_bytes=$((max_num * 1048576)) ;;
        K) max_bytes=$((max_num * 1024)) ;;
        *) max_bytes=$((max_num)) ;;
    esac

    if [[ "$journal_bytes" -gt "$max_bytes" ]]; then
        local savings=$((journal_bytes - max_bytes))
        add_reclaimable "$savings"
        printf "  %-30s %s\n" "Reclaimable:" "$(human_bytes "$savings")"
    fi
}

clean_journal() {
    if ! command -v journalctl &>/dev/null; then
        return
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would vacuum journal to ${JOURNAL_MAX}"
        return
    fi

    journalctl --vacuum-size="$JOURNAL_MAX" 2>/dev/null || warn "Journal vacuum failed"
    log "Vacuumed journal to ${JOURNAL_MAX}"
}

# ══════════════════════════════════════════════════════════════════════
#                       TEMP FILES
# ══════════════════════════════════════════════════════════════════════

scan_tmp() {
    section_header "Temp Files (> ${TMP_AGE_DAYS} days)"

    local total_size=0
    local count=0

    for dir in /tmp /var/tmp; do
        if [[ -d "$dir" ]]; then
            while IFS= read -r -d '' file; do
                local size
                size=$(stat -c%s "$file" 2>/dev/null || echo 0)
                total_size=$((total_size + size))
                ((count++)) || true
            done < <(find "$dir" -maxdepth 2 -type f -mtime +"$TMP_AGE_DAYS" -print0 2>/dev/null)
        fi
    done

    printf "  %-30s %s (%d files)\n" "Old temp files:" "$(human_bytes "$total_size")" "$count"
    add_reclaimable "$total_size"
}

clean_tmp() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would delete temp files older than ${TMP_AGE_DAYS} days"
        return
    fi

    for dir in /tmp /var/tmp; do
        if [[ -d "$dir" ]]; then
            find "$dir" -maxdepth 2 -type f -mtime +"$TMP_AGE_DAYS" -delete 2>/dev/null || true
        fi
    done
    log "Cleaned temp files"
}

# ══════════════════════════════════════════════════════════════════════
#                       PACKAGE CACHE
# ══════════════════════════════════════════════════════════════════════

scan_package_cache() {
    section_header "Package Cache"

    if command -v apt-get &>/dev/null; then
        local apt_size
        apt_size=$(du -sb /var/cache/apt/archives/ 2>/dev/null | awk '{print $1}' || echo "0")
        printf "  %-30s %s\n" "APT cache:" "$(human_bytes "$apt_size")"
        add_reclaimable "$apt_size"
    fi

    if command -v yum &>/dev/null || command -v dnf &>/dev/null; then
        local yum_size
        yum_size=$(du -sb /var/cache/yum/ /var/cache/dnf/ 2>/dev/null | awk '{total+=$1} END {print total+0}')
        printf "  %-30s %s\n" "YUM/DNF cache:" "$(human_bytes "$yum_size")"
        add_reclaimable "$yum_size"
    fi
}

clean_package_cache() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would clean package cache"
        return
    fi

    if command -v apt-get &>/dev/null; then
        apt-get clean -y 2>/dev/null || warn "apt-get clean failed"
        log "Cleaned APT cache"
    fi

    if command -v dnf &>/dev/null; then
        dnf clean all 2>/dev/null || warn "dnf clean failed"
        log "Cleaned DNF cache"
    elif command -v yum &>/dev/null; then
        yum clean all 2>/dev/null || warn "yum clean failed"
        log "Cleaned YUM cache"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       OLD KERNELS
# ══════════════════════════════════════════════════════════════════════

scan_old_kernels() {
    section_header "Old Kernels"

    local current_kernel
    current_kernel=$(uname -r)

    local old_count=0
    local total_size=0

    if command -v dpkg &>/dev/null; then
        while IFS= read -r pkg; do
            [[ -z "$pkg" ]] && continue
            local pkg_version
            pkg_version=$(echo "$pkg" | sed 's/linux-image-//' | sed 's/-generic//' | sed 's/-unsigned//')
            if [[ "$current_kernel" != *"$pkg_version"* ]]; then
                local size
                size=$(dpkg-query -W --showformat='${Installed-Size}' "$pkg" 2>/dev/null || echo "0")
                total_size=$((total_size + size * 1024))
                ((old_count++)) || true
                verbose "Old kernel: ${pkg} ($(human_bytes $((size * 1024))))"
            fi
        done < <(dpkg --list 'linux-image-*' 2>/dev/null | grep '^ii' | awk '{print $2}' | grep -v "$current_kernel")
    elif command -v rpm &>/dev/null; then
        while IFS= read -r pkg; do
            [[ -z "$pkg" ]] && continue
            if [[ "$pkg" != *"$current_kernel"* ]]; then
                local size
                size=$(rpm -q --queryformat '%{SIZE}' "$pkg" 2>/dev/null || echo "0")
                total_size=$((total_size + size))
                ((old_count++)) || true
                verbose "Old kernel: ${pkg}"
            fi
        done < <(rpm -qa kernel 2>/dev/null)
    fi

    printf "  %-30s %s\n" "Current kernel:" "$current_kernel"
    printf "  %-30s %d ($(human_bytes "$total_size"))\n" "Old kernels:" "$old_count"
    add_reclaimable "$total_size"
}

# ══════════════════════════════════════════════════════════════════════
#                       DOCKER CLEANUP
# ══════════════════════════════════════════════════════════════════════

scan_docker() {
    if ! command -v docker &>/dev/null; then
        return
    fi

    if ! docker info &>/dev/null 2>&1; then
        return
    fi

    section_header "Docker"

    # Dangling images
    local dangling_count
    dangling_count=$(docker images -f "dangling=true" -q 2>/dev/null | wc -l)

    printf "  %-30s %d\n" "Dangling images:" "$dangling_count"

    # Stopped containers
    local stopped_count
    stopped_count=$(docker ps -f "status=exited" -q 2>/dev/null | wc -l)
    printf "  %-30s %d\n" "Stopped containers:" "$stopped_count"

    # Unused volumes
    local vol_count
    vol_count=$(docker volume ls -f "dangling=true" -q 2>/dev/null | wc -l)
    printf "  %-30s %d\n" "Unused volumes:" "$vol_count"

    # Build cache
    if docker builder prune --dry-run 2>/dev/null | grep -q "Total:"; then
        local build_cache
        build_cache=$(docker builder prune --dry-run 2>/dev/null | grep "Total:" | awk '{print $2}')
        printf "  %-30s %s\n" "Build cache:" "${build_cache:-0}"
    fi

    # Docker system df
    echo ""
    docker system df 2>/dev/null | while IFS= read -r line; do
        printf "    %s\n" "$line"
    done
}

clean_docker() {
    if ! command -v docker &>/dev/null || ! docker info &>/dev/null 2>&1; then
        return
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        log "[DRY-RUN] Would prune Docker system (stopped containers, dangling images, unused networks, build cache)"
        return
    fi

    docker system prune -f 2>/dev/null || warn "Docker prune failed"
    log "Pruned Docker system"
}

# ══════════════════════════════════════════════════════════════════════
#                       LARGE FILES
# ══════════════════════════════════════════════════════════════════════

scan_large_files() {
    section_header "Large Files (> ${LARGE_FILE_MIN})"

    printf "  ${BOLD}%-12s  %s${RESET}\n" "SIZE" "FILE"
    printf "  %s\n" "$(printf '%.0s─' {1..70})"

    local count=0
    for dir in $LARGE_FILE_DIRS; do
        [[ -d "$dir" ]] || continue
        find "$dir" -xdev -type f -size +"$LARGE_FILE_MIN" -printf '%s %p\n' 2>/dev/null
    done | sort -rn | head -20 | while IFS=' ' read -r size path; do
        printf "  %10s  %s\n" "$(human_bytes "$size")" "$path"
        ((count++)) || true
    done

    if [[ "$count" -eq 0 ]]; then
        echo "  No files larger than ${LARGE_FILE_MIN} found"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       SUMMARY
# ══════════════════════════════════════════════════════════════════════

print_summary() {
    echo ""
    echo -e "  ${BOLD}══════════════════════════════════════════${RESET}"
    echo -e "  ${BOLD}Disk Cleanup Summary${RESET}"
    echo -e "  ${BOLD}══════════════════════════════════════════${RESET}"

    # Current disk usage
    local root_pct
    root_pct=$(df / 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%')
    printf "  %-20s %s%%\n" "Root disk usage:" "${root_pct:-?}"
    printf "  %-20s %s\n" "Reclaimable:" "$(human_bytes "$TOTAL_RECLAIMABLE")"

    if [[ "$TOTAL_CLEANED" -gt 0 ]]; then
        printf "  %-20s %s\n" "Cleaned:" "$(human_bytes "$TOTAL_CLEANED")"
    fi

    if [[ "$DRY_RUN" == "true" && "$RUN_MODE" == *"clean"* ]]; then
        echo ""
        echo -e "  ${YELLOW}Dry-run mode — nothing was deleted${RESET}"
        echo -e "  Run with --force to actually clean"
    fi

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Find and clean disk space hogs on Linux servers

USAGE:
  ${SCRIPT_NAME} [OPTIONS]

MODES:
  --scan                Scan and report reclaimable space (default)
  --clean               Scan and clean (dry-run by default)
  --large-files         Show largest files only

OPTIONS:
  --force               Actually delete files (disables dry-run)
  --log-age DAYS        Age threshold for old logs (default: ${LOG_AGE_DAYS})
  --tmp-age DAYS        Age threshold for temp files (default: ${TMP_AGE_DAYS})
  --journal-max SIZE    Journal vacuum target (default: ${JOURNAL_MAX})
  --min-size SIZE       Minimum size for large file scan (default: ${LARGE_FILE_MIN})
  --verbose             Show individual files found
  --no-color            Disable colored output
  --help                Show this help

ENVIRONMENT VARIABLES:
  LOG_AGE_DAYS          Age threshold for old logs (default: 30)
  TMP_AGE_DAYS          Age threshold for temp files (default: 7)
  JOURNAL_MAX           Journal vacuum target (default: 500M)
  LARGE_FILE_MIN        Large file threshold (default: 100M)
  DRY_RUN               Dry-run mode (default: true)

EXAMPLES:
  # Scan for reclaimable space
  ./disk-cleanup.sh --scan

  # Scan with file details
  ./disk-cleanup.sh --scan --verbose

  # Clean (dry-run, shows what would be deleted)
  ./disk-cleanup.sh --clean

  # Actually clean
  sudo ./disk-cleanup.sh --clean --force

  # Just show large files
  ./disk-cleanup.sh --large-files

  # Custom thresholds
  ./disk-cleanup.sh --clean --force --log-age 14 --journal-max 200M
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    RUN_MODE="scan"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --scan)
                RUN_MODE="scan"; shift ;;
            --clean)
                RUN_MODE="clean"; shift ;;
            --large-files)
                RUN_MODE="large-files"; shift ;;
            --force)
                DRY_RUN="false"; shift ;;
            --log-age)
                LOG_AGE_DAYS="$2"; shift 2 ;;
            --tmp-age)
                TMP_AGE_DAYS="$2"; shift 2 ;;
            --journal-max)
                JOURNAL_MAX="$2"; shift 2 ;;
            --min-size)
                LARGE_FILE_MIN="$2"; shift 2 ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                setup_colors
                usage
                exit 0 ;;
            *)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 1 ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    setup_colors

    echo ""
    echo -e "${BOLD}Disk Cleanup — $(hostname -f 2>/dev/null || hostname)${RESET}"
    echo -e "Mode:   ${RUN_MODE}"
    if [[ "$RUN_MODE" == "clean" ]]; then
        if [[ "$DRY_RUN" == "true" ]]; then
            echo -e "Safety: ${YELLOW}dry-run (use --force to delete)${RESET}"
        else
            echo -e "Safety: ${RED}LIVE — files will be deleted${RESET}"
        fi
    fi
    echo -e "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"

    case "$RUN_MODE" in
        scan)
            scan_old_logs
            scan_journal
            scan_tmp
            scan_package_cache
            scan_old_kernels
            scan_docker
            scan_large_files
            print_summary
            ;;
        clean)
            scan_old_logs
            clean_old_logs
            scan_journal
            clean_journal
            scan_tmp
            clean_tmp
            scan_package_cache
            clean_package_cache
            scan_docker
            clean_docker
            scan_large_files
            print_summary
            ;;
        large-files)
            scan_large_files
            ;;
    esac
}

main "$@"