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

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
}

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() {
    log_info "Cleaning up Docker resources..."
    
    if [ "$DRY_RUN" = "true" ]; then
        log_warning "DRY RUN MODE - No changes will be made"
        echo "Would run: docker image prune -f (excluding critical infrastructure)"
        echo "Would run: docker volume prune -f (excluding critical infrastructure)"
        echo "Would run: docker network prune -f (excluding critical infrastructure)"
        return
    fi
    
    # Remove unused images (excluding critical infrastructure)
    log_info "Removing unused Docker images (excluding critical infrastructure)..."
    # Use protection labels to exclude critical images
    docker image prune -f --filter "label!=critical=infrastructure" --filter "label!=protected=true"
    
    # Remove unused volumes (excluding critical infrastructure)
    log_info "Removing unused Docker volumes (excluding critical infrastructure)..."
    # Use protection labels to exclude critical volumes
    docker volume prune -f --filter "label!=critical=infrastructure" --filter "label!=protected=true"
    
    # Remove unused networks (excluding critical infrastructure)
    log_info "Removing unused Docker networks (excluding critical infrastructure)..."
    # Use protection labels to exclude critical networks
    docker network prune -f --filter "label!=critical=infrastructure" --filter "label!=protected=true"
    
    log_success "Docker resources cleanup completed (critical infrastructure protected)"
}

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
    
    # Clean up Docker resources
    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"