#!/usr/bin/env bash
# Remove directories below a specified size (in KB)
# Usage: remove-small-dirs DIRECTORY [SIZE_THRESHOLD]
# Default SIZE_THRESHOLD is 1000 KB

set -euo pipefail

usage() {
    cat <<-EOF
	Usage: remove-small-dirs [OPTIONS] DIRECTORY [SIZE_THRESHOLD]
	
	Remove directories below a specified size (default: 1000 KB).
	
	OPTIONS:
	  -n, --dry-run    Show what would be removed without deleting
	  -h, --help       Display this help message
	
	ARGUMENTS:
	  DIRECTORY        Directory to search for small directories
	  SIZE_THRESHOLD   Maximum size in KB (default: 1000)
	EOF
}

# Parse options
DRY_RUN=false
while [[ $# -gt 0 ]]; do
    case "$1" in
        -n|--dry-run)
            DRY_RUN=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        -*)
            echo "[ERROR] Unknown option: $1" >&2
            usage
            exit 1
            ;;
        *)
            break
            ;;
    esac
done

if [[ $# -lt 1 ]]; then
    echo "[ERROR] You must provide a directory." >&2
    usage
    exit 1
fi

dir="$1"
SIZE="${2:-1000}"

if [[ ! -d "$dir" ]]; then
    echo "[ERROR] Directory does not exist: $dir" >&2
    exit 1
fi

if ! [[ "$SIZE" =~ ^[0-9]+$ ]] || [[ "$SIZE" -lt 1 ]]; then
    echo "[ERROR] SIZE_THRESHOLD must be a positive integer" >&2
    exit 1
fi

echo "[INFO] Searching for directories <= $SIZE KB in '$dir'..."

# Find directories with size less or equal to SIZE
# Sort by depth (deepest first) to avoid removing parent before child
small_dirs=$(find "$dir" -mindepth 1 -type d -exec du -ks {} + | \
    awk -v size="$SIZE" '$1 <= size {print $2}' | \
    awk '{ print length, $0 }' | sort -rn | cut -d' ' -f2-)

if [[ -z "$small_dirs" ]]; then
    echo "[INFO] No directories with size <= $SIZE KB found in '$dir'."
    exit 0
fi

echo "[INFO] Found $(echo "$small_dirs" | wc -l) directories to remove:"
echo "$small_dirs"

if [[ "$DRY_RUN" == true ]]; then
    echo "[DRY-RUN] Would remove the above directories."
    exit 0
fi

read -r -p "Remove these directories? [y/N] " response
response="${response,,}"  # Convert to lowercase

if [[ ! "$response" =~ ^(yes|y)$ ]]; then
    echo "[INFO] Exiting, no changes were made."
    exit 0
fi

count=0
while IFS= read -r small_dir; do
    if [[ -d "$small_dir" ]]; then
        echo "[INFO] Removing: $small_dir"
        if rm -rf "$small_dir"; then
            ((count++))
        else
            echo "[ERROR] Failed to remove: $small_dir" >&2
        fi
    fi
done <<< "$small_dirs"

echo "[SUCCESS] Removed $count directories."