#!/bin/bash

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

# 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 hello-world images with force
log_message "Attempting to remove hello-world images"
HELLO_WORLD_IMAGES=$($DOCKER_PATH images --format "{{.Repository}}:{{.Tag}}" | grep "^hello-world")
if [ ! -z "$HELLO_WORLD_IMAGES" ]; then
    echo "$HELLO_WORLD_IMAGES" | while read -r image; do
        log_message "Force removing hello-world image: $image"
        if $DOCKER_PATH rmi -f "$image" > /dev/null 2>&1; then
            log_message "Successfully removed: $image"
            ((REMOVED_COUNT++))
        else
            # If force removal fails, try removing any stopped containers using this image
            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"
                    ((REMOVED_COUNT++))
                else
                    log_message "Failed to remove: $image even after container cleanup"
                    ((FAILED_COUNT++))
                fi
            else
                log_message "Failed to remove: $image"
                ((FAILED_COUNT++))
            fi
        fi
    done
fi

# Get list of Pterodactyl images
PTERO_IMAGES=$($DOCKER_PATH images --format "{{.Repository}}:{{.Tag}}" | grep -E "pterodactyl|pterodactyl/yolks|ghcr.io/pterodactyl")

if [ ! -z "$PTERO_IMAGES" ]; then
    log_message "Processing Pterodactyl images"
    # Remove each unused Pterodactyl image
    echo "$PTERO_IMAGES" | while read -r image; do
        # Check if image is being used by any container
        if ! $DOCKER_PATH ps -a --format "{{.Image}}" | grep -q "^${image}$"; then
            log_message "Attempting to remove image: $image"
            if $DOCKER_PATH rmi "$image" > /dev/null 2>&1; then
                log_message "Successfully removed: $image"
                ((REMOVED_COUNT++))
            else
                log_message "Failed to remove: $image (might be in use)"
                ((FAILED_COUNT++))
            fi
        else
            log_message "Skipping $image - currently in use"
        fi
    done
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