#!/bin/bash
# Omarchy Repository Cleaner
# Removes old package versions keeping only the latest

set -e

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

# Repository configuration
REPO_DIR="/home/ryan/Dev/omarchy-repo"
ARCH_DIR="$REPO_DIR/x86_64"
KEEP_VERSIONS=2

# Function to print colored messages
print_msg() {
    local color=$1
    local msg=$2
    echo -e "${color}${msg}${NC}"
}

# Function to clean old packages
clean_packages() {
    local keep=$1
    local dry_run=$2
    
    print_msg "$BLUE" "==> Cleaning old packages (keeping $keep versions)..."
    
    cd "$ARCH_DIR"
    
    # Group packages by name
    declare -A packages
    
    for pkg in *.pkg.tar.*; do
        [[ -f "$pkg" ]] || continue
        
        # Skip signature files
        [[ "$pkg" == *.sig ]] && continue
        
        # Extract package name (remove version and architecture)
        # Format: name-version-release-arch.pkg.tar.*
        local pkgname=$(echo "$pkg" | sed -E 's/-[0-9]+.*-(any|x86_64|i686)\.pkg\.tar\..*//')
        
        # Add to array
        if [[ -n "${packages[$pkgname]}" ]]; then
            packages[$pkgname]="${packages[$pkgname]} $pkg"
        else
            packages[$pkgname]="$pkg"
        fi
    done
    
    # Process each package group
    local total_removed=0
    local space_freed=0
    
    for pkgname in "${!packages[@]}"; do
        # Convert string to array and sort by modification time (newest first)
        IFS=' ' read -ra pkg_files <<< "${packages[$pkgname]}"
        
        # Sort by modification time
        local sorted_files=($(ls -t "${pkg_files[@]}" 2>/dev/null))
        local count=${#sorted_files[@]}
        
        if [[ $count -gt $keep ]]; then
            print_msg "$GREEN" "  -> Processing $pkgname ($count versions found)"
            
            # Remove old versions
            for ((i=$keep; i<$count; i++)); do
                local file="${sorted_files[$i]}"
                local size=$(stat -c%s "$file" 2>/dev/null || echo 0)
                space_freed=$((space_freed + size))
                
                if [[ "$dry_run" == true ]]; then
                    print_msg "$YELLOW" "      Would remove: $file ($(numfmt --to=iec-i --suffix=B $size))"
                else
                    print_msg "$RED" "      Removing: $file ($(numfmt --to=iec-i --suffix=B $size))"
                    rm -f "$file"
                    rm -f "${file}.sig" 2>/dev/null || true
                fi
                
                total_removed=$((total_removed + 1))
            done
        fi
    done
    
    # Summary
    if [[ $total_removed -gt 0 ]]; then
        if [[ "$dry_run" == true ]]; then
            print_msg "$BLUE" "\n  -> Would remove $total_removed old package(s)"
            print_msg "$BLUE" "  -> Would free $(numfmt --to=iec-i --suffix=B $space_freed)"
        else
            print_msg "$GREEN" "\n  -> Removed $total_removed old package(s)"
            print_msg "$GREEN" "  -> Freed $(numfmt --to=iec-i --suffix=B $space_freed)"
        fi
    else
        print_msg "$YELLOW" "  -> No old packages to remove"
    fi
}

# Function to remove all packages
remove_all_packages() {
    local dry_run=$1
    
    print_msg "$RED" "==> WARNING: Removing ALL packages from repository!"
    
    if [[ "$dry_run" != true ]]; then
        read -p "Are you sure? Type 'yes' to confirm: " confirm
        if [[ "$confirm" != "yes" ]]; then
            print_msg "$YELLOW" "Cancelled."
            exit 0
        fi
    fi
    
    cd "$ARCH_DIR"
    
    local count=$(ls -1 *.pkg.tar.* 2>/dev/null | wc -l)
    
    if [[ $count -eq 0 ]]; then
        print_msg "$YELLOW" "No packages found to remove"
        return
    fi
    
    if [[ "$dry_run" == true ]]; then
        print_msg "$YELLOW" "Would remove $count package(s):"
        ls -1 *.pkg.tar.* 2>/dev/null | while read -r pkg; do
            echo "  - $pkg"
        done
    else
        rm -f *.pkg.tar.*
        rm -f *.db* *.files*
        print_msg "$GREEN" "Removed $count package(s) and database files"
    fi
}

# Function to show repository disk usage
show_disk_usage() {
    print_msg "$BLUE" "\n==> Repository Disk Usage:"
    
    cd "$ARCH_DIR"
    
    local total_size=$(du -sh . 2>/dev/null | cut -f1)
    local package_count=$(ls -1 *.pkg.tar.* 2>/dev/null | grep -v '\.sig$' | wc -l)
    
    print_msg "$GREEN" "  -> Total size: $total_size"
    print_msg "$GREEN" "  -> Package count: $package_count"
    
    if [[ $package_count -gt 0 ]]; then
        print_msg "$GREEN" "  -> Average package size: $(du -b *.pkg.tar.* 2>/dev/null | awk '{sum+=$1; count++} END {print sum/count}' | numfmt --to=iec-i --suffix=B)"
    fi
}

# Main execution
main() {
    print_msg "$BLUE" "================================"
    print_msg "$BLUE" "Omarchy Repository Cleaner"
    print_msg "$BLUE" "================================"
    
    # Parse command line arguments
    DRY_RUN=false
    REMOVE_ALL=false
    SHOW_USAGE=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            --keep)
                KEEP_VERSIONS="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --all)
                REMOVE_ALL=true
                shift
                ;;
            --usage)
                SHOW_USAGE=true
                shift
                ;;
            -h|--help)
                echo "Usage: $0 [OPTIONS]"
                echo ""
                echo "Options:"
                echo "  --keep N      Keep N versions of each package (default: 2)"
                echo "  --dry-run     Preview what would be removed without removing"
                echo "  --all         Remove ALL packages (dangerous!)"
                echo "  --usage       Show disk usage statistics"
                echo "  -h, --help    Show this help message"
                echo ""
                echo "Examples:"
                echo "  $0                    # Keep 2 versions (default)"
                echo "  $0 --keep 3          # Keep 3 versions"
                echo "  $0 --dry-run         # Preview cleanup"
                echo "  $0 --all --dry-run   # Preview complete removal"
                exit 0
                ;;
            *)
                print_msg "$RED" "Unknown option: $1"
                exit 1
                ;;
        esac
    done
    
    # Validate keep versions
    if ! [[ "$KEEP_VERSIONS" =~ ^[0-9]+$ ]] || [[ "$KEEP_VERSIONS" -lt 1 ]]; then
        print_msg "$RED" "Error: --keep must be a positive integer"
        exit 1
    fi
    
    # Show dry run notice
    if [[ "$DRY_RUN" == true ]]; then
        print_msg "$YELLOW" "DRY RUN MODE - No changes will be made"
    fi
    
    # Execute based on options
    if [[ "$SHOW_USAGE" == true ]]; then
        show_disk_usage
    elif [[ "$REMOVE_ALL" == true ]]; then
        remove_all_packages "$DRY_RUN"
    else
        clean_packages "$KEEP_VERSIONS" "$DRY_RUN"
        
        # Update database after cleaning (unless dry run)
        if [[ "$DRY_RUN" != true ]]; then
            print_msg "$BLUE" "\n==> Updating repository database..."
            "$REPO_DIR/scripts/update-repo.sh" >/dev/null 2>&1 && {
                print_msg "$GREEN" "  -> Database updated successfully!"
            } || {
                print_msg "$YELLOW" "  -> Database update failed (may need manual update)"
            }
        fi
        
        show_disk_usage
    fi
    
    print_msg "$GREEN" "\n==> Cleanup complete!"
}

# Run main function
main "$@"