#!/usr/bin/env bash

# Self-cleaning script for fetch_ml Docker resources
# Usage: ./scripts/cleanup.sh [--dry-run] [--force] [--all]

set -euo pipefail

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

# Default options
DRY_RUN=false
FORCE=false
ALL=false

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        --force)
            FORCE=true
            shift
            ;;
        --all)
            ALL=true
            shift
            ;;
        --help|-h)
            echo "Usage: $0 [--dry-run] [--force] [--all]"
            echo ""
            echo "Options:"
            echo "  --dry-run  Show what would be cleaned without actually doing it"
            echo "  --force    Skip confirmation prompts"
            echo "  --all      Clean everything including images"
            echo ""
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# Helper functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Docker check
if ! command -v docker &> /dev/null; then
    log_error "Docker is not installed or not in PATH"
    exit 1
fi

# Check if Docker is running
if ! docker info &> /dev/null; then
    log_error "Docker daemon is not running"
    exit 1
fi

log_info "Starting cleanup of fetch_ml Docker resources..."

# Function to count resources
count_resources() {
    local containers=$(docker ps -a --filter "name=ml-" --format "{{.Names}}" | wc -l)
    local images=$(docker images --filter "reference=fetch_ml-*" --format "{{.Repository}}" | wc -l)
    local networks=$(docker network ls --filter "name=ml-" --format "{{.Name}}" | wc -l)
    local volumes=$(docker volume ls --filter "name=ml-" --format "{{.Name}}" | wc -l)
    
    echo "Containers: $containers, Images: $images, Networks: $networks, Volumes: $volumes"
}

# Show current state
log_info "Current resources: $(count_resources)"

# Dry run mode
if [ "$DRY_RUN" = true ]; then
    log_info "DRY RUN MODE - No actual cleanup will be performed"
    
    echo ""
    log_info "Containers that would be stopped and removed:"
    docker ps -a --filter "name=ml-" --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}" || echo "No containers found"
    
    echo ""
    log_info "Images that would be removed:"
    if [ "$ALL" = true ]; then
        docker images --filter "reference=fetch_ml-*" --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}" || echo "No images found"
    else
        echo "Use --all to remove images"
    fi
    
    echo ""
    log_info "Networks that would be removed:"
    docker network ls --filter "name=ml-" --format "table {{.Name}}\t{{.Driver}}" || echo "No networks found"
    
    echo ""
    log_info "Volumes that would be removed:"
    docker volume ls --filter "name=ml-" --format "table {{.Name}}\t{{.Driver}}" || echo "No volumes found"
    
    exit 0
fi

# Confirmation prompt
if [ "$FORCE" = false ]; then
    echo ""
    read -p "This will stop and remove all fetch_ml containers. Continue? (y/N): " -n 1 -r
    echo ""
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "Cleanup cancelled"
        exit 0
    fi
fi

# Stop containers
log_info "Stopping containers..."
containers=$(docker ps -q --filter "name=ml-")
if [ -n "$containers" ]; then
    if [ "$DRY_RUN" = false ]; then
        echo "$containers" | xargs docker stop
        log_success "Containers stopped"
    fi
else
    log_info "No running containers found"
fi

 # Remove containers
log_info "Removing containers..."
containers=$(docker ps -aq --filter "name=ml-")
if [ -n "$containers" ]; then
    if [ "$DRY_RUN" = false ]; then
        echo "$containers" | xargs docker rm
        log_success "Containers removed"
    fi
else
    log_info "No containers found"
fi

# Remove networks
log_info "Removing networks..."
networks=$(docker network ls -q --filter "name=ml-")
if [ -n "$networks" ]; then
    if [ "$DRY_RUN" = false ]; then
        echo "$networks" | xargs docker network rm
        log_success "Networks removed"
    fi
else
    log_info "No networks found"
fi

 # Remove volumes (with caution)
log_warning "Skipping volumes by default (use --all to remove them)"
if [ "$ALL" = true ]; then
    volumes=$(docker volume ls -q --filter "name=ml-")
    if [ -n "$volumes" ]; then
        if [ "$DRY_RUN" = false ]; then
            echo "$volumes" | xargs docker volume rm
            log_success "Volumes removed"
        fi
    else
        log_info "No volumes found"
    fi
else
    log_info "Skipping volumes"
fi

 # Remove images if requested
if [ "$ALL" = true ]; then
    log_info "Removing images..."
    images=$(docker images -q --filter "reference=fetch_ml-*")
    if [ -n "$images" ]; then
        if [ "$DRY_RUN" = false ]; then
            echo "$images" | xargs docker rmi
            log_success "Images removed"
        fi
    else
        log_info "No images found"
    fi
else
    log_info "Skipping images (use --all to remove them)"
fi

 # General Docker cleanup
if [ "$ALL" = true ]; then
    log_info "Running general Docker cleanup (docker system prune)..."
    if [ "$DRY_RUN" = false ]; then
        docker system prune -f
        log_success "General cleanup completed"
    fi
else
    log_info "Skipping docker system prune (use --all to enable)"
fi

# Show final state
log_info "Final resources: $(count_resources)"

# Space reclaimed
if [ "$DRY_RUN" = false ]; then
    log_info "Docker system info:"
    docker system df
fi

log_success "Cleanup completed!"