package eu.gity.jmeter.ddos.simpleddosthreadgrp;

import eu.gity.jmeter.ddos.common.DDoSThreadGroup;
import eu.gity.jmeter.ddos.common.Trafgen;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.engine.TreeCloner;
import org.apache.jmeter.gui.GuiPackage;
import org.apache.jmeter.gui.action.ActionRouter;
import org.apache.jmeter.gui.action.EditCommand;
import org.apache.jmeter.gui.tree.JMeterTreeListener;
import org.apache.jmeter.gui.tree.JMeterTreeNode;
import org.apache.jmeter.testelement.TestElement;
import org.apache.jmeter.testelement.property.BooleanProperty;
import org.apache.jmeter.testelement.property.LongProperty;
import org.apache.jmeter.threads.*;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.collections.ListedHashTree;
import org.apache.jorphan.util.JMeterStopTestException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.swing.event.TreeModelEvent;
import javax.swing.event.TreeModelListener;
import javax.swing.tree.TreeNode;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Enumeration;
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;

/**
 * Simple DDoS Thread group test element class which is modification of classic Thread Group shaped to needs of DDoS attacks samplers
 * - modified start() method
 * - added checkDDoSPacketCount() and stopDDoSSamplers() *
 */

public class SimpleDDoSThreadGroup extends AbstractThreadGroup implements DDoSThreadGroup {
    /**
     * Whether scheduler is being used
     */
    public static final String SCHEDULER = "ThreadGroup.scheduler";
    /**
     * Scheduler absolute start time
     */
    public static final String START_TIME = "ThreadGroup.start_time";
    /**
     * Scheduler absolute end time
     */
    public static final String END_TIME = "ThreadGroup.end_time";
    /**
     * Scheduler duration, overrides end time
     */
    public static final String DURATION = "ThreadGroup.duration";
    /**
     * Scheduler start delay, overrides start time
     */
    public static final String DELAY = "ThreadGroup.delay";

    //+ JMX entries - do not change the string values
    private static final long serialVersionUID = 280L;
    private static final Logger log = LoggerFactory.getLogger(SimpleDDoSThreadGroup.class);
    private static final String DATE_FIELD_FORMAT = "yyyy/MM/dd HH:mm:ss"; //$NON-NLS-1$
    private static final long WAIT_TO_DIE = JMeterUtils.getPropDefault("jmeterengine.threadstop.wait", 5 * 1000); // 5 seconds
    /**
     * How often to check for shutdown during ramp-up, default 1000ms
     */
    private static final int RAMPUP_GRANULARITY = JMeterUtils.getPropDefault("jmeterthread.rampup.granularity", 1000); // $NON-NLS-1$
    // List of active threads
    private final Map<JMeterThread, Thread> allThreads = new ConcurrentHashMap<>();

    /**
     * Is test (still) running?
     */
    private volatile boolean running = false;

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

    /**
     * Get whether scheduler is being used
     *
     * @return true if scheduler is being used
     */
    public boolean getScheduler() {
        return getPropertyAsBoolean(SCHEDULER);
    }

    /**
     * Set whether scheduler is being used
     *
     * @param Scheduler true is scheduler is to be used
     */
    public void setScheduler(boolean Scheduler) {
        setProperty(new BooleanProperty(SCHEDULER, Scheduler));
    }

    /**
     * Get the absolute start time value.
     *
     * @return the start time value.
     */
    public long getStartTime() {
        return getPropertyAsLong(START_TIME);
    }

    /**
     * Set the absolute StartTime value.
     *
     * @param stime the StartTime value.
     */
    public void setStartTime(long stime) {
        setProperty(new LongProperty(START_TIME, stime));
    }

    /**
     * Get the desired duration of the thread group test run
     *
     * @return the duration (in secs)
     */
    public long getDuration() {
        return getPropertyAsLong(DURATION);
    }

    /**
     * Set the desired duration of the thread group test run
     *
     * @param duration in seconds
     */
    public void setDuration(long duration) {
        setProperty(new LongProperty(DURATION, duration));
    }

    /**
     * Get the startup delay
     *
     * @return the delay (in secs)
     */
    public long getDelay() {
        return getPropertyAsLong(DELAY);
    }

    /**
     * Set the startup delay
     *
     * @param delay in secondsstopDDoSSamplers(jmThread)
     */
    public void setDelay(long delay) {
        setProperty(new LongProperty(DELAY, delay));
    }

    /**
     * Get the end time value.
     *
     * @return the end time value.
     */
    public long getEndTime() {
        return getPropertyAsLong(END_TIME);
    }

    /**
     * Set the EndTime value.
     *
     * @param etime the EndTime value.
     */
    public void setEndTime(long etime) {
        setProperty(new LongProperty(END_TIME, etime));
    }

    /**
     * This will schedule the time for the JMeterThread.
     *
     * @param thread JMeterThread
     */
    private void scheduleThread(JMeterThread thread, long now) {
        // if true the Scheduler is enabled
        if (getScheduler()) {
            // set the start time for the Thread
            if (getDelay() > 0) {// Duration is in seconds
                thread.setStartTime(getDelay() * 1000 + now);
            } else {
                long start = getStartTime();
                if (start < now) {
                    start = now; // Force a sensible start time
                }
                thread.setStartTime(start);
            }
            // set the endtime for the Thread
            if (getDuration() > 0) {// Duration is in seconds
                thread.setEndTime(getDuration() * 1000 + (thread.getStartTime()));
            } else {
                if (getEndTime() <= now) {
                    SimpleDateFormat sdf = new SimpleDateFormat(DATE_FIELD_FORMAT);
                    throw new JMeterStopTestException("End Time (" + sdf.format(new Date(getEndTime())) + ") of Scheduler for Thread Group " + getName() + " is in the past, fix value of End Time field");
                }
                StandardJMeterEngine.stopEngine();
                thread.setEndTime(getEndTime());
            }
            // Enables the scheduler
            thread.setScheduled(true);
        }
    }

    /**
     * Wait for delay with RAMPUP_GRANULARITY
     *
     * @param delay delay in ms
     */
    private void delayBy(long delay) {
        if (delay > 0) {
            long start = System.currentTimeMillis();
            long end = start + delay;
            long now = 0;
            long pause = RAMPUP_GRANULARITY; // maximum pause to use
            while (running && (now = System.currentTimeMillis()) < end) {
                long togo = end - now;
                if (togo < pause) {
                    pause = togo;
                }
                pause(pause); // delay between checks
            }
        }
    }

    @Override
    public void start(int groupCount, ListenerNotifier notifier, ListedHashTree threadGroupTree, StandardJMeterEngine engine) {
        running = true;
        log.info("Starting thread group number " + groupCount);
        long now = System.currentTimeMillis();
        final JMeterContext context = JMeterContextService.getContext();
        final boolean usingScheduler = getScheduler();

        JMeterThread jmThread = makeThread(groupCount, notifier, threadGroupTree, engine, 1, context);
        scheduleThread(jmThread, now); // set start and end time
        Thread newThread = new Thread(jmThread, jmThread.getThreadName());
        registerStartedThread(jmThread, newThread);

        if (usingScheduler) {
            long duration = jmThread.getEndTime() - jmThread.getStartTime();
            checkDDoSPacketCount(jmThread, duration);
            long delay = jmThread.getStartTime() - now;
            if (delay > 0) {
                delayBy(delay);
            }
            //Ensures stopping of DDoS samplers when time is out
            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).stopTrafgen();
                        }
                    }
                    service.shutdown();
                }
            }, duration, TimeUnit.MILLISECONDS);
        }
        newThread.start();
        log.info("Started thread group number " + groupCount);
    }

    /**
     * // else start immediately
     * }
     * Register Thread when it starts
     *
     * @param jMeterThread {@link JMeterThread}
     * @param newThread    Thread
     */
    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;
    }

    /**
     * Stop thread called threadName:
     * <ol>
     *  <li>stop JMeter thread</li>
     *  <li>interrupt JMeter thread</li>
     *  <li>interrupt underlying thread</li>
     * </ol>
     *
     * @param threadName String thread name
     * @param now        boolean for stop
     * @return true if thread stopped
     */
    @Override
    public boolean stopThread(String threadName, boolean now) {
        for (Entry<JMeterThread, Thread> entry : allThreads.entrySet()) {
            JMeterThread thrd = entry.getKey();
            if (thrd.getThreadName().equals(threadName)) {
                stopThread(thrd, entry.getValue(), now);
                return true;
            }
        }
        return false;
    }

    /**
     * @param thrd      JMeterThread
     * @param t         Thread
     * @param interrupt Interrup thread or not
     */
    private void stopThread(JMeterThread thrd, Thread t, boolean interrupt) {
        thrd.stop();
        thrd.interrupt(); // interrupt sampler if possible
        if (interrupt) {
            if (t != null) { // Bug 49734
                t.interrupt(); // also interrupt JVM thread
            }
        }
    }

    /**
     * Verifies if packet count set in each DDoS sampler will be enough for the test duration
     *
     * @param thrd     thread from which a list of samplers will be extracted
     * @param duration planned test duration
     */
    private void checkDDoSPacketCount(JMeterThread thrd, long duration) {
        //There's always only a single thread in the DDoS Thread Group
        SimpleDDoSThreadGroup ddostg = (SimpleDDoSThreadGroup) thrd.getTestTree().getArray()[0];

        //Gets the children (only DDoS samplers are allowed to reside inside og DDoS Thread Group)
        Object[] ddosSamplers = thrd.getTestTree().getArray(ddostg);

        for (Object sampler : ddosSamplers) {
            if (sampler instanceof Trafgen) {
                ((Trafgen) sampler).checkPacketCount(getName(), duration);
            }
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void broadcastSamplerRateStrategy() {
        SimpleDDoSThreadGroup tg = this;

        //Broadcast every time a DDoS node is added
        GuiPackage.getInstance().getTreeModel().addTreeModelListener(new TreeModelListener() {
            @Override
            public void treeNodesChanged(TreeModelEvent e) {
            }

            @Override
            public void treeNodesInserted(TreeModelEvent e) {
                Object[] children = e.getChildren();
                for (Object child : children) {
                    if (((JMeterTreeNode) child).getUserObject() instanceof Trafgen) {
                        //If parent of this DDoS node is this object...
                        if (((JMeterTreeNode) ((JMeterTreeNode) child).getParent()).getUserObject().equals(tg)) {
                            JMeterTreeNode tgNode = GuiPackage.getInstance().getNodeOf(tg);
                            if (tgNode != null) {
                                @SuppressWarnings("unchecked") Enumeration<TreeNode> kids = tgNode.children();
                                while (kids.hasMoreElements()) {
                                    JMeterTreeNode childTN = (JMeterTreeNode) kids.nextElement();
                                    TestElement childTE = childTN.getTestElement();
                                    if (childTE instanceof Trafgen) {
                                        ((Trafgen) childTE).setDynamicRate(false);
                                    }
                                }
                            }
                        }
                    }
                }
            }

            @Override
            public void treeNodesRemoved(TreeModelEvent e) {
            }

            @Override
            public void treeStructureChanged(TreeModelEvent e) {
            }
        });

        //Broadcast every time on "EditCommand" event
        ActionRouter.getInstance().addPreActionListener(EditCommand.class, new ActionListener() {
            //TODO move to DDoSUtils??
            @Override
            public void actionPerformed(ActionEvent e) {
                if (e.getSource() instanceof JMeterTreeListener) {
                    JMeterTreeNode currentNode = ((JMeterTreeListener) e.getSource()).getCurrentNode();
                    JMeterTreeNode parentNode = ((JMeterTreeNode) currentNode.getParent());
                    if (parentNode != null) {
                        //If my own Trafgen child sampler is being edited
                        if (currentNode.getTestElement() instanceof Trafgen && parentNode.getTestElement().equals(tg)) {
                            @SuppressWarnings("unchecked") Enumeration<TreeNode> children = parentNode.children();
                            while (children.hasMoreElements()) {
                                JMeterTreeNode childTN = (JMeterTreeNode) children.nextElement();
                                TestElement childTE = childTN.getTestElement();
                                if (childTE instanceof Trafgen) {
                                    ((Trafgen) childTE).setDynamicRate(false);
                                }
                            }
                        }
                    }
                }
            }
        });
    }

    /**
     * Called by JMeterThread when it finishes
     */
    @Override
    public void threadFinished(JMeterThread thread) {
        log.debug("Ending thread " + thread.getThreadName());
        allThreads.remove(thread);
    }

    /**
     * For each thread, invoke:
     * <ul>
     * <li>{@link JMeterThread#stop()} - set stop flag</li>
     * <li>{@link JMeterThread#interrupt()} - interrupt sampler</li>
     * <li>{@link Thread#interrupt()} - interrupt JVM thread</li>
     * </ul>
     */
    @Override
    public void tellThreadsToStop() {
        running = false;
        for (Entry<JMeterThread, Thread> entry : allThreads.entrySet()) {
            stopThread(entry.getKey(), entry.getValue(), true);
        }
    }

    /**
     * For each thread, invoke:
     * <ul>
     * <li>{@link JMeterThread#stop()} - set stop flag</li>
     * </ul>
     */
    @Override
    public void stop() {
        running = false;
        for (JMeterThread item : allThreads.keySet()) {
            item.stop();
        }
    }

    /**
     * @return number of active threads
     */
    @Override
    public int numberOfActiveThreads() {
        return allThreads.size();
    }

    /**
     * @return boolean true if all threads stopped
     */
    @Override
    public boolean verifyThreadsStopped() {
        boolean stoppedAll = true;
        for (Thread t : allThreads.values()) {
            stoppedAll = stoppedAll && verifyThreadStopped(t);
        }
        return stoppedAll;
    }

    /**
     * Verify thread stopped and return true if stopped successfully
     *
     * @param thread Thread
     * @return boolean
     */
    private boolean verifyThreadStopped(Thread thread) {
        boolean stopped = true;
        if (thread != null) {
            if (thread.isAlive()) {
                try {
                    thread.join(WAIT_TO_DIE);
                } catch (InterruptedException e) {
                }
                if (thread.isAlive()) {
                    stopped = false;
                    log.warn("Thread won't exit: " + thread.getName());
                }
            }
        }
        return stopped;
    }

    /**
     * Wait for all Group Threads to stop
     */
    @Override
    public void waitThreadsStopped() {
        for (Thread t : allThreads.values()) {
            waitThreadStopped(t);
        }
    }

    /**
     * Wait for thread to stop
     *
     * @param thread Thread
     */
    private void waitThreadStopped(Thread thread) {
        if (thread != null) {
            while (thread.isAlive()) {
                try {
                    thread.join(WAIT_TO_DIE);
                } catch (InterruptedException e) {
                }
            }
        }
    }

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

    private void pause(long ms) {
        try {
            TimeUnit.MILLISECONDS.sleep(ms);
        } catch (InterruptedException e) {
            // TODO Is this silent exception intended
        }
    }

    @Override
    public JMeterThread addNewThread(int delay, StandardJMeterEngine engine) {
        return null;
    }
}
