#!/bin/bash
# Docker Cleanup Script
# Enhanced version matching Go docwell functionality

set -euo pipefail

VERSION="2.6.2"

# Resolve script directory and source shared library
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib/common.sh
source "$SCRIPT_DIR/lib/common.sh"

# Enable debug trace after sourcing common.sh (so DEBUG variable is available)
enable_debug_trace

# CLI flags
CLEANUP_CONTAINERS=false
CLEANUP_IMAGES=false
CLEANUP_VOLUMES=false
CLEANUP_NETWORKS=false
CLEANUP_BUILDCACHE=false
CLEANUP_SYSTEM=false
CLEANUP_PACKAGES=false
CLEANUP_ALL=false
LOG_FILE="${LOG_FILE:-/tmp/docwell/docker-cleanup.log}"

# Load config file overrides
load_config

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --version)
                echo "docker-cleanup.sh v$VERSION"
                exit 0
                ;;
            --quiet|-q)
                QUIET=true
                shift
                ;;
            --yes|-y)
                YES=true
                shift
                ;;
            --containers)
                CLEANUP_CONTAINERS=true
                shift
                ;;
            --images)
                CLEANUP_IMAGES=true
                shift
                ;;
            --volumes)
                CLEANUP_VOLUMES=true
                shift
                ;;
            --networks)
                CLEANUP_NETWORKS=true
                shift
                ;;
            --buildcache)
                CLEANUP_BUILDCACHE=true
                shift
                ;;
            --system)
                CLEANUP_SYSTEM=true
                shift
                ;;
            --packages)
                CLEANUP_PACKAGES=true
                shift
                ;;
            --all)
                CLEANUP_ALL=true
                shift
                ;;
            --install-deps)
                AUTO_INSTALL=true
                shift
                ;;
            --log)
                LOG_FILE="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --debug|-d)
                DEBUG=true
                shift
                ;;
            --verbose|-v)
                VERBOSE=true
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                echo -e "${RED}Unknown option: $1${NC}" >&2
                show_help
                exit 1
                ;;
        esac
    done
}

show_help() {
    cat << EOF
Docker Cleanup Script v$VERSION

Usage: $0 [OPTIONS]

Options:
    --containers          Remove stopped containers
    --images             Remove unused images
    --volumes            Remove unused volumes (DATA LOSS WARNING)
    --networks           Remove unused networks
    --buildcache         Clear build cache
    --system             System-wide cleanup
    --packages           Clean OS package cache
    --all                Run all cleanup operations
    --install-deps       Auto-install missing dependencies (requires root)
    --log FILE           Log file path (default: /tmp/docwell/docker-cleanup.log)
    --dry-run            Show what would be done without actually doing it
    --debug, -d          Enable debug mode (verbose output + xtrace)
    --verbose, -v        Enable verbose logging
    --quiet, -q          Suppress non-error output
    --yes, -y            Auto-confirm all prompts
    --version            Show version information
    --help, -h           Show this help message

Examples:
    $0 --containers              # Remove stopped containers
    $0 --all --yes               # Run all cleanups non-interactively
    $0 --images --volumes        # Remove images and volumes
    $0 --all --dry-run           # Preview cleanup without executing
EOF
}

show_docker_usage() {
    log_header "Docker Disk Usage"
    docker system df -v 2>/dev/null || {
        log_error "Docker not running or not accessible"
        exit 1
    }
    echo
}

cleanup_containers() {
    log_header "Container Cleanup"

    local stopped_count
    stopped_count=$(docker ps -aq --filter status=exited 2>/dev/null | wc -l || echo "0")

    if [[ "$stopped_count" == "0" ]] || [[ -z "$stopped_count" ]]; then
        log_info "No stopped containers found"
        echo
        return 0
    fi

    if [[ "$QUIET" == "false" ]]; then
        echo -e "${YELLOW}Stopped containers to remove ($stopped_count):${NC}"
        docker ps -a --filter status=exited --format "table {{.Names}}\t{{.Image}}\t{{.Status}}" 2>/dev/null || true
        echo
    fi

    if ! confirm "Remove stopped containers?"; then
        log_info "Skipping container cleanup"
        echo
        return 0
    fi

    if run_or_dry "prune stopped containers" docker container prune -f &>/dev/null; then
        log_success "Removed stopped containers"
    else
        log_error "Failed to remove containers"
        return 1
    fi
    echo
}

cleanup_images() {
    local remove_all="${1:-}"
    log_header "Image Cleanup"

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

    if [[ "$dangling_count" != "0" ]] && [[ -n "$dangling_count" ]]; then
        if [[ "$QUIET" == "false" ]]; then
            echo -e "${YELLOW}Dangling images to remove ($dangling_count):${NC}"
            docker images -f "dangling=true" --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}" 2>/dev/null || true
            echo
        fi

        if ! confirm "Remove dangling images?"; then
            log_info "Skipping dangling image cleanup"
            echo
            return 0
        fi

        if run_or_dry "prune dangling images" docker image prune -f &>/dev/null; then
            log_success "Removed dangling images"
        else
            log_error "Failed to remove dangling images"
            return 1
        fi
    else
        log_info "No dangling images found"
    fi

    if [[ "$QUIET" == "false" ]] && [[ "$remove_all" != "--all" ]]; then
        local unused_count
        unused_count=$(docker images -q 2>/dev/null | wc -l || echo "0")
        if [[ "$unused_count" != "0" ]]; then
            echo -e "${YELLOW}Unused images (showing first 20):${NC}"
            docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}" 2>/dev/null | head -20 || true
            echo
        fi
    fi

    if [[ "$remove_all" != "--all" ]]; then
        if ! confirm "Remove ALL unused images?"; then
            log_info "Skipping unused image cleanup"
            echo
            return 0
        fi
    fi

    if run_or_dry "prune all unused images" docker image prune -a -f &>/dev/null; then
        log_success "Removed all unused images"
    else
        log_error "Failed to remove unused images"
        return 1
    fi
    echo
}

cleanup_volumes() {
    log_header "Volume Cleanup"

    local unused_count
    unused_count=$(docker volume ls -q --filter dangling=true 2>/dev/null | wc -l || echo "0")

    if [[ "$unused_count" == "0" ]] || [[ -z "$unused_count" ]]; then
        log_info "No unused volumes found"
        echo
        return 0
    fi

    if [[ "$QUIET" == "false" ]]; then
        echo -e "${YELLOW}Unused volumes to remove ($unused_count):${NC}"
        docker volume ls --filter dangling=true --format "table {{.Name}}\t{{.Driver}}" 2>/dev/null || true
        echo
    fi

    log_warn "This will permanently delete volume data!"
    if ! confirm "Remove unused volumes?"; then
        log_info "Skipping volume cleanup"
        echo
        return 0
    fi

    if run_or_dry "prune unused volumes" docker volume prune -f &>/dev/null; then
        log_success "Removed unused volumes"
    else
        log_error "Failed to remove volumes"
        return 1
    fi
    echo
}

cleanup_networks() {
    log_header "Network Cleanup"

    if [[ "$QUIET" == "false" ]]; then
        echo -e "${YELLOW}Unused networks:${NC}"
        docker network ls --format "table {{.Name}}\t{{.Driver}}\t{{.Scope}}"
        echo
    fi

    if ! confirm "Remove unused networks?"; then
        log_info "Skipping network cleanup"
        echo
        return 0
    fi

    if run_or_dry "prune unused networks" docker network prune -f &>/dev/null; then
        log_success "Removed unused networks"
    else
        log_error "Failed to remove networks"
        return 1
    fi
    echo
}

cleanup_build_cache() {
    log_header "Build Cache Cleanup"

    if [[ "$QUIET" == "false" ]]; then
        echo -e "${YELLOW}Build cache usage:${NC}"
        docker system df 2>/dev/null | grep "Build Cache" || echo "No build cache info available"
        echo
    fi

    if ! confirm "Clear build cache?"; then
        log_info "Skipping build cache cleanup"
        echo
        return 0
    fi

    if run_or_dry "prune build cache" docker builder prune -f &>/dev/null; then
        log_success "Cleared build cache"
    else
        log_error "Failed to clear build cache"
        return 1
    fi
    echo
}

cleanup_system() {
    log_header "System-wide Cleanup"

    log_warn "This will remove:"
    echo "- All stopped containers"
    echo "- All networks not used by at least one container"
    echo "- All dangling images"
    echo "- All dangling build cache"
    echo

    if ! confirm "Perform system cleanup?"; then
        log_info "Skipping system cleanup"
        echo
        return 0
    fi

    if ! run_or_dry "system prune" docker system prune -f &>/dev/null; then
        log_error "Failed to prune system"
        return 1
    fi

    log_success "System cleanup completed"

    if confirm "Also remove unused images?"; then
        if run_or_dry "aggressive system prune" docker system prune -a -f &>/dev/null; then
            log_success "Aggressive cleanup completed"
        else
            log_error "Failed to prune system with images"
            return 1
        fi
    fi
    echo
}

cleanup_packages() {
    log_header "Package Manager Cleanup"

    if ! check_root; then
        log_error "sudo privileges required for package cleanup"
        echo
        return 1
    fi

    if command -v pacman &> /dev/null; then
        log_info "Detected Arch Linux"

        local orphaned
        orphaned=$(pacman -Qtdq 2>/dev/null || true)
        if [[ -n "$orphaned" ]]; then
            if [[ "$QUIET" == "false" ]]; then
                echo -e "${YELLOW}Orphaned packages:${NC}"
                echo "$orphaned"
                echo
            fi

            if confirm "Remove orphaned packages?"; then
                local orphaned_packages=()
                while IFS= read -r pkg; do
                    [[ -n "$pkg" ]] && orphaned_packages+=("$pkg")
                done < <(echo "$orphaned")

                if (( ${#orphaned_packages[@]} > 0 )); then
                    if run_or_dry "remove orphaned packages" sudo pacman -Rns "${orphaned_packages[@]}" &>/dev/null; then
                        log_success "Orphaned packages removed"
                    else
                        log_warn "Failed to remove orphaned packages"
                    fi
                fi
            fi
        else
            log_info "No orphaned packages found"
        fi

        if confirm "Clear package cache?"; then
            if run_or_dry "clear package cache" sudo pacman -Scc --noconfirm &>/dev/null; then
                log_success "Package cache cleared"
            else
                log_error "Failed to clear package cache"
                return 1
            fi
        fi

    elif command -v apt &> /dev/null; then
        log_info "Detected Debian/Ubuntu"

        if confirm "Remove unused packages?"; then
            if run_or_dry "autoremove packages" sudo apt autoremove -y &>/dev/null; then
                log_success "Removed unused packages"
            else
                log_error "Failed to remove unused packages"
                return 1
            fi
        fi

        if confirm "Clean package cache?"; then
            if run_or_dry "clean package cache" sudo apt autoclean &>/dev/null; then
                log_success "Cache cleaned"
            else
                log_error "Failed to clean cache"
                return 1
            fi
        fi
    else
        log_info "No supported package manager found"
    fi
    echo
}

cleanup_all() {
    show_docker_usage
    cleanup_containers
    cleanup_images
    cleanup_networks
    cleanup_build_cache

    if confirm "Also remove unused volumes? (DATA LOSS WARNING)"; then
        cleanup_volumes
    fi

    if confirm "Also remove ALL unused images?"; then
        cleanup_images --all
    fi

    if confirm "Also clean package cache?"; then
        cleanup_packages
    fi
}

main_menu() {
    while true; do
        log_header "Docker Cleanup Menu"
        echo -e "${GRAY}DocWell Cleanup v$VERSION${NC}"
        echo
        echo "1) Show Docker disk usage"
        echo "2) Clean containers (stopped)"
        echo "3) Clean images (dangling/unused)"
        echo "4) Clean volumes (unused)"
        echo "5) Clean networks (unused)"
        echo "6) Clean build cache"
        echo "7) System-wide cleanup"
        echo "8) Package manager cleanup"
        echo "9) Run all cleanups"
        echo "0) Exit"
        echo

        read -rp "Select option [0-9]: " choice

        if [[ ! "$choice" =~ ^[0-9]$ ]]; then
            log_error "Invalid option. Please enter a number between 0-9."
            echo
            read -rp "Press Enter to continue..."
            clear_screen
            continue
        fi

        case $choice in
            1) show_docker_usage ;;
            2) cleanup_containers ;;
            3) cleanup_images ;;
            4) cleanup_volumes ;;
            5) cleanup_networks ;;
            6) cleanup_build_cache ;;
            7) cleanup_system ;;
            8) cleanup_packages ;;
            9) cleanup_all ;;
            0)
                log_info "Exiting..."
                exit 0
                ;;
            *)
                log_error "Invalid option. Please try again."
                ;;
        esac

        echo
        read -rp "Press Enter to continue..."
        clear_screen
    done
}

main() {
    parse_args "$@"
    check_docker || exit 1

    # Handle CLI flags
    if [[ "$CLEANUP_ALL" == "true" ]]; then
        cleanup_all
        exit 0
    fi

    [[ "$CLEANUP_CONTAINERS" == "true" ]] && cleanup_containers
    [[ "$CLEANUP_IMAGES" == "true" ]] && cleanup_images
    [[ "$CLEANUP_VOLUMES" == "true" ]] && cleanup_volumes
    [[ "$CLEANUP_NETWORKS" == "true" ]] && cleanup_networks
    [[ "$CLEANUP_BUILDCACHE" == "true" ]] && cleanup_build_cache
    [[ "$CLEANUP_SYSTEM" == "true" ]] && cleanup_system
    [[ "$CLEANUP_PACKAGES" == "true" ]] && cleanup_packages

    # If no flags provided, run interactive menu
    if [[ "$CLEANUP_CONTAINERS" == "false" ]] && \
       [[ "$CLEANUP_IMAGES" == "false" ]] && \
       [[ "$CLEANUP_VOLUMES" == "false" ]] && \
       [[ "$CLEANUP_NETWORKS" == "false" ]] && \
       [[ "$CLEANUP_BUILDCACHE" == "false" ]] && \
       [[ "$CLEANUP_SYSTEM" == "false" ]] && \
       [[ "$CLEANUP_PACKAGES" == "false" ]]; then
        clear_screen
        main_menu
    fi
}

main "$@"