package com.engebretson;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.junit.Test;

import org.apache.logging.log4j.ThreadContext;

public class CloseableThreadContextSpeedTest {
    public static class OptimizedInstance implements AutoCloseable {

        private final Map<String, String> originalValues = new HashMap<>();

        private OptimizedInstance() {
        }

        @Override
        public void close() {
            closeMap();
        }

        private void closeMap() {
            Map<String, String> valuesToReplace = new HashMap<>(originalValues.size());
            List<String> keysToRemove = new ArrayList<>(originalValues.size());
            for (Map.Entry<String, String> entry : originalValues.entrySet()) {
                final String key = entry.getKey();
                final String originalValue = entry.getValue();
                if (null == originalValue) {
                    keysToRemove.add(key);
                } else {
                    valuesToReplace.put(key, originalValue);
                }
            }
            if (!valuesToReplace.isEmpty()) {
                ThreadContext.putAll(valuesToReplace);
            }
            if (!keysToRemove.isEmpty()) {
                ThreadContext.removeAll(keysToRemove);
            }
        }

        /**
         * Populates the Thread Context Map with the supplied key/value pairs. Any
         * existing keys in the {@link ThreadContext} will be replaced with the supplied
         * values, and restored back to their original value when the instance is
         * closed.
         *
         * @param values The map of key/value pairs to be added
         * @return a new instance that will back out the changes when closed.
         * @since 2.8
         */
        public OptimizedInstance putAll(final Map<String, String> values) {
            final Map<String, String> currentValues = ThreadContext.getContext();
            ThreadContext.putAll(values);
            for (final String key : values.keySet()) {
                if (!originalValues.containsKey(key)) {
                    originalValues.put(key, currentValues.get(key));
                }
            }
            return this;
        }
    }

    public static class OriginalInstance implements AutoCloseable {

        private final Map<String, String> originalValues = new HashMap<>();

        private OriginalInstance() {
        }

        @Override
        public void close() {
            closeMap();
        }

        private void closeMap() {
            for (final Iterator<Map.Entry<String, String>> it = originalValues.entrySet().iterator(); it.hasNext();) {
                final Map.Entry<String, String> entry = it.next();
                final String key = entry.getKey();
                final String originalValue = entry.getValue();
                if (null == originalValue) {
                    ThreadContext.remove(key);
                } else {
                    ThreadContext.put(key, originalValue);
                }
                it.remove();
            }
        }

        /**
         * Populates the Thread Context Map with the supplied key/value pairs. Any
         * existing keys in the {@link ThreadContext} will be replaced with the supplied
         * values, and restored back to their original value when the instance is
         * closed.
         *
         * @param values The map of key/value pairs to be added
         * @return a new instance that will back out the changes when closed.
         * @since 2.8
         */
        public OriginalInstance putAll(final Map<String, String> values) {
            final Map<String, String> currentValues = ThreadContext.getContext();
            ThreadContext.putAll(values);
            for (final String key : values.keySet()) {
                if (!originalValues.containsKey(key)) {
                    originalValues.put(key, currentValues.get(key));
                }
            }
            return this;
        }
    }

    private static final int NUM_TEST_ITERATIONS = 400000;

    private static final int NUM_TESTS = 5;

    private final Map<String, String> contextValuesToAdd;

    public CloseableThreadContextSpeedTest() {
        contextValuesToAdd = new HashMap<>();
        // Setting six parameters, which matches our production application
        contextValuesToAdd.put("arg1", "1");
        contextValuesToAdd.put("arg2", "2");
        contextValuesToAdd.put("arg3", "3");
        contextValuesToAdd.put("arg4", "4");
        contextValuesToAdd.put("arg5", "5");
        contextValuesToAdd.put("arg6", "6");
    }

    protected void runNewTest() {
        try (OptimizedInstance instance = new OptimizedInstance()) {
            instance.putAll(contextValuesToAdd);
        }
    }

    protected void runOldTest() {
        try (OriginalInstance instance = new OriginalInstance()) {
            instance.putAll(contextValuesToAdd);
        }
    }

    @Test
    public void runTests() {
        // warmup
        for (int i = 0; i < NUM_TEST_ITERATIONS; i++) {
            runOldTest();
            runNewTest();
        }

        for (int i = 0; i < NUM_TESTS; i++) {
            System.gc();
            long startOld = System.currentTimeMillis();
            for (int j = 0; j < NUM_TEST_ITERATIONS; j++) {
                runOldTest();
            }
            System.out.println("Done with old in " + (System.currentTimeMillis() - startOld));
        }

        for (int i = 0; i < NUM_TESTS; i++) {
            System.gc();
            long startNew = System.currentTimeMillis();
            for (int j = 0; j < NUM_TEST_ITERATIONS; j++) {
                runNewTest();
            }
            System.out.println("Done with new in " + (System.currentTimeMillis() - startNew));
        }
    }

}