#!/bin/bash

# Sharenet Cleanup Script
# This script cleans up Docker resources and registry images

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

# Configuration
CLEANUP_TYPE="${CLEANUP_TYPE:-ci-cd}"  # ci-cd or production
REGISTRY_DIR="${REGISTRY_DIR:-/opt/registry}"
DRY_RUN="${DRY_RUN:-false}"

# Critical infrastructure protection
CRITICAL_CONTAINERS="harbor-core,harbor-db,harbor-jobservice,harbor-log,harbor-portal,nginx,redis,registry,registryctl,trivy-adapter"
CRITICAL_IMAGES="goharbor,forgejo-runner"
CRITICAL_VOLUMES="harbor"
CRITICAL_NETWORKS="harbor"

# 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"
}

get_unused_images() {
    # Get unused images (excluding Harbor images) with detailed information
    # This includes both dangling images and tagged images not used by running containers
    
    # Get all images (excluding Harbor images)
    local all_images=$(docker image ls --format "{{.Repository}}:{{.Tag}} ({{.ID}}) - Size: {{.Size}}, Created: {{.CreatedAt}}" | grep -v "goharbor" | grep -v "harbor" || true)
    
    if [ -n "$all_images" ]; then
        echo "$all_images" | while read -r image; do
            if [ -n "$image" ]; then
                # Extract image ID from the format
                local image_id=$(echo "$image" | sed 's/.*(\([a-f0-9]*\)).*/\1/')
                
                if [ -n "$image_id" ]; then
                    # Check if this image is used by any running containers
                    local used_by_containers=$(docker ps --format "{{.Image}}" | grep -q "$image_id" && echo "used" || echo "unused")
                    
                    # Check if this image is used by any stopped containers
                    local used_by_stopped=$(docker ps -a --format "{{.Image}}" | grep -q "$image_id" && echo "used" || echo "unused")
                    
                    # If image is not used by any containers, it's safe to remove
                    if [ "$used_by_containers" = "unused" ] && [ "$used_by_stopped" = "unused" ]; then
                        echo "$image (unused - safe to remove)"
                    fi
                fi
            fi
        done
    fi
}

get_unused_volumes() {
    # Get unused volumes (excluding Harbor volumes) - just IDs for cleanup
    docker volume ls -q --filter "dangling=true" | grep -v "harbor" | grep -v "registry" || true
}

get_unused_volumes_with_usage() {
    # Get unused volumes (excluding Harbor volumes) with container usage info
    local volumes=$(docker volume ls -q --filter "dangling=true" | grep -v "harbor" | grep -v "registry" || true)
    
    if [ -n "$volumes" ]; then
        echo "$volumes" | while read -r volume; do
            if [ -n "$volume" ]; then
                # Get containers using this volume
                local containers=$(docker ps -a --filter "volume=$volume" --format "{{.Names}}" 2>/dev/null || echo "")
                # Get volume creation time
                local created=$(docker volume inspect "$volume" --format "{{.CreatedAt}}" 2>/dev/null || echo "unknown")
                # Get volume driver
                local driver=$(docker volume inspect "$volume" --format "{{.Driver}}" 2>/dev/null || echo "local")
                
                # Format the output with volume ID and usage info
                if [ -n "$containers" ]; then
                    echo "$volume (used by: $containers, created: $created, driver: $driver)"
                else
                    echo "$volume (unused - safe to remove, created: $created, driver: $driver)"
                fi
            fi
        done
    fi
}

get_unused_volumes_detailed() {
    # Get unused volumes (excluding Harbor volumes) with detailed information for display
    local volumes=$(docker volume ls -q --filter "dangling=true" | grep -v "harbor" | grep -v "registry" || true)
    
    if [ -n "$volumes" ]; then
        echo "$volumes" | while read -r volume; do
            if [ -n "$volume" ]; then
                # Get containers using this volume
                local containers=$(docker ps -a --filter "volume=$volume" --format "{{.Names}}" 2>/dev/null || echo "")
                # Get volume creation time
                local created=$(docker volume inspect "$volume" --format "{{.CreatedAt}}" 2>/dev/null || echo "unknown")
                # Get volume driver
                local driver=$(docker volume inspect "$volume" --format "{{.Driver}}" 2>/dev/null || echo "local")
                
                # Format the output with volume ID and usage info
                if [ -n "$containers" ]; then
                    echo "$volume (used by: $containers, created: $created, driver: $driver)"
                else
                    echo "$volume (unused - safe to remove, created: $created, driver: $driver)"
                fi
            fi
        done
    fi
}

get_unused_networks() {
    # Get unused networks (excluding Harbor networks) with detailed information
    docker network ls --filter "type=custom" --format "{{.Name}} ({{.Driver}}) - Created: {{.CreatedAt}}" | grep -v "harbor" | grep -v "registry" || true
}

get_protected_containers() {
    # Get all critical containers that are running
    local protected_containers=""
    IFS=',' read -ra CONTAINERS <<< "$CRITICAL_CONTAINERS"
    for container in "${CONTAINERS[@]}"; do
        if docker ps --format "{{.Names}}" | grep -q "^${container}$"; then
            if [ -n "$protected_containers" ]; then
                protected_containers="$protected_containers
  - $container (Harbor container)"
            else
                protected_containers="  - $container (Harbor container)"
            fi
        fi
    done
    echo "$protected_containers"
}

get_protected_volumes() {
    # Get Harbor volumes that would be protected
    docker volume ls -q --filter "dangling=true" | grep -E "(harbor|registry)" || true
}

get_protected_images() {
    # Get Harbor images that would be protected
    docker images --format "{{.Repository}}:{{.Tag}}" | grep -E "(goharbor|harbor)" || true
}

get_protected_networks() {
    # Get Harbor networks that would be protected
    docker network ls --format "{{.Name}}" | grep -E "(harbor|registry)" || true
}

check_critical_infrastructure() {
    log_info "Checking critical infrastructure status..."
    
    local missing_containers=""
    IFS=',' read -ra CONTAINERS <<< "$CRITICAL_CONTAINERS"
    
    for container in "${CONTAINERS[@]}"; do
        if ! docker ps --format "{{.Names}}" | grep -q "^${container}$"; then
            if [ -n "$missing_containers" ]; then
                missing_containers="$missing_containers, $container"
            else
                missing_containers="$container"
            fi
        fi
    done
    
    if [ -n "$missing_containers" ]; then
        log_warning "Some critical containers are not running: $missing_containers"
        log_warning "This may indicate infrastructure issues. Proceed with caution."
    else
        log_success "All critical infrastructure containers are running"
    fi
    
    # Check Forgejo runner service
    if systemctl is-active --quiet forgejo-runner.service; then
        log_success "Forgejo runner service is running"
    else
        log_warning "Forgejo runner service is not running"
    fi
    
    echo
}

label_critical_containers() {
    log_info "Labeling critical containers for protection..."
    
    IFS=',' read -ra CONTAINERS <<< "$CRITICAL_CONTAINERS"
    
    for container in "${CONTAINERS[@]}"; do
        if docker ps --format "{{.Names}}" | grep -q "^${container}$"; then
            # Add protection labels
            docker update --label critical=infrastructure "$container" 2>/dev/null || true
            docker update --label protected=true "$container" 2>/dev/null || true
            
            # Add specific labels based on container type
            if [[ "$container" == harbor* ]]; then
                docker update --label service=harbor "$container" 2>/dev/null || true
            elif [[ "$container" == forgejo* ]]; then
                docker update --label service=forgejo "$container" 2>/dev/null || true
            fi
        fi
    done
    
    log_success "Critical containers labeled for protection"
    echo
}

protect_harbor_volumes() {
    log_info "Identifying and protecting Harbor volumes..."
    
    # Use the same logic as the dry-run section to find Harbor volumes
    HARBOR_VOLUMES=$(docker volume ls -q --filter "dangling=true" | grep -E "(harbor|registry)" || true)
    
    if [ -n "$HARBOR_VOLUMES" ]; then
        log_info "Found Harbor volumes:"
        echo "$HARBOR_VOLUMES" | while read -r vol; do
            if [ -n "$vol" ]; then
                echo "  - $vol (PROTECTED)"
            fi
        done
    else
        log_info "No Harbor volumes found"
    fi
    
    log_success "Harbor volumes identified for protection"
    echo
}

show_help() {
    cat << EOF
Sharenet Cleanup Script

Usage: $0 [OPTIONS]

Options:
    --type TYPE        Cleanup type: ci-cd or production (default: ci-cd)
    --registry-dir DIR Registry directory (default: /opt/registry)
    --dry-run          Show what would be done without executing
    --help             Show this help message

Environment Variables:
    CLEANUP_TYPE       Set cleanup type (ci-cd/production)
    REGISTRY_DIR       Set registry directory path
    DRY_RUN            Set to 'true' for dry run mode

Examples:
    $0                                    # Cleanup CI/CD environment
    $0 --type production                  # Cleanup production environment
    $0 --dry-run                         # Show what would be cleaned
    DRY_RUN=true $0                      # Dry run mode
EOF
}

cleanup_docker_resources() {
    if [ "$DRY_RUN" = "true" ]; then
        echo
        echo "=================================================================================="
        echo "                            🚨 DRY RUN MODE 🚨"
        echo "                    No resources will be deleted"
        echo "=================================================================================="
        echo
        log_info "Checking what would be cleaned up..."
        echo

        # Get lists using shared functions
        local unused_images=$(get_unused_images)
        local unused_volumes=$(get_unused_volumes_with_usage)
        local unused_networks=$(get_unused_networks)
        local protected_containers=$(get_protected_containers)
        local protected_volumes=$(get_protected_volumes)
        local protected_images=$(get_protected_images)
        local protected_networks=$(get_protected_networks)

        # Section 1: ITEMS PROTECTED
        echo "=================================================================================="
        echo "                    🛡️  ITEMS PROTECTED 🛡️"
        echo "=================================================================================="
        echo

        if [ -n "$protected_containers" ]; then
            log_info "Protected Containers:"
            echo "$protected_containers"
            echo
        else
            log_info "Protected Containers: None found"
            echo
        fi

        if [ -n "$protected_volumes" ]; then
            log_info "Protected Volumes:"
            echo "$protected_volumes" | while read -r volume; do
                echo "  - $volume (Harbor volume)"
            done
            echo
        else
            log_info "Protected Volumes: None found"
            echo
        fi

        if [ -n "$protected_images" ]; then
            log_info "Protected Images:"
            echo "$protected_images" | while read -r image; do
                echo "  - $image (Harbor image)"
            done
            echo
        else
            log_info "Protected Images: None found"
            echo
        fi

        if [ -n "$protected_networks" ]; then
            log_info "Protected Networks:"
            echo "$protected_networks" | while read -r network; do
                echo "  - $network (Harbor network)"
            done
            echo
        else
            log_info "Protected Networks: None found"
            echo
        fi

        # Section 2: ITEMS REMOVED
        echo "=================================================================================="
        echo "                    🗑️  ITEMS REMOVED 🗑️"
        echo "=================================================================================="
        echo

        # Clean up unused images
        if [ -n "$unused_images" ]; then
            log_info "Images removed:"
            echo "$unused_images" | while read -r image; do
                echo "  - $image"
            done
            echo
            
            log_info "Removing unused images..."
            echo "$unused_images" | while read -r image; do
                # Extract image ID from the format "repo:tag (id) - Size: size, Created: date"
                local image_id=$(echo "$image" | sed 's/.*(\([a-f0-9]*\)).*/\1/')
                if [ -n "$image_id" ]; then
                    # First try to remove by image ID (this will remove all tags)
                    docker rmi "$image_id" 2>/dev/null || {
                        # If that fails, try to remove by repository:tag
                        local repo_tag=$(echo "$image" | cut -d' ' -f1)
                        if [ -n "$repo_tag" ]; then
                            docker rmi "$repo_tag" 2>/dev/null || log_warning "Failed to remove image: $repo_tag"
                        fi
                    }
                fi
            done
            log_success "Unused images removed"
        else
            log_info "Images removed: None found"
            log_info "No unused images to remove"
        fi

        # Clean up unused volumes
        if [ -n "$unused_volumes" ]; then
            log_info "Volumes removed:"
            local volumes_with_usage=$(get_unused_volumes_with_usage)
            if [ -n "$volumes_with_usage" ]; then
                echo "$volumes_with_usage" | while read -r volume_info; do
                    if [ -n "$volume_info" ]; then
                        echo "  - $volume_info"
                    fi
                done
            fi
            echo
            
            log_info "Removing unused volumes..."
            # Now remove the volumes
            echo "$unused_volumes" | while read -r volume; do
                if [ -n "$volume" ]; then
                    docker volume rm "$volume" 2>/dev/null || log_warning "Failed to remove volume: $volume"
                fi
            done
            log_success "Unused volumes removed"
        else
            log_info "Volumes removed: None found"
            log_info "No unused volumes to remove"
        fi

        # Clean up unused networks
        if [ -n "$unused_networks" ]; then
            log_info "Networks removed:"
            echo "$unused_networks" | while read -r network; do
                echo "  - $network"
            done
            echo
            
            log_info "Removing unused networks..."
            echo "$unused_networks" | while read -r network; do
                # Extract network name from the format "name (driver) - Created: date"
                local network_name=$(echo "$network" | cut -d' ' -f1)
                if [ -n "$network_name" ]; then
                    docker network rm "$network_name" 2>/dev/null || log_warning "Failed to remove network: $network_name"
                fi
            done
            log_success "Unused networks removed"
        else
            log_info "Networks removed: None found"
            log_info "No unused networks to remove"
        fi

        echo "=================================================================================="
        log_success "Dry run completed. No resources were deleted."
        echo "=================================================================================="
        return
    fi
    
    # Get lists using shared functions
    local unused_images=$(get_unused_images)
    local unused_volumes=$(get_unused_volumes)
    local unused_networks=$(get_unused_networks)
    local protected_containers=$(get_protected_containers)
    local protected_volumes=$(get_protected_volumes)
    local protected_images=$(get_protected_images)
    local protected_networks=$(get_protected_networks)

    # Section 1: ITEMS PROTECTED
    echo "=================================================================================="
    echo "                    🛡️  ITEMS PROTECTED 🛡️"
    echo "=================================================================================="
    echo

    if [ -n "$protected_containers" ]; then
        log_info "Protected Containers:"
        echo "$protected_containers"
        echo
    else
        log_info "Protected Containers: None found"
        echo
    fi

    if [ -n "$protected_volumes" ]; then
        log_info "Protected Volumes:"
        echo "$protected_volumes" | while read -r volume; do
            echo "  - $volume (Harbor volume)"
        done
        echo
    else
        log_info "Protected Volumes: None found"
        echo
    fi

    if [ -n "$protected_images" ]; then
        log_info "Protected Images:"
        echo "$protected_images" | while read -r image; do
            echo "  - $image (Harbor image)"
        done
        echo
    else
        log_info "Protected Images: None found"
        echo
    fi

    if [ -n "$protected_networks" ]; then
        log_info "Protected Networks:"
        echo "$protected_networks" | while read -r network; do
            echo "  - $network (Harbor network)"
        done
        echo
    else
        log_info "Protected Networks: None found"
        echo
    fi

    # Section 2: ITEMS REMOVED
    echo "=================================================================================="
    echo "                    🗑️  ITEMS REMOVED 🗑️"
    echo "=================================================================================="
    echo

    # Clean up unused images
    if [ -n "$unused_images" ]; then
        log_info "Images removed:"
        echo "$unused_images" | while read -r image; do
            echo "  - $image"
        done
        echo
        
        log_info "Removing unused images..."
        echo "$unused_images" | while read -r image; do
            # Extract image ID from the format "repo:tag (id) - Size: size, Created: date"
            local image_id=$(echo "$image" | sed 's/.*(\([a-f0-9]*\)).*/\1/')
            if [ -n "$image_id" ]; then
                # First try to remove by image ID (this will remove all tags)
                docker rmi "$image_id" 2>/dev/null || {
                    # If that fails, try to remove by repository:tag
                    local repo_tag=$(echo "$image" | cut -d' ' -f1)
                    if [ -n "$repo_tag" ]; then
                        docker rmi "$repo_tag" 2>/dev/null || log_warning "Failed to remove image: $repo_tag"
                    fi
                }
            fi
        done
        log_success "Unused images removed"
    else
        log_info "Images removed: None found"
        log_info "No unused images to remove"
    fi

    # Clean up unused volumes
    if [ -n "$unused_volumes" ]; then
        log_info "Volumes removed:"
        local volumes_with_usage=$(get_unused_volumes_with_usage)
        if [ -n "$volumes_with_usage" ]; then
            echo "$volumes_with_usage" | while read -r volume_info; do
                if [ -n "$volume_info" ]; then
                    echo "  - $volume_info"
                fi
            done
        fi
        echo
        
        log_info "Removing unused volumes..."
        # Now remove the volumes
        echo "$unused_volumes" | while read -r volume; do
            if [ -n "$volume" ]; then
                docker volume rm "$volume" 2>/dev/null || log_warning "Failed to remove volume: $volume"
            fi
        done
        log_success "Unused volumes removed"
    else
        log_info "Volumes removed: None found"
        log_info "No unused volumes to remove"
    fi

    # Clean up unused networks
    if [ -n "$unused_networks" ]; then
        log_info "Networks removed:"
        echo "$unused_networks" | while read -r network; do
            echo "  - $network"
        done
        echo
        
        log_info "Removing unused networks..."
        echo "$unused_networks" | while read -r network; do
            # Extract network name from the format "name (driver) - Created: date"
            local network_name=$(echo "$network" | cut -d' ' -f1)
            if [ -n "$network_name" ]; then
                docker network rm "$network_name" 2>/dev/null || log_warning "Failed to remove network: $network_name"
            fi
        done
        log_success "Unused networks removed"
    else
        log_info "Networks removed: None found"
        log_info "No unused networks to remove"
    fi

    echo "=================================================================================="
    log_success "Docker resources cleanup completed (critical infrastructure protected)"
    echo "=================================================================================="
}

cleanup_registry() {
    if [ "$CLEANUP_TYPE" != "ci-cd" ]; then
        log_info "Skipping registry cleanup (not CI/CD environment)"
        return
    fi
    
    log_info "Cleaning up Harbor registry..."
    
    # Check if Harbor containers are running
    if ! docker ps --format "{{.Names}}" | grep -q harbor; then
        log_warning "Harbor containers are not running, skipping registry cleanup"
        return
    fi
    
    if [ "$DRY_RUN" = "true" ]; then
        log_warning "DRY RUN MODE - No changes will be made"
        echo "Would run: Harbor registry garbage collection via API"
        return
    fi
    
    # Harbor garbage collection is typically done via the Harbor UI or API
    # For now, we'll just log that manual cleanup may be needed
    log_info "Harbor registry cleanup: Use Harbor UI to clean up old images"
    log_info "Manual cleanup: Go to Harbor UI → Projects → Select project → Artifacts → Delete old tags"
    
    log_success "Harbor registry cleanup info provided"
}

cleanup_production() {
    log_info "Cleaning up production environment..."
    
    # Check if we're in the application directory
    if [ -f "docker-compose.yml" ]; then
        log_info "Found docker-compose.yml, cleaning up application resources..."
        
        if [ "$DRY_RUN" = "true" ]; then
            log_warning "DRY RUN MODE - No changes will be made"
            echo "Would run: docker-compose down"
            echo "Would run: docker image prune -f"
            return
        fi
        
        # Stop containers to free up resources
        log_info "Stopping application containers..."
        docker-compose down
        
        # Clean up Docker resources
        cleanup_docker_resources
        
        # Start containers again
        log_info "Starting application containers..."
        docker-compose up -d
        
        log_success "Production cleanup completed"
    else
        log_warning "Not in application directory (docker-compose.yml not found)"
        cleanup_docker_resources
    fi
}

cleanup_ci_cd() {
    log_info "Cleaning up CI/CD environment..."
    
    # Check critical infrastructure before cleanup
    check_critical_infrastructure
    
    # Label critical containers for protection
    label_critical_containers
    
    # Protect Harbor volumes
    protect_harbor_volumes
    
    # Clean up Docker resources (this will show the organized sections)
    cleanup_docker_resources
    
    # Clean up registry
    cleanup_registry
    
    log_success "CI/CD cleanup completed"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --type)
            CLEANUP_TYPE="$2"
            shift 2
            ;;
        --registry-dir)
            REGISTRY_DIR="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN="true"
            shift
            ;;
        --help|-h)
            show_help
            exit 0
            ;;
        *)
            log_error "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
done

# Main cleanup logic
case "$CLEANUP_TYPE" in
    production)
        cleanup_production
        ;;
    ci-cd)
        cleanup_ci_cd
        ;;
    *)
        log_error "Invalid cleanup type: $CLEANUP_TYPE"
        log_error "Valid types: production, ci-cd"
        exit 1
        ;;
esac

log_success "Cleanup completed successfully"