package org.apache.jmeter;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileFilter;
import java.io.FileInputStream;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.Thread.UncaughtExceptionHandler;
import java.security.Permission;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.regex.Pattern;

import org.apache.commons.io.IOUtils;
import org.apache.log4j.Logger;
import org.apache.maven.plugin.AbstractMojo;
import org.apache.maven.plugin.MojoExecutionException;
import org.apache.maven.plugin.MojoFailureException;
import org.apache.tools.ant.DirectoryScanner;

/**
 * JMeter Maven plugin.
 * 
 * @author Tim McCune
 * @goal jmeter
 * @requiresDependencyResolution test
 */
public class JMeterMojo extends AbstractMojo {

    private static final Pattern PAT_ERROR = Pattern
        .compile(".*\\s+ERROR\\s+.*");

    /**
     * @parameter
     */
    private List<String> includes;

    /**
     * @parameter
     */
    private List<String> excludes;

    /**
     * @parameter expression="${basedir}/src/test/jmeter"
     */
    private File srcDir;

    /**
     * @parameter expression="jmeter-reports"
     */
    private File reportDir;

    /**
     * @parameter expression="${basedir}/src/test/jmeter/jmeter.properties"
     */
    private File jmeterProps;

    /**
     * @parameter expression="${settings.localRepository}"
     */
    @SuppressWarnings("unused")
    private File repoDir;

    /**
     * JMeter Properties to be overridden
     * 
     * @parameter
     */
    private Map<String, ?> jmeterUserProperties;

    /**
     * @parameter
     */
    private boolean remote;

    /**
     * The dir where the compiled classes are placed
     * 
     * @parameter expression="${jmeter.compiledClasses}"
     *            default-value="${basedir}/target/classes"
     */
    private File compiledClasses;

    /**
     * The dir where the compiled test classes are placed
     * 
     * @parameter expression="${jmeter.compiledTestClasses}"
     *            default-value="${basedir}/target/test-classes"
     */
    private File compiledTestClasses;

    /**
     * The file where the classpath is written
     * 
     * @parameter expression="${jmeter.classpathDump}"
     *            default-value="${basedir}/target/classpath"
     */
    private File classpathDump;

    /**
     * The file where the logs are created
     * 
     * @parameter expression="${jmeter.workDir}"
     *            default-value="${basedir}/target/jmeter"
     */
    private File workDir;

    /**
     * Sets whether ErrorScanner should ignore failures in JMeter result file.
     * 
     * @parameter expression="${jmeter.ignore.failure}" default-value=false
     */
    private boolean jmeterIgnoreFailure;

    private final Logger logger = Logger.getLogger(JMeterMojo.class);

    private File saveServiceProps;

    private File upgradeProps;

    private File jmeterLog;

    private static final String BASE_DIR_PATH = System
        .getProperty("user.dir");

    private static final String JAVA_CLASS_PATH = "java.class.path";

    private static final String DOT_JAR = ".jar";

    /**
     * Run all JMeter tests.
     */
    public void execute()
            throws MojoExecutionException, MojoFailureException {
        addDependenciesToClasspath();
        initSystemProps();

        try {
            DirectoryScanner scanner = new DirectoryScanner();
            scanner.setBasedir(srcDir);
            scanner
                .setIncludes(includes == null ? new String[] {"**/*.jmx" }
                        : includes.toArray(new String[] {}));
            if (excludes != null) {
                scanner.setExcludes(excludes.toArray(new String[] {}));
            }
            scanner.scan();
            for (String file : scanner.getIncludedFiles()) {
                executeTest(new File(srcDir, file));
            }
            checkForErrors();
        } finally {
            saveServiceProps.delete();
            upgradeProps.delete();
        }
    }

    private void checkForErrors()
            throws MojoExecutionException, MojoFailureException {
        try {
            BufferedReader in =
                new BufferedReader(new FileReader(jmeterLog));
            String line;
            while ((line = in.readLine()) != null) {
                if (PAT_ERROR.matcher(line).find()) {

                    if (jmeterIgnoreFailure) {
                        getLog()
                            .warn(
                                "There were test errors, continuing because jmeterIgnoreFailure = 'true', see logfile '"
                                    + jmeterLog
                                    + "' for further information.");
                        break;
                    } else {
                        throw new MojoFailureException(
                            "There were test errors, see logfile '"
                                + jmeterLog + "' for further information");
                    }
                }
            }
            in.close();
        } catch (IOException e) {
            throw new MojoExecutionException("Can't read log file", e);
        }
    }

    private void initSystemProps() throws MojoExecutionException {
        workDir.mkdirs();
        createSaveServiceProps();
        jmeterLog = new File(workDir, "jmeter.log");
        try {
            System.setProperty("log_file", jmeterLog.getCanonicalPath());
        } catch (IOException e) {
            throw new MojoExecutionException(
                "Can't get canonical path for log file", e);
        }
    }

    /**
     * This mess is necessary because JMeter must load this info from a file. Do
     * the same for the upgrade.properties Resources won't work.
     */
    private void createSaveServiceProps() throws MojoExecutionException {
        saveServiceProps = new File(workDir, "saveservice.properties");
        upgradeProps = new File(workDir, "upgrade.properties");
        try {
            FileWriter out = new FileWriter(saveServiceProps);
            IOUtils.copy(Thread.currentThread().getContextClassLoader()
                .getResourceAsStream("saveservice.properties"), out);
            out.flush();
            out.close();
            System.setProperty("saveservice_properties", File.separator
                + "target" + File.separator + "jmeter" + File.separator
                + "saveservice.properties");

            out = new FileWriter(upgradeProps);

            IOUtils.copy(Thread.currentThread().getContextClassLoader()
                .getResourceAsStream("upgrade.properties"), out);

            out.flush();
            out.close();
            System.setProperty("upgrade_properties", File.separator
                + "target" + File.separator + "jmeter" + File.separator
                + "upgrade.properties");

        } catch (IOException e) {
            throw new MojoExecutionException(
                "Could not create temporary saveservice.properties", e);
        }
    }

    /**
     * Executes a single JMeter test by building up a list of command line
     * parameters to pass to JMeter.start().
     */
    private void executeTest(final File test)
            throws MojoExecutionException {
        try {
            getLog().info("Executing test: " + test.getCanonicalPath());
            String reportFileName =
                test.getName().substring(0,
                    test.getName().lastIndexOf("."))
                    + ".csv";

            List<String> argsTmp =
                Arrays
                    .asList("-n", "-t", test.getCanonicalPath(), "-l",
                        reportDir.toString() + File.separator
                            + reportFileName, "-p",
                        jmeterProps.toString(), "-d", BASE_DIR_PATH);

            List<String> args = new ArrayList<String>();
            args.addAll(argsTmp);
            args.addAll(getUserProperties());

            if (remote) {
                args.add("-r");
            }

            // This mess is necessary because JMeter likes to use System.exit.
            // We need to trap the exit call.
            SecurityManager oldManager = System.getSecurityManager();
            System.setSecurityManager(new SecurityManager() {
                @Override
                public void checkExit(final int status) {
                    throw new ExitException(status);
                }

                @Override
                public void checkPermission(final Permission perm,
                        final Object context) {
                }

                @Override
                public void checkPermission(final Permission perm) {
                }
            });
            UncaughtExceptionHandler oldHandler =
                Thread.getDefaultUncaughtExceptionHandler();
            Thread
                .setDefaultUncaughtExceptionHandler(new UncaughtExceptionHandler() {
                    public void uncaughtException(final Thread t,
                            final Throwable e) {
                        if (e instanceof ExitException
                            && ((ExitException) e).getCode() == 0) {
                            return; // Ignore
                        }
                        getLog().error("Error in thread " + t.getName());
                    }
                });
            try {
                // This mess is necessary because the only way to know when
                // JMeter
                // is done is to wait for all of the threads that it spawned to
                // exit.
                new JMeter().start(args.toArray(new String[] {}));
                BufferedReader in =
                    new BufferedReader(new FileReader(jmeterLog));
                while (!checkForEndOfTest(in)) {
                    try {
                        Thread.sleep(10000);
                    } catch (InterruptedException e) {
                        break;
                    }
                }
                in.close();
            } catch (ExitException e) {
                if (e.getCode() != 0) {
                    throw new MojoExecutionException("Test failed", e);
                }
            } finally {
                System.setSecurityManager(oldManager);
                Thread.setDefaultUncaughtExceptionHandler(oldHandler);
            }
        } catch (IOException e) {
            throw new MojoExecutionException("Can't execute test", e);
        }
    }

    private ArrayList<String> getUserProperties() {
        ArrayList<String> propsList = new ArrayList<String>();
        if (jmeterUserProperties == null) {
            return propsList;
        }
        Set<String> keySet = jmeterUserProperties.keySet();

        for (String key : keySet) {

            propsList.add("-J");
            propsList.add(key + "=" + jmeterUserProperties.get(key));
        }

        return propsList;
    }

    private static class ExitException extends SecurityException {
        private static final long serialVersionUID = 5544099211927987521L;

        public int _rc;

        public ExitException(final int rc) {
            super(Integer.toString(rc));
            _rc = rc;
        }

        public int getCode() {
            return _rc;
        }
    }

    private void addDependenciesToClasspath() {

        final String initialClassptah =
            System.getProperty(JAVA_CLASS_PATH, "");

        // Load jmeter driver class and add it to java class path property
        try {
            File jmeterDriverJar =
                new File(BASE_DIR_PATH + File.separator + "bin"
                    + File.separator + "ApacheJMeter.jar");
            if (!jmeterDriverJar.exists()) {
                logger.warn("Cannot found jmeter driver jar: "
                    + jmeterDriverJar);
            } else {
                logger.debug("Adding depedency: " + jmeterDriverJar);
                ClassPathHacker.addFile(jmeterDriverJar);
                System.setProperty(JAVA_CLASS_PATH,
                    jmeterDriverJar.getPath());
            }
        } catch (IOException e) {
            logger.warn("An error occured reading jmeter driver jar file",
                e);
        }

        // Load jmeter libraries
        try {
            File[] jarDirs =
                {
                        new File(BASE_DIR_PATH + File.separator + "lib"),
                        new File(BASE_DIR_PATH + File.separator + "lib"
                            + File.separator + "ext"),
                        new File(BASE_DIR_PATH + File.separator + "lib"
                            + File.separator + "junit") };
            List<File> jars = getJarsFromDirs(jarDirs);

            for (File jar : jars) {
                ClassPathHacker.addFile(jar);
                addToClasspath(jar.getAbsolutePath());
            }

        } catch (IOException e) {
            logger.warn("An error occured reading jmeter jar modules", e);
        }

        try {

            if (compiledClasses != null) {
                logger.debug("Adding depedency: " + compiledClasses);
                ClassPathHacker.addFile(compiledClasses);
            } else {
                logger.warn("No compiled classes directory specified");
            }

            if (compiledTestClasses != null) {
                logger.debug("Adding depedency: " + compiledTestClasses);
                ClassPathHacker.addFile(compiledTestClasses);
            } else {
                logger
                    .warn("No compiled test classes directory specified");
            }

            String[] dependencies = getClasspathFiles();

            String[] javaClassPath =
                System.getProperty(JAVA_CLASS_PATH, "").split(
                    File.pathSeparator);

            if (dependencies != null) {
                for (String dependency : dependencies) {
                    if (!dependency.contains("maven-")
                        && jarNotIn(dependency, javaClassPath)) {
                        logger.debug("Adding depedency: " + dependency);
                        ClassPathHacker.addFile(dependency);
                    } else {
                        logger.warn("Avoiding depedency: " + dependency);
                    }
                }
            }

        } catch (IOException e) {
            logger.warn(
                "An error occured reading the classpath dump file", e);
        }

        // Adding initial java class path at the end
        if (!initialClassptah.isEmpty()) {
            System.setProperty(JAVA_CLASS_PATH,
                System.getProperty(JAVA_CLASS_PATH) + File.pathSeparator
                    + initialClassptah);
        }
    }

    private boolean jarNotIn(final String jarPath, final String[] paths) {
        String jarName = new File(jarPath).getName();

        for (String p : paths) {
            if (p.endsWith(jarName)) {
                logger.warn(jarName + " is already loaded");
                return false;
            }
        }
        return true;
    }

    private void addToClasspath(final String path) {
        StringBuilder classpath =
            new StringBuilder(System.getProperty(JAVA_CLASS_PATH, ""));
        if (!classpath.toString().isEmpty()
            && !classpath.toString().endsWith(File.pathSeparator)) {
            classpath.append(File.pathSeparator);
        }
        classpath.append(path);
        System.setProperty(JAVA_CLASS_PATH, classpath.toString());
    }

    private List<File> getJarsFromDirs(final File[] dirs) {
        List<File> ret = new LinkedList<File>();
        if (dirs != null) {
            for (File dir : dirs) {
                if (!dir.exists()) {
                    logger.warn(dir + " directory does not exist");
                } else if (!dir.isDirectory()) {
                    logger.warn(dir + " is not a directory");
                } else if (!dir.canRead()) {
                    logger.warn(dir + " cannot read");
                } else {
                    File[] jars = dir.listFiles(new FileFilter() {
                        public boolean accept(final File pathname) {
                            return pathname.getName().endsWith(DOT_JAR);
                        }
                    });

                    if (jars.length == 0) {
                        logger.warn(dir
                            + " directory does not contain jar files");
                    } else {
                        for (File jar : jars) {
                            ret.add(jar);
                        }
                    }
                }
            }
        }
        return ret;
    }

    private String[] getClasspathFiles() throws IOException {
        if (classpathDump == null) {
            logger.warn("The classpath dump file has not been set");

            return null;
        }

        StringBuffer fileData = new StringBuffer(1000);
        BufferedReader reader =
            new BufferedReader(new InputStreamReader(new FileInputStream(
                classpathDump)));

        char[] buf = new char[1024];
        int numRead = 0;
        String readData;
        while ((numRead = reader.read(buf)) != -1) {
            readData = String.valueOf(buf, 0, numRead);
            fileData.append(readData);
            buf = new char[1024];
        }
        reader.close();

        String rtn[] =
            fileData.toString()
                .split(System.getProperty("path.separator"));
        logger.debug("Dependencies found : " + rtn.length);

        return rtn;
    }

    private boolean checkForEndOfTest(final BufferedReader in)
            throws MojoExecutionException {
        boolean testEnded = false;
        try {
            String line;
            while ((line = in.readLine()) != null) {
                if (line.indexOf("Test has ended") != -1) {
                    testEnded = true;
                    break;
                }
            }
        } catch (IOException e) {
            throw new MojoExecutionException("Can't read log file", e);
        }
        return testEnded;
    }

    public boolean isJmeterIgnoreFailure() {
        return jmeterIgnoreFailure;
    }

    public void setJmeterIgnoreFailure(final boolean jmeterIgnoreFailure) {
        this.jmeterIgnoreFailure = jmeterIgnoreFailure;
    }

}
