package org.testcontainers.utility;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.DockerException;
import com.github.dockerjava.api.InternalServerErrorException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.DockerClientFactory;

import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Component that responsible for container removal and automatic cleanup of dead containers at JVM shutdown.
 */
public final class ContainerReaper {

    private static final Logger LOGGER = LoggerFactory.getLogger(ContainerReaper.class);
    private static ContainerReaper instance;
    private final DockerClient dockerClient;
    private Map<String, String> registeredContainers = new ConcurrentHashMap<>();

    private ContainerReaper() {
        dockerClient = DockerClientFactory.instance().client();

        // If the JVM stops without containers being stopped, try and stop the container.
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            for (Map.Entry<String, String> entry : new HashSet<>(registeredContainers.entrySet())) {
                stopAndRemoveContainer(entry.getKey(), entry.getValue());
            }
        }));
    }

    public synchronized static ContainerReaper instance() {
        if (instance == null) {
            instance = new ContainerReaper();
        }

        return instance;
    }

    /**
     * Register a container to be cleaned up, either on explicit call to stopAndRemoveContainer, or at JVM shutdown.
     * @param containerId the ID of the container
     * @param imageName the image name of the container (used for logging)
     */
    public void registerContainerForCleanup(String containerId, String imageName) {
        registeredContainers.put(containerId, imageName);
    }

    /**
     * Stop a potentially running container and remove it, including associated volumes.
     * @param containerId the ID of the container
     * @param imageName the image name of the container (used for logging)
     */
    public void stopAndRemoveContainer(String containerId, String imageName) {
        try {
            LOGGER.trace("Stopping container: {}", containerId);
            dockerClient.killContainerCmd(containerId).exec();
            LOGGER.trace("Stopped container: {}", imageName);
        } catch (DockerException e) {
            LOGGER.trace("Error encountered shutting down container (ID: {}) - it may not have been stopped, or may already be stopped: {}", containerId, e.getMessage());
        }

        try {
            LOGGER.trace("Stopping container: {}", containerId);
            try {
                dockerClient.removeContainerCmd(containerId).withRemoveVolumes(true).withForce(true).exec();
                LOGGER.info("Removed container and associated volume(s): {}", imageName);
            } catch (InternalServerErrorException e) {
                LOGGER.trace("Exception when removing container with associated volume(s): {} (due to {})", imageName, e.getMessage());
            }
        } catch (DockerException e) {
            LOGGER.trace("Error encountered shutting down container (ID: {}) - it may not have been stopped, or may already be stopped: {}", containerId, e.getMessage());
        }

        registeredContainers.remove(containerId);
    }
}
@@ -299,34 +299,8 @@
 y;%0A%0A
-import java.util.HashSet;%0A
 impo
@@ -317,16 +317,16 @@
 il.Map;%0A
+
 import j
@@ -1050,184 +1050,58 @@
     
-for (Map.Entry%3CString, String%3E entry : new HashSet%3C%3E(registeredContainers.entrySet())) %7B%0A                stopAndRemoveContainer(entry.getKey(), entry.getValue());%0A            %7D
+registeredContainers.forEach(this::stopContainer);
 %0A   
@@ -2006,32 +2006,208 @@
 ng imageName) %7B%0A
+        stopContainer(containerId, imageName);%0A%0A        registeredContainers.remove(containerId);%0A    %7D%0A%0A    private void stopContainer(String containerId, String imageName) %7B%0A
         try %7B%0A  
@@ -3316,32 +3316,32 @@
 .getMessage());%0A
+
         %7D%0A%0A     
@@ -3337,59 +3337,8 @@
    %7D
-%0A%0A        registeredContainers.remove(containerId);
 %0A   

Tidy up removal code so that containers to be removed is not modified during JVM shutdown hook, as it doesn't need to and causes concurrent modification issues.