package eu.gity.jmeter.ddos.common;

import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.gui.GuiPackage;
import org.apache.jmeter.gui.action.ActionNames;
import org.apache.jmeter.gui.action.ActionRouter;
import org.apache.jmeter.gui.action.Start;
import org.apache.jmeter.gui.tree.JMeterTreeModel;
import org.apache.jmeter.gui.tree.JMeterTreeNode;
import org.apache.jmeter.testelement.TestElement;
import org.apache.jmeter.util.JMeterUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.swing.*;
import javax.swing.event.TreeModelEvent;
import javax.swing.event.TreeModelListener;
import javax.swing.event.TreeSelectionEvent;
import javax.swing.event.TreeSelectionListener;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static eu.gity.jmeter.ddos.properties.strings.DdosPropertiesStrings.*;

/**
 * Control class for trafgen generator which is used to configure, start and stop trafgen
 *
 * @author Gity a.s.
 */
public class TrafgenControl {

    public static final String configRelativePath = "/lib/ext/ddos/";
    private static final Logger log = LoggerFactory.getLogger(TrafgenControl.class);
    private static final String processName = "trafgen";
    private TrafgenThread trafgenThread;
    private String timestamp;
    private String configFile;
    private String configFileSuffix;
    private TestElement te;
    private final DDoSUtils ddosUtils = new DDoSUtils();

    private JMeterTreeNode previouslySelectedNode;
    private boolean nodePreviouslyEnabled = true;

    private String correctedPacketCount; //Packet count after correction

    //Synced to GUI to disable rate and number setting if rate will be set dynamically by thread group 
    private boolean dynamicRate;

    private int newRate; //Used if dynamicRate is true
    private int oldRate; //Used if dynamicRate is true

    /**
     * A simple ctor for special DDoS plugins like Trafgen Custom Config
     */
    public TrafgenControl() {
    }

    /**
     * Ctor which performs initialization tasks including checking of abstract config file and action router listeners registration
     *
     * @param configFileSuffix name of the configuration file for particular DDoS sampler which uses this class, including extension excluding path
     * @param testElement      Test Element object of the sampler
     */
    public TrafgenControl(String configFileSuffix, TestElement testElement) {
        this.configFileSuffix = configFileSuffix;
        this.te = testElement;
        checkAbstractConfig();
        addActionRouterListeners();
    }

    /**
     * Allows a plugin to set a timestamp for config file (for special DDoS plugins like Trafgen Custom Config)
     */
    public void setTimestamp(String timestamp) {
        this.timestamp = timestamp;
    }

    /**
     * Verifies if abstract configuration file for DDoS sampler which instantiated this class contains only valid property references
     *
     * @return true if abstract config contains only valid property references
     */
    private boolean checkAbstractConfig() {
        if (te == null) { //Instantiated by DDoS which does not have any abstract config file
            return true;
        }
        String jmeterHomeDir = ddosUtils.getJmeterHomeDir();
        String abstractConfigFile = jmeterHomeDir + configRelativePath + "Abstract" + configFileSuffix;

        List<String> lines = null;
        try {
            lines = new ArrayList<String>(Files.readAllLines(Paths.get(abstractConfigFile)));
        } catch (IOException e) {
            JOptionPane.showMessageDialog(null, JMeterUtils.getLocaleString(DDOS_ERROR_ABSTRACT_CONFIG_FILE_1OF2) + " " +
                            abstractConfigFile + "\n" + JMeterUtils.getLocaleString(DDOS_ERROR_ABSTRACT_CONFIG_FILE_2OF2),
                    JMeterUtils.getLocaleString(DDOS_ERROR), JOptionPane.ERROR_MESSAGE);
            return false;
        }

        Pattern pattern = Pattern.compile("##([A-Za-z0-9]+\\.[A-Za-z0-9]+)");
        for (String line : lines) {
            Matcher matcher = pattern.matcher(line);
            while (matcher.find()) {
                if (te.getPropertyAsString(matcher.group(1)).equals("")) {
                    JOptionPane.showMessageDialog(null, JMeterUtils.getLocaleString(DDOS_ERROR_ABSTRACT_CONFIG_VARIABLE_1OF2) + " " +
                                    te.getName() + " : " + matcher.group(1) + "\n" + JMeterUtils.getLocaleString(DDOS_ERROR_ABSTRACT_CONFIG_VARIABLE_2OF2),
                            JMeterUtils.getLocaleString(DDOS_ERROR), JOptionPane.ERROR_MESSAGE);
                    return false;
                }
            }
        }
        return true;
    }

    /**
     * Generates configuration file based on properties acquired from Test Element. It sets configFile field to either path to successfuly
     * created config file or null if error occurred which will then indicate trafgen should not be launched
     */
    public void generateConfig() {
        if (!checkAbstractConfig()) {
            log.error("Error occured while checking Trafgen abstract config file.");
            this.configFile = null;
        }

        String jmeterHomeDir = ddosUtils.getJmeterHomeDir();
        String abstractConfigFile = jmeterHomeDir + configRelativePath + "Abstract" + configFileSuffix;

        DateFormat dateFormat = new SimpleDateFormat("yyyy_MM_dd_HH_mm_ss");
        Date date = new Date();

        String nanos = Long.toString(System.nanoTime());
        timestamp = dateFormat.format(date) + "_" + nanos.substring(nanos.length() - 6);

        String configFile = jmeterHomeDir + configRelativePath + timestamp + configFileSuffix;

        List<String> lines = null;
        try {
            lines = new ArrayList<String>(Files.readAllLines(Paths.get(abstractConfigFile)));
        } catch (IOException e) {
            log.error("Error while reading abstract configuration from file " + abstractConfigFile + ":\n" + e.getStackTrace().toString());
            this.configFile = null;
        }

        Pattern pattern = Pattern.compile("##([A-Za-z0-9]+\\.[A-Za-z0-9]+)");
        List<String> newLines = new ArrayList<String>();
        for (String line : lines) {
            Matcher matcher = pattern.matcher(line);
            while (matcher.find()) {
                line = line.replace("##" + matcher.group(1), te.getPropertyAsString(matcher.group(1)));
            }
            newLines.add(line);
        }

        try {
            Files.write(Paths.get(configFile), newLines, StandardOpenOption.CREATE_NEW);
        } catch (IOException e) {
            log.error("Error while writting configuration to file " + configFile + ":\n" + e.getStackTrace().toString());
            this.configFile = null;
        }

        this.configFile = configFile;
    }

    /**
     * Verifies if packet count will be enough for the test duration
     *
     * @param tgName   name of the DDoS Thread Group in which the DDoS sampler resides
     * @param duration duration of the test calculated by DDoS Thread Group scheduler
     * @param rate     packet rate set in the sampler
     * @param number   packet count set in the sampler
     */
    public void verifyPacketCount(String tgName, long duration, long rate, long number) {
        long samplerMaxDuration = number / rate;
        long threadGroupMaxCount = duration / 1000 * rate;
        if (duration / 1000 > samplerMaxDuration) {
            Object[] options = {JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_OPTION_YES),
                    JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_OPTION_NO),
                    JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_OPTION_STOP_TEST)};

            int answer = JOptionPane.showOptionDialog(GuiPackage.getInstance().getMainFrame(),
                    te.getName() + " " + JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_1OF6) + " " + number + " " +
                            JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_2OF6) + " " + rate + " pps = " + samplerMaxDuration + " " +
                            JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_3OF6) + "\n" +
                            tgName + " " + JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_4OF6) + " " + duration / 1000 + "\n" +
                            JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_5OF6) + " " + duration / 1000 + " * " + rate + " = " + threadGroupMaxCount + " " +
                            JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_6OF6) + " " + duration / 1000 + " sec?",
                    JMeterUtils.getLocaleString(DDOS_LOW_PACKET_COUNT_TITLE), JOptionPane.YES_NO_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE,
                    null, options, options[0]);

            if (answer == JOptionPane.YES_OPTION) {
                correctedPacketCount = String.valueOf(threadGroupMaxCount);
            } else if (answer == JOptionPane.CANCEL_OPTION) {
                StandardJMeterEngine.stopEngine();
            }
        }
    }

    /**
     * Starts trafgen process in a separate thread or restarts it if dynamic rate setting is being used and the setting has changed since last run
     *
     * @return -2 when no rate change was required, -1 on error, 0 when rate is set to 0 (trafgen is terminated) or any other number which corresponds to a new rate
     */
    public int runTrafgen() {
        if (configFile == null) {
            log.error("Error occured while creating Trafgen config file. Aborting!");
            return -1;
        }

        String netInterface = te.getPropertyAsString(te.getClass().getSimpleName() + ".interface");
        String number = null;
        String rate = null;

        if (!dynamicRate) {
            if (trafgenThread != null) { //Already running
                if (trafgenThread.isAlive()) {
                    return -2; //No action required
                } else {
                    return -3; //Stop the JMeter thread
                }
            }
            String packetCount = correctedPacketCount == null ? te.getPropertyAsString(te.getClass().getSimpleName() + ".number") : correctedPacketCount;
            if (packetCount.isEmpty()) {
                number = "";
            } else {
                number = "--num " + packetCount;
            }

            String packetRate = te.getPropertyAsString(te.getClass().getSimpleName() + ".rate");
            if (packetRate.isEmpty()) {
                rate = "";
            } else {
                rate = " --rate " + packetRate + "pps";
            }
        } else {
            if (newRate != oldRate) {
                oldRate = newRate;
                terminateTrafgen(false);
                if (newRate == 0) {
                    return 0;
                }
                number = "";
                rate = " --rate " + newRate + "pps";
            } else {
                return -2; //Rate has not changed so trafgen doesn't have to be restarted
            }
        }
        String command = "pkexec trafgen --in " + configFile + " --out "
                + netInterface + " --cpp " + number + rate;

        String[] args = new String[]{"/bin/bash", "-c", command};

        trafgenThread = new TrafgenThread(args);
        trafgenThread.start();

        return newRate;
    }

    /**
     * Searches for running trafgen processes for this Trafgen Control instance
     *
     * @param timestamp config file timestamp which identifies config file and also trafgen process of particular TrafgenControl instance
     * @return a space delimited String containing PIDs of trafgen processes belonging to this Trafgen Control instance
     */
    private String getSpecificTrafgenPIDs(String timestamp) {
        StringBuilder PIDs = new StringBuilder();
        try {
            String line;
            Process p = Runtime.getRuntime().exec("ps -eo pid,command");
            BufferedReader input = new BufferedReader(new InputStreamReader(p.getInputStream()));

            while ((line = input.readLine()) != null) {
                if (line.contains(processName) && line.contains(timestamp)) {
                    line = line.substring(0, 6).trim();
                    PIDs.append(line + " ");
                }
            }
            input.close();

        } catch (Exception e) {
            log.error("Error occured while searching for trafgen process PIDs:\n" + e.getStackTrace().toString() + "\n Killing all trafgens");
            terminateTrafgen(true);
            return "";
        }

        return PIDs.toString();
    }

    /**
     * Terminates trafgen processes for a specific sampler instance (i.e. when DDoS Thread Group is terminated by scheduler)
     * or all processes (i.e. when "Stop Test" button is pressed)
     *
     * @param killall determines if all processes regardless of Thread Group should be terminated
     */
    public void terminateTrafgen(boolean killall) {
        String PIDs;
        String termCommand;
        String killCommand;
        if (killall) {
            termCommand = "killall -SIGTERM " + processName;
            killCommand = "killall -SIGKILL " + processName;
        } else {
            PIDs = getSpecificTrafgenPIDs(timestamp);
            if (PIDs.equals("")) {
                return; //Because of error getting PIDs, all trafgens were already killed so no need to do anything
            }
            termCommand = "pkexec kill -SIGTERM " + PIDs;
            killCommand = "pkexec kill -SIGKILL " + PIDs;
        }
        try {
            //Tries to terminate specific/all trafgen processes
            String[] termArgs = new String[]{"/bin/bash", "-c", termCommand};
            Process termPr = new ProcessBuilder(termArgs).start();
            termPr.waitFor();
            if (termPr.exitValue() != 0) {
                //Tries to kill specific/all trafgen processes if termination did not succeeded
                String[] killArgs = new String[]{"/bin/bash", "-c", killCommand};
                Process killPr = new ProcessBuilder(killArgs).start();
                killPr.waitFor();
                if (killPr.exitValue() != 0) {
                    //Tries to kill all trafgen processes if killing specific processes did not succeeded (i.e. wrong PIDs)
                    String[] killAllArgs = new String[]{"/bin/bash", "-c", "killall -SIGKILL " + processName};
                    Process killAllPr = new ProcessBuilder(killAllArgs).start();
                    killAllPr.waitFor();
                }
            }
        } catch (IOException e) {
            log.error("Error occured while terminating trafgen processes:\n" + e.getStackTrace().toString());
        } catch (InterruptedException e) {
            log.error("Termination of trafgen processes was interrupted:\n" + e.getStackTrace().toString());
        }
    }

    /**
     * Attaches necessary listeners to JMeter ActionRouter to initialize checking of DDoS plugin's abstract configuration file
     */
    void addActionRouterListeners() {
        GuiPackage gp = GuiPackage.getInstance();

        if (gp == null) { //Non-GUI mode
            return;
        }
        //Reacts on selection event
        gp.getMainFrame().getTree().addTreeSelectionListener(new TreeSelectionListener() {

            @Override
            public void valueChanged(TreeSelectionEvent e) {
                JMeterTreeNode currentNode = (JMeterTreeNode) e.getPath().getLastPathComponent();
                if (currentNode.getUserObject() == te && !currentNode.equals(previouslySelectedNode) && currentNode.isEnabled()) {
                    checkAbstractConfig();
                }
                previouslySelectedNode = currentNode;
            }
        });

        //Reacts on enable event
        JMeterTreeModel jmeterTreeModel = (JMeterTreeModel) gp.getMainFrame().getTree().getModel();
        jmeterTreeModel.addTreeModelListener(new TreeModelListener() {

            @Override
            public void treeNodesChanged(TreeModelEvent e) {
                if (e.getSource() instanceof JMeterTreeModel) {
                    JMeterTreeNode currentNode = (JMeterTreeNode) e.getChildren()[0];
                    if (currentNode.getUserObject().equals(te)) {
                        if (currentNode.isEnabled() && !nodePreviouslyEnabled) {
                            checkAbstractConfig();
                        }
                        nodePreviouslyEnabled = currentNode.isEnabled();
                    }
                }
            }

            @Override
            public void treeNodesInserted(TreeModelEvent e) {
            }

            @Override
            public void treeNodesRemoved(TreeModelEvent e) {
            }

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

    /**
     * Provides read access to thread which runs trafgen process to allow checking of it's status
     *
     * @return TrafgenThread object
     */
    public TrafgenThread getTrafgenThread() {
        return trafgenThread;
    }

    /**
     * Provides write access to thread which runs trafgen process to allow external Thread instantiation for plugins which do not
     * use the internal runTrafgen() method
     *
     * @return TrafgenThread object
     */
    public void setTrafgenThread(TrafgenThread trafgenThread) {
        this.trafgenThread = trafgenThread;
    }

    /**
     * Provides read access to correctedPacketCount for plugins which do not use the internal runTrafgen() method
     *
     * @return correctedPacketCount
     */
    public String getCorrectedPacketCount() {
        return correctedPacketCount;
    }

    /**
     * Provides read access to oldRate for plugins which do not use the internal runTrafgen() method
     *
     * @return oldRate
     */
    public int getOldRate() {
        return oldRate;
    }

    /**
     * Provides write access to oldRate for plugins which do not use the internal runTrafgen() method
     *
     * @return oldRate
     */
    public void setOldRate(int oldRate) {
        this.oldRate = oldRate;
    }

    /**
     * Provides read access to newRate for plugins which do not use the internal runTrafgen() method
     *
     * @return newRate
     */
    public int getNewRate() {
        return newRate;
    }

    /**
     * Provides write access to configFile for plugins which do not use the internal runTrafgen() method
     *
     * @return configFile
     */
    public String getConfigFile() {
        return configFile;
    }

    /**
     * Provides write access to configFile for plugins which do not use the internal runTrafgen() method
     *
     * @param configFile
     */
    public void setConfigFile(String configFile) {
        this.configFile = configFile;
    }

    /**
     * @return true if rate will be changed dynamically during the test
     */
    public boolean isDynamicRate() {
        return dynamicRate;
    }

    /**
     * @param dynamicRate if rate will be changed dynamically during the test
     */
    public void setDynamicRate(boolean dynamicRate) {
        this.dynamicRate = dynamicRate;
    }

    /**
     * Allow DDoS thread groups to change packet rate dynamically during the test
     *
     * @param rateChange either positive or negative change in pps rate
     */
    public void changeRate(int rateChange) {
        newRate += rateChange;
    }

    /**
     * Resets some instance variables for next test run
     */
    public void reset() {
        oldRate = 0;
        newRate = 0;
        trafgenThread = null;
        if (configFile != null) {
            try {
                Files.delete(Paths.get(configFile));
            } catch (IOException e1) {
                log.error("Error occured while deleting trafgen configuration file " + configFile + ":\n" + e1.getStackTrace().toString());
            }
            configFile = null;
        }
    }

    /**
     * A class which implements thread for running trafgen process
     *
     * @author Gity a.s.
     */
    public class TrafgenThread extends Thread {
        private final String[] args;
        private boolean result;

        /**
         * @param args arguments for process creation containing mainly command for starting trafgen including parameters
         */
        public TrafgenThread(String[] args) {
            this.args = args;
        }

        /**
         * Allows access to result of the trafgen run
         *
         * @return true if there were no errors
         */
        public boolean getResult() {
            return result;
        }

        /**
         * {@inheritDoc}
         */
        public void run() {
            ActionRouter.getInstance().addPostActionListener(Start.class, new ActionListener() {

                @Override
                public void actionPerformed(ActionEvent e) {
                    if (e.getActionCommand().equals(ActionNames.ACTION_STOP) || e.getActionCommand().equals(ActionNames.ACTION_SHUTDOWN)) {
                        terminateTrafgen(true);
                        if (configFile != null) {
                            try {
                                Files.delete(Paths.get(configFile));
                            } catch (IOException e1) {
                                log.error("Error occured while deleting trafgen configuration file " + configFile + ":\n" + e1.getStackTrace().toString());
                            }
                            configFile = null;
                        }
                    }
                }
            });
            try {
                Runtime rt = Runtime.getRuntime();
                Process process;
                process = rt.exec(args);

                //Redirects errors from stdout to JMeter log
                BufferedReader errorStream = new BufferedReader(new InputStreamReader(process.getErrorStream()));
                ArrayList<String> errorList = new ArrayList<String>();
                String line = null;
                while ((line = errorStream.readLine()) != null) {
                    errorList.add(line);
                }
                if (errorList.size() > 0) {
                    log.error("Trafgen error: \n" + errorList);
                }
                process.waitFor();
				result = process.exitValue() == 0;
            } catch (IOException e) {
                log.error("Error occured while generating traffic:\n" + e.getStackTrace().toString());
                result = false;
            } catch (InterruptedException e) {
                log.warn("Trafgen was interrupted");
                result = true;
            }
        }
    }
}
