package eu.gity.jmeter.ddos.stairsddosthreadgrp;

import eu.gity.jmeter.ddos.common.Trafgen;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.engine.TreeCloner;
import org.apache.jmeter.threads.*;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.collections.ListedHashTree;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * A superclass for DDoS Thread Groups which uses advanced load scheduling
 */

public abstract class AbstractDDoSThreadGroup extends AbstractThreadGroup {
    private static final long serialVersionUID = 240L;

    private static final Logger log = LoggerFactory.getLogger(AbstractDDoSThreadGroup.class);
    private static final long WAIT_TO_DIE = JMeterUtils.getPropDefault("jmeterengine.threadstop.wait", 5 * 1000); // 5 seconds

    // List of active threads
    private final Map<JMeterThread, Thread> allThreads = new ConcurrentHashMap<JMeterThread, Thread>();

    /**
     * No-arg constructor.
     */
    public AbstractDDoSThreadGroup() {
    }

    /**
     * Serves for purpose of generation chart on GUI
     *
     * @param thread
     * @param now
     */
    protected abstract void scheduleThread(JMeterThread thread, long now);

    /**
     * Generates a list of setting change points in a form of time-change pair
     *
     * @return a list of setting changes
     */
    protected abstract ArrayList<SettingChange> generateSettingChanges();

    @Override
    public void start(int groupCount, ListenerNotifier notifier, ListedHashTree threadGroupTree, StandardJMeterEngine engine) {

        log.info("Starting thread group number " + groupCount
                + " threads " + 1);
        long now = System.currentTimeMillis(); // needs to be same time for all threads in the group
        final JMeterContext context = JMeterContextService.getContext();

        JMeterThread jmThread = makeThread(groupCount, notifier, threadGroupTree, engine, 1, context);

        jmThread.setStartTime(now);

        long endTime = scheduleSettingChanges(generateSettingChanges(), threadGroupTree) + now;

        jmThread.setEndTime(endTime);
        jmThread.setScheduled(true);
        Thread newThread = new Thread(jmThread, jmThread.getThreadName());
        registerStartedThread(jmThread, newThread);
        newThread.start();

        log.info("Started thread group number " + groupCount);
    }

    /**
     * @param settingChanges
     * @return endTime relative to test start
     */
    private long scheduleSettingChanges(ArrayList<SettingChange> settingChanges, ListedHashTree threadGroupTree) {
        for (SettingChange sc : settingChanges) {
            final ScheduledExecutorService service = Executors.newSingleThreadScheduledExecutor();
            service.schedule(new Runnable() {
                @Override
                public void run() {
                    Object[] tgElements = threadGroupTree.getArray(threadGroupTree.getArray()[0]);
                    for (Object tge : tgElements) {
                        if (tge instanceof Trafgen) {
                            ((Trafgen) tge).changeRate(sc.ppsChange);
                        }
                    }
                    service.shutdown();
                }
            }, sc.time, TimeUnit.MILLISECONDS);
        }
        return settingChanges.get(settingChanges.size() - 1).time;
    }

    private void registerStartedThread(JMeterThread jMeterThread, Thread newThread) {
        allThreads.put(jMeterThread, newThread);
    }

    private JMeterThread makeThread(int groupCount,
                                    ListenerNotifier notifier, ListedHashTree threadGroupTree,
                                    StandardJMeterEngine engine, int i,
                                    JMeterContext context) { // N.B. Context needs to be fetched in the correct thread
        boolean onErrorStopTest = getOnErrorStopTest();
        boolean onErrorStopTestNow = getOnErrorStopTestNow();
        boolean onErrorStopThread = getOnErrorStopThread();
        boolean onErrorStartNextLoop = getOnErrorStartNextLoop();
        String groupName = getName();
        final JMeterThread jmeterThread = new JMeterThread(cloneTree(threadGroupTree), this, notifier);
        jmeterThread.setThreadNum(i);
        jmeterThread.setThreadGroup(this);
        jmeterThread.setInitialContext(context);
        final String threadName = groupName + " " + (groupCount) + "-" + (i + 1);
        jmeterThread.setThreadName(threadName);
        jmeterThread.setEngine(engine);
        jmeterThread.setOnErrorStopTest(onErrorStopTest);
        jmeterThread.setOnErrorStopTestNow(onErrorStopTestNow);
        jmeterThread.setOnErrorStopThread(onErrorStopThread);
        jmeterThread.setOnErrorStartNextLoop(onErrorStartNextLoop);
        return jmeterThread;
    }

    @Override
    public boolean stopThread(String threadName, boolean now) {
        for (Entry<JMeterThread, Thread> entry : allThreads.entrySet()) {
            JMeterThread thrd = entry.getKey();
            if (thrd.getThreadName().equals(threadName)) {
                thrd.stop();
                thrd.interrupt();
                if (now) {
                    Thread t = entry.getValue();
                    if (t != null) {
                        t.interrupt();
                    }
                }
                return true;
            }
        }
        return false;
    }

    @Override
    public void threadFinished(JMeterThread thread) {
        log.debug("Ending thread " + thread.getThreadName());
        allThreads.remove(thread);
    }

    @Override
    public void tellThreadsToStop() {
        for (Entry<JMeterThread, Thread> entry : allThreads.entrySet()) {
            JMeterThread item = entry.getKey();
            item.stop(); // set stop flag
            item.interrupt(); // interrupt sampler if possible
            Thread t = entry.getValue();
            if (t != null) { // Bug 49734
                t.interrupt(); // also interrupt JVM thread
            }
        }
    }

    @Override
    public void stop() {
        for (JMeterThread item : allThreads.keySet()) {
            item.stop();
        }
    }

    @Override
    public int numberOfActiveThreads() {
        return allThreads.size();
    }

    @Override
    public boolean verifyThreadsStopped() {
        boolean stoppedAll = true;
        for (Thread t : allThreads.values()) {
            stoppedAll = stoppedAll && verifyThreadStopped(t);
        }
        return stoppedAll;
    }

    private boolean verifyThreadStopped(Thread thread) {
        boolean stopped = true;
        if (thread != null) {
            if (thread.isAlive()) {
                try {
                    thread.join(WAIT_TO_DIE);
                } catch (InterruptedException e) {
                    log.warn("Interrupted while waiting for thread {} to finish", thread.getName());
                }
                if (thread.isAlive()) {
                    stopped = false;
                    log.warn("Thread won't exit: " + thread.getName());
                }
            }
        }
        return stopped;
    }

    @Override
    public void waitThreadsStopped() {
        for (Thread t : allThreads.values()) {
            waitThreadStopped(t);
        }
    }

    private void waitThreadStopped(Thread thread) {
        if (thread != null) {
            while (thread.isAlive()) {
                try {
                    thread.join(WAIT_TO_DIE);
                } catch (InterruptedException e) {
                    log.warn("Interrupted while waiting for thread {} to finish", thread.getName());
                }
            }
        }
    }

    private ListedHashTree cloneTree(ListedHashTree tree) {
        TreeCloner cloner = new TreeCloner(true);
        tree.traverse(cloner);
        return cloner.getClonedTree();
    }
}

