#!/bin/bash

###########################################
# Configuration - Add new patterns here
###########################################
# Add one pattern per line between the quotes
# Use regex patterns (grep -E compatible)
# Examples:
#   "^nginx"      - matches images starting with nginx
#   "alpine:.*"   - matches all alpine tagged images
#   "mysql|redis" - matches mysql or redis images
IMAGE_PATTERNS=(
    "^hello-world"
    "pterodactyl"
    "pterodactyl/yolks"
    "ghcr.io/pterodactyl"
    "ghcr.io/parkervcp/games"
    "ghcr.io/parkervcp/yolks"
    "ghcr.io/parkervcp/steamcmd"
    # Add new patterns below this line
)

# Log file setup
LOG_FILE="/var/log/ptero-cleanup.log"
DOCKER_PATH=$(which docker)

# Function to log messages
log_message() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

# Function to force remove image and its containers
force_remove_image() {
    local image=$1
    log_message "Force removing image: $image"
    
    # First try to remove dangling images specifically
    if [[ $image == *"<none>"* ]]; then
        local IMAGE_ID=$(echo $image | awk '{print $3}')
        log_message "Detected dangling image with ID: $IMAGE_ID"
        if $DOCKER_PATH rmi -f $IMAGE_ID > /dev/null 2>&1; then
            log_message "Successfully removed dangling image: $IMAGE_ID"
            return 0
        fi
    fi
    
    # Regular force remove attempt
    if $DOCKER_PATH rmi -f "$image" > /dev/null 2>&1; then
        log_message "Successfully removed: $image"
        return 0
    else
        # If force removal fails, try removing any stopped containers using this image
        local CONTAINERS=$($DOCKER_PATH ps -a --filter "ancestor=$image" --format "{{.ID}}")
        if [ ! -z "$CONTAINERS" ]; then
            log_message "Removing stopped containers using $image"
            echo "$CONTAINERS" | xargs -r $DOCKER_PATH rm -f
            # Try removing the image again
            if $DOCKER_PATH rmi -f "$image" > /dev/null 2>&1; then
                log_message "Successfully removed: $image after container cleanup"
                return 0
            else
                log_message "Failed to remove: $image even after container cleanup"
                return 1
            fi
        else
            log_message "Failed to remove: $image"
            return 1
        fi
    fi
}

# Check if docker is installed
if [ -z "$DOCKER_PATH" ]; then
    log_message "ERROR: Docker is not installed or not in PATH"
    exit 1
fi

# Check if script is run as root
if [ "$EUID" -ne 0 ]; then
    log_message "ERROR: Please run as root"
    exit 1
fi

log_message "Starting image cleanup"

# Counter for removed images
REMOVED_COUNT=0
FAILED_COUNT=0

# First, try to remove all dangling images specifically
log_message "Checking for dangling images first..."
DANGLING_IMAGES=$($DOCKER_PATH images -f "dangling=true" --format "{{.Repository}} {{.Tag}} {{.ID}}")
if [ ! -z "$DANGLING_IMAGES" ]; then
    echo "$DANGLING_IMAGES" | while read -r image; do
        if force_remove_image "$image"; then
            ((REMOVED_COUNT++))
        else
            ((FAILED_COUNT++))
        fi
    done
fi

# Build grep pattern from IMAGE_PATTERNS array
GREP_PATTERN=$(IFS="|"; echo "${IMAGE_PATTERNS[*]}")
log_message "Using pattern: $GREP_PATTERN"

# Get all target images
TARGET_IMAGES=$($DOCKER_PATH images --format "{{.Repository}}:{{.Tag}}" | grep -E "$GREP_PATTERN")

if [ ! -z "$TARGET_IMAGES" ]; then
    # Process all images
    log_message "Processing target images"
    echo "$TARGET_IMAGES" | while read -r image; do
        if force_remove_image "$image"; then
            ((REMOVED_COUNT++))
        else
            ((FAILED_COUNT++))
        fi
    done
else
    log_message "No target images found matching patterns"
fi

# Display summary
log_message "Cleanup completed"
log_message "Images removed: $REMOVED_COUNT"
log_message "Failed removals: $FAILED_COUNT"

# Clean up old log files (keep last 5 days)
find "$(dirname "$LOG_FILE")" -name "$(basename "$LOG_FILE")*" -mtime +5 -delete

exit 0