#!/bin/bash
set -euo pipefail

# Check for -rt flag
REMOVE_TAR=false
if [ "${1:-}" = "-rt" ]; then
    REMOVE_TAR=true
fi

# Images and tags
BASE_IMAGE="jobarg-server-base"
BASE_TAG="rhel9"
COMMON_TAG="rhel9-6.2.2"

MARIA_IMAGE="jobarg-server-maria"
MYSQL_IMAGE="jobarg-server-mysql"
PGSQL_IMAGE="jobarg-server-postgres"

IMAGES=(
    "${MARIA_IMAGE}:${COMMON_TAG}"
    "${MYSQL_IMAGE}:${COMMON_TAG}"
    "${PGSQL_IMAGE}:${COMMON_TAG}"
    "${BASE_IMAGE}:${BASE_TAG}"
)

echo "Starting cleanup..."

for IMAGE in "${IMAGES[@]}"; do
    # Stop running containers
    RUNNING_CONTAINERS=$(docker ps -q --filter "ancestor=$IMAGE")
    if [ -n "$RUNNING_CONTAINERS" ]; then
        echo "Stopping running containers for image $IMAGE..."
        docker stop $RUNNING_CONTAINERS
    fi

    # Remove all containers for this image
    ALL_CONTAINERS=$(docker ps -aq --filter "ancestor=$IMAGE")
    if [ -n "$ALL_CONTAINERS" ]; then
        echo "Removing containers for image $IMAGE..."
        docker rm $ALL_CONTAINERS
    fi

    # Remove the image
    if docker image inspect "$IMAGE" >/dev/null 2>&1; then
        echo "Removing image $IMAGE..."
        docker rmi -f "$IMAGE"
    fi

    # Remove tar.gz if -rt flag is set
    if [ "$REMOVE_TAR" = true ]; then
        case "$IMAGE" in
            "${MARIA_IMAGE}:${COMMON_TAG}") DIR="db/maria" ;;
            "${MYSQL_IMAGE}:${COMMON_TAG}") DIR="db/mysql" ;;
            "${PGSQL_IMAGE}:${COMMON_TAG}") DIR="db/postgres" ;;
            "${BASE_IMAGE}:${BASE_TAG}") DIR="db/base" ;;
            *) DIR="" ;;
        esac

        if [ -n "$DIR" ] && [ -d "$DIR" ]; then
            # Find tar.gz files that match the image prefix
            PREFIX=$(echo "$IMAGE" | cut -d':' -f1)  # e.g., jobarg-server-maria
            for TAR_FILE in "$DIR"/"$PREFIX"*".tar.gz"; do
                # Check if file exists
                if [ -f "$TAR_FILE" ]; then
                    echo "Removing tar file $TAR_FILE..."
                    rm -f "$TAR_FILE"
                fi
            done
        fi
    fi
done

docker images

echo "Cleanup complete!"