package com.blazemeter.jmeter.control;

import kg.apc.emulators.TestJMeterUtils;
import org.apache.jmeter.JMeter;
import org.apache.jmeter.control.GenericController;
import org.apache.jmeter.control.LoopController;
import org.apache.jmeter.control.TransactionController;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.gui.util.PowerTableModel;
import org.apache.jmeter.reporters.ResultCollector;
import org.apache.jmeter.sampler.DebugSampler;
import org.apache.jmeter.samplers.SampleEvent;
import org.apache.jmeter.samplers.SampleListener;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.samplers.Sampler;
import org.apache.jmeter.save.SaveService;
import org.apache.jmeter.threads.JMeterContextService;
import org.apache.jmeter.threads.JMeterThread;
import org.apache.jmeter.threads.JMeterVariables;
import org.apache.jmeter.threads.ListenerNotifier;
import org.apache.jmeter.threads.TestCompiler;
import org.apache.jmeter.threads.ThreadGroup;
import org.apache.jorphan.collections.HashTree;
import org.apache.jorphan.collections.ListedHashTree;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;
import org.junit.BeforeClass;
import org.junit.Test;

import java.io.File;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

public class WeightedSwitchControllerTest {

    @BeforeClass
    public static void setUpClass()
            throws Exception {
        TestJMeterUtils.createJmeterEnv();
    }

    private static final Logger log = LoggingManager.getLoggerForClass();

    @Test
    public void testGeneric() {
        WeightedSwitchController obj = new WeightedSwitchController();
        obj.addTestElement(getSampler("0"));
        obj.addTestElement(getSampler("1"));
        obj.addTestElement(getSampler("2"));
        obj.addTestElement(getSampler("5"));
        PowerTableModel mdl = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        mdl.addRow(new String[]{"0", "0"});
        mdl.addRow(new String[]{"1", "1"});
        mdl.addRow(new String[]{"2", "2"});
        mdl.addRow(new String[]{"5", "5"});
        obj.setData(mdl);

        for (int n = 0; n < 800; n++) {
            Sampler s = obj.next();
            log.info("Sampler: " + s.getName());
            assertNotNull(s);
            assertNull(obj.next());
            assertNotEquals(s.getName(), "0");
        }
    }

    private DebugSampler getSampler(String s) {
        DebugSampler debugSampler = new DebugSampler();
        debugSampler.setName(s);
        return debugSampler;
    }

    @Test
    public void testComb() {
        long[] sums = new long[]{0L, 0L};
        long[] prio = new long[]{2, 3};
        long[] choices = new long[]{0, 0};
        for (int i = 0; i < 100; i++) {
            // inc counts
            for (int k = 0; k < sums.length; k++) {
                sums[k] += prio[k];
            }

            if (sums[0] >= sums[1]) {
                // 2
                sums[0] -= prio[0];
                choices[0]++;
            } else {
                // 3
                sums[1] -= prio[1];
                choices[1]++;
            }

            // minimize nums
            long min;
            if (sums[0] >= sums[1]) {
                min = sums[1];
            } else {
                min = sums[0];
            }
            for (int z = 0; z < sums.length; z++) {
                sums[z] -= min;
            }

        }
        System.out.print(Arrays.toString(choices));
    }


    @Test
    public void testNestedWSC() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());


        TestSampleListener listener = new TestSampleListener();

        // top WSC
        WeightedSwitchController topWSC = new WeightedSwitchController();
        PowerTableModel topPTM = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        topPTM.addRow(new String[]{"wsc1", "2"});
        topPTM.addRow(new String[]{"wsc2", "2"});
        topPTM.addRow(new String[]{"D_#1", "1"});
        topPTM.addRow(new String[]{"D_#2", "1"});
        topWSC.setData(topPTM);

        DebugSampler d1 = new DebugSampler();
        d1.setName("D_#1");
        DebugSampler d2 = new DebugSampler();
        d2.setName("D_#2");

        // first child WSC of top WSC
        WeightedSwitchController childWSC1 = new WeightedSwitchController();
        childWSC1.setName("wsc1");
        PowerTableModel childPTM1 = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        childPTM1.addRow(new String[]{"D1_#1", "1"});
        childPTM1.addRow(new String[]{"D1_#2", "1"});
        childWSC1.setData(childPTM1);

        DebugSampler d1_1 = new DebugSampler();
        d1_1.setName("D1_#1");
        DebugSampler d1_2 = new DebugSampler();
        d1_2.setName("D1_#2");

        // second child WSC of top WSC
        WeightedSwitchController childWSC2 = new WeightedSwitchController();
        childWSC2.setName("wsc2");
        PowerTableModel childPTM2 = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        childPTM2.addRow(new String[]{"D2_#1", "1"});
        childPTM2.addRow(new String[]{"D2_#2", "1"});
        childWSC2.setData(childPTM2);

        DebugSampler d2_1 = new DebugSampler();
        d2_1.setName("D2_#1");
        DebugSampler d2_2 = new DebugSampler();
        d2_2.setName("D2_#2");

        // main loop
        LoopController loop = new LoopController();
        loop.setLoops(6);
        loop.setContinueForever(false);

        // test tree
        ListedHashTree hashTree = new ListedHashTree();
        hashTree.add(loop);
        hashTree.add(loop, topWSC);
        hashTree.add(topWSC, listener);
        hashTree.add(topWSC, childWSC1);
        hashTree.add(childWSC1, d1_1);
        hashTree.add(childWSC1, d1_2);
        hashTree.add(childWSC1, listener);
        hashTree.add(topWSC, childWSC2);
        hashTree.add(childWSC2, d2_1);
        hashTree.add(childWSC2, d2_2);
        hashTree.add(childWSC2, listener);
        hashTree.add(topWSC, d1);
        hashTree.add(topWSC, d2);

        TestCompiler compiler = new TestCompiler(hashTree);
        hashTree.traverse(compiler);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setNumThreads(1);

        ListenerNotifier notifier = new ListenerNotifier();

        JMeterThread thread = new JMeterThread(hashTree, threadGroup, notifier);
        thread.setThreadGroup(threadGroup);
        thread.setOnErrorStopThread(true);
        thread.run();

        assertEquals(6, listener.events.size());
        List<String> labels = new ArrayList<>();
        labels.add("D_#1");
        labels.add("D_#2");
        labels.add("D1_#1");
        labels.add("D1_#2");
        labels.add("D2_#1");
        labels.add("D2_#2");
        for (SampleEvent event : listener.events) {
            assertTrue(labels.contains(event.getResult().getSampleLabel()));
        }
    }

    // https://groups.google.com/forum/?hl=ru#!searchin/jmeter-plugins/Weighted$20Switch$20Controller|sort:relevance/jmeter-plugins/P9Nx9OqgWj4/aeC-RaNgCAAJ
    @Test
    public void testNestedSimpleControllers() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());


        TestSampleListener listener = new TestSampleListener();

        // top WSC
        WeightedSwitchController topWSC = new WeightedSwitchController();
        PowerTableModel topPTM = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        topPTM.addRow(new String[]{"ex1", "10"});
        topPTM.addRow(new String[]{"ex2", "20"});
        topWSC.setData(topPTM);


        // first child: simple controller
        GenericController ex1 = new GenericController();
        ex1.setName("ex1");
        DebugSampler example1_1 = new DebugSampler();
        example1_1.setName("example1_1");
        DebugSampler example1_2 = new DebugSampler();
        example1_2.setName("example1_2");

        // second child: simple controller
        GenericController ex2 = new GenericController();
        ex2.setName("ex2");
        DebugSampler example2_1 = new DebugSampler();
        example2_1.setName("example2_1");
        DebugSampler example2_2 = new DebugSampler();
        example2_2.setName("example2_2");

        // main loop
        LoopController loop = new LoopController();
        loop.setLoops(60);
        loop.setContinueForever(false);

        // test tree
        ListedHashTree hashTree = new ListedHashTree();
        hashTree.add(loop);
        hashTree.add(loop, topWSC);
        hashTree.add(topWSC, listener);
        hashTree.add(topWSC, ex1);
        hashTree.add(ex1, example1_1);
        hashTree.add(ex1, example1_2);
        hashTree.add(ex1, listener);
        hashTree.add(topWSC, ex2);
        hashTree.add(ex2, example2_1);
        hashTree.add(ex2, example2_2);
        hashTree.add(ex2, listener);

        TestCompiler compiler = new TestCompiler(hashTree);
        hashTree.traverse(compiler);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setNumThreads(1);

        ListenerNotifier notifier = new ListenerNotifier();

        JMeterThread thread = new JMeterThread(hashTree, threadGroup, notifier);
        thread.setThreadGroup(threadGroup);
        thread.setOnErrorStopThread(true);
        thread.run();

        Map<String, Integer> totalResults = new HashMap<>();
        for (SampleEvent event : listener.events) {
            String label = event.getResult().getSampleLabel();
            if (totalResults.containsKey(label)) {
                totalResults.put(label, totalResults.get(label) + 1);
            } else {
                totalResults.put(label, 1);
            }
        }

        assertEquals(120, listener.events.size());
        assertEquals(20, (int) totalResults.get("example1_1"));
        assertEquals(20, (int) totalResults.get("example1_2"));
        assertEquals(40, (int) totalResults.get("example2_1"));
        assertEquals(40, (int) totalResults.get("example2_2"));
    }

    @Test
    public void testNestedTransactionControllers() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());


        TestSampleListener listener = new TestSampleListener();

        // top WSC
        WeightedSwitchController topWSC = new WeightedSwitchController();
        PowerTableModel topPTM = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        topPTM.addRow(new String[]{"ex1", "10"});
        topPTM.addRow(new String[]{"ex2", "20"});
        topWSC.setData(topPTM);


        // first child: transaction controller
        TransactionController ex1 = new TransactionController();
        ex1.setName("ex1");
        DebugSampler example1_1 = new DebugSampler();
        example1_1.setName("example1_1");
        DebugSampler example1_2 = new DebugSampler();
        example1_2.setName("example1_2");

        // second child: transaction controller
        TransactionController ex2 = new TransactionController();
        ex2.setName("ex2");
        DebugSampler example2_1 = new DebugSampler();
        example2_1.setName("example2_1");
        DebugSampler example2_2 = new DebugSampler();
        example2_2.setName("example2_2");

        // main loop
        LoopController loop = new LoopController();
        loop.setLoops(3);
        loop.setContinueForever(false);

        // test tree
        ListedHashTree hashTree = new ListedHashTree();
        hashTree.add(loop);
        hashTree.add(loop, topWSC);
        hashTree.add(topWSC, listener);
        hashTree.add(topWSC, ex1);
        hashTree.add(ex1, example1_1);
        hashTree.add(ex1, example1_2);
        hashTree.add(ex1, listener);
        hashTree.add(topWSC, ex2);
        hashTree.add(ex2, example2_1);
        hashTree.add(ex2, example2_2);
        hashTree.add(ex2, listener);

        TestCompiler compiler = new TestCompiler(hashTree);
        hashTree.traverse(compiler);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setNumThreads(1);

        ListenerNotifier notifier = new ListenerNotifier();

        JMeterThread thread = new JMeterThread(hashTree, threadGroup, notifier);
        thread.setThreadGroup(threadGroup);
        thread.setOnErrorStopThread(true);
        thread.run();

        Map<String, Integer> totalResults = new HashMap<>();
        for (SampleEvent event : listener.events) {
            String label = event.getResult().getSampleLabel();
            if (totalResults.containsKey(label)) {
                totalResults.put(label, totalResults.get(label) + 1);
            } else {
                totalResults.put(label, 1);
            }
        }

        assertEquals(9, listener.events.size());
        assertEquals(1, (int) totalResults.get("example1_1"));
        assertEquals(1, (int) totalResults.get("example1_2"));
        assertEquals(2, (int) totalResults.get("example2_1"));
        assertEquals(2, (int) totalResults.get("example2_2"));
        assertEquals(1, (int) totalResults.get("ex1")); // transaction result
        assertEquals(2, (int) totalResults.get("ex2")); // transaction result
    }


    @Test
    public void testNestedTransactionControllersWithParentSample() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());

        TestSampleListener listener = new TestSampleListener();

        // top WSC
        WeightedSwitchController topWSC = new WeightedSwitchController();
        PowerTableModel topPTM = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        topPTM.addRow(new String[]{"ex1", "10"});
        topPTM.addRow(new String[]{"ex2", "20"});
        topPTM.addRow(new String[]{"ex3", "30"});
        topWSC.setData(topPTM);


        // first child: transaction controller
        TransactionController ex1 = new TransactionController();
        ex1.setName("ex1");
        ex1.setParent(true);
        DebugSampler example1_1 = new DebugSampler();
        example1_1.setName("example1_1");
        DebugSampler example1_2 = new DebugSampler();
        example1_2.setName("example1_2");

        // second child: transaction controller
        TransactionController ex2 = new TransactionController();
        ex2.setName("ex2");
        ex2.setParent(true);
        DebugSampler example2_1 = new DebugSampler();
        example2_1.setName("example2_1");
        DebugSampler example2_2 = new DebugSampler();
        example2_2.setName("example2_2");

        // third child: transaction controller
        TransactionController ex3 = new TransactionController();
        ex3.setName("ex3");
        ex3.setParent(true);
        DebugSampler example3_1 = new DebugSampler();
        example3_1.setName("example3_1");
        DebugSampler example3_2 = new DebugSampler();
        example3_2.setName("example3_2");

        // main loop
        LoopController loop = new LoopController();
        loop.setLoops(6);
        loop.setContinueForever(false);

        // test tree
        ListedHashTree hashTree = new ListedHashTree();
        hashTree.add(loop);
        hashTree.add(loop, topWSC);
        hashTree.add(topWSC, listener);
        hashTree.add(topWSC, ex1);
        hashTree.add(ex1, example1_1);
        hashTree.add(ex1, example1_2);
        hashTree.add(ex1, listener);
        hashTree.add(topWSC, ex2);
        hashTree.add(ex2, example2_1);
        hashTree.add(ex2, example2_2);
        hashTree.add(ex2, listener);
        hashTree.add(topWSC, ex3);
        hashTree.add(ex3, example3_1);
        hashTree.add(ex3, example3_2);
        hashTree.add(ex3, listener);

        TestCompiler compiler = new TestCompiler(hashTree);
        hashTree.traverse(compiler);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setNumThreads(1);

        ListenerNotifier notifier = new ListenerNotifier();

        JMeterThread thread = new JMeterThread(hashTree, threadGroup, notifier);
        thread.setThreadGroup(threadGroup);
        thread.setOnErrorStopThread(true);
        thread.run();

        Map<String, Integer> totalResults = new HashMap<>();
        for (SampleEvent event : listener.events) {
            String label = event.getResult().getSampleLabel();
            if (totalResults.containsKey(label)) {
                totalResults.put(label, totalResults.get(label) + 1);
            } else {
                totalResults.put(label, 1);
            }
        }

        assertEquals(6, listener.events.size());
        assertEquals(1, (int) totalResults.get("ex1")); // transaction result
        assertEquals(2, (int) totalResults.get("ex2")); // transaction result
        assertEquals(3, (int) totalResults.get("ex3")); // transaction result

        for (SampleEvent event : listener.events) {
            SampleResult[] subResults = event.getResult().getSubResults();
            assertNotNull(subResults);
            assertEquals(2, subResults.length);
        }
    }

    @Test
    public void testResetTransactionControllerUnderSimpleController() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());


        TestSampleListener listener = new TestSampleListener();

        // top WSC
        WeightedSwitchController topWSC = new WeightedSwitchController();
        PowerTableModel topPTM = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
        topPTM.addRow(new String[]{"simple1", "100"});
        topPTM.addRow(new String[]{"simple2", "100"});
        topWSC.setData(topPTM);

        GenericController simple1 = new GenericController();
        simple1.setName("simple1");

        GenericController simple2 = new GenericController();
        simple2.setName("simple2");

        // first child: transaction controller
        TransactionController ex1 = new TransactionController();
        ex1.setName("ex1");
        DebugSampler example1_1 = new DebugSampler();
        example1_1.setName("example1_1");
        DebugSampler example1_2 = new DebugSampler();
        example1_2.setName("example1_2");

        // second child: transaction controller
        TransactionController ex2 = new TransactionController();
        ex2.setName("ex2");
        DebugSampler example2_1 = new DebugSampler();
        example2_1.setName("example2_1");
        DebugSampler example2_2 = new DebugSampler();
        example2_2.setName("example2_2");

        // main loop
        LoopController loop = new LoopController();
        loop.setLoops(4);
        loop.setContinueForever(false);

        // test tree
        ListedHashTree hashTree = new ListedHashTree();
        hashTree.add(loop);
        hashTree.add(loop, topWSC);
        hashTree.add(topWSC, listener);
        hashTree.add(topWSC, simple1);
        hashTree.add(simple1, ex1);
        hashTree.add(ex1, example1_1);
        hashTree.add(ex1, example1_2);
        hashTree.add(ex1, listener);
        hashTree.add(topWSC, simple2);
        hashTree.add(simple2, ex2);
        hashTree.add(ex2, example2_1);
        hashTree.add(ex2, example2_2);
        hashTree.add(ex2, listener);

        TestCompiler compiler = new TestCompiler(hashTree);
        hashTree.traverse(compiler);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setNumThreads(1);

        ListenerNotifier notifier = new ListenerNotifier();

        JMeterThread thread = new JMeterThread(hashTree, threadGroup, notifier);
        thread.setThreadGroup(threadGroup);
        thread.setOnErrorStopThread(true);
        thread.run();

        Map<String, Integer> totalResults = new HashMap<>();
        for (SampleEvent event : listener.events) {
            String label = event.getResult().getSampleLabel();
            if (totalResults.containsKey(label)) {
                totalResults.put(label, totalResults.get(label) + 1);
            } else {
                totalResults.put(label, 1);
            }
        }

        assertEquals(12, listener.events.size());
        assertEquals(2, (int) totalResults.get("example1_1"));
        assertEquals(2, (int) totalResults.get("example1_2"));
        assertEquals(2, (int) totalResults.get("example2_1"));
        assertEquals(2, (int) totalResults.get("example2_2"));
        assertEquals(2, (int) totalResults.get("ex1")); // transaction result
        assertEquals(2, (int) totalResults.get("ex2")); // transaction result
    }

    @Test
    public void testLoopControllerWithNestedDisableController() throws Exception {
        JMeterContextService.getContext().setVariables(new JMeterVariables());
        TestSampleListener listener = new TestSampleListener();

        URL resource = WeightedSwitchControllerTest.class.getResource("/disableController.jmx");
        File testPlan = new File(resource.getFile());

        HashTree hashTree = SaveService.loadTree(testPlan);

        hashTree.add(hashTree.getArray()[0], listener);
        JMeter.convertSubTree(hashTree);
        StandardJMeterEngine engine = new StandardJMeterEngine();
        engine.configure(hashTree);
        engine.run();

        Map<String, Integer> totalResults = new HashMap<>();
        for (SampleEvent event : listener.events) {
            String label = event.getResult().getSampleLabel();
            if (totalResults.containsKey(label)) {
                totalResults.put(label, totalResults.get(label) + 1);
            } else {
                totalResults.put(label, 1);
            }
        }

        assertEquals(100, listener.events.size());
        assertEquals(50, (int) totalResults.get("Debug 1"));
        assertEquals(50, (int) totalResults.get("Debug 3"));
        assertNull(totalResults.get("Debug 2"));
    }

    public class TestSampleListener extends ResultCollector implements SampleListener {
        public List<SampleEvent> events = new ArrayList<>();

        @Override
        public void sampleOccurred(SampleEvent e) {
            events.add(e);
        }

        @Override
        public void sampleStarted(SampleEvent e) {
            events.add(e);
        }

        @Override
        public void sampleStopped(SampleEvent e) {
            events.add(e);
        }
    }

    @Test
    public void testRandomChoice() {
        LoggingManager.setPriority("INFO");
        try {
            final int iterations = 1000000;
            WeightedSwitchController obj = new WeightedSwitchController();
            obj.setIsRandomChoice(true);
            obj.addTestElement(getSampler("0"));
            obj.addTestElement(getSampler("1"));
            PowerTableModel mdl = new PowerTableModel(new String[]{"name", WeightedSwitchController.WEIGHTS}, new Class[]{String.class, String.class});
            mdl.addRow(new String[]{"0", "95"});
            mdl.addRow(new String[]{"1", "5"});
            obj.setData(mdl);

            int sampler0Counter = 0;
            int sampler1Counter = 0;
            for (int n = 0; n < iterations; n++) {
                Sampler sampler = obj.next();
                assertNotNull(sampler);
                String name = sampler.getName();
                if ("0".equals(name)) {
                    sampler0Counter++;
                } else {
                    sampler1Counter++;
                }
                assertNull(obj.next());
            }

            double sampler0Percent = ((double) sampler0Counter / iterations) * 100;
            String msg0 = "Sampler 0 was ran " + sampler0Counter + " times. Total percent: " + sampler0Percent;
            assertTrue(msg0, 94.5 < sampler0Percent);
            assertTrue(msg0, 95.5 > sampler0Percent);

            double sampler1Percent = ((double) sampler1Counter / iterations) * 100;
            String msg1 = "Sampler 1 was ran " + sampler1Counter + " times. Total percent: " + sampler1Percent;
            assertTrue(msg1, 4.5 < sampler1Percent);
            assertTrue(msg1, 5.5 > sampler1Percent);
        } finally {
            LoggingManager.setPriority("DEBUG");
        }
    }
}