/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.maven.surefire.junitcore.pc;

import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumMap;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;

import org.apache.maven.plugin.surefire.log.api.ConsoleLogger;
import org.apache.maven.surefire.api.testset.TestSetFailedException;
import org.apache.maven.surefire.api.util.internal.DaemonThreadFactory;
import org.apache.maven.surefire.junitcore.JUnitCoreParameters;
import org.junit.internal.runners.ErrorReportingRunner;
import org.junit.runner.Description;
import org.junit.runner.Runner;
import org.junit.runner.manipulation.Filter;
import org.junit.runner.manipulation.NoTestsRemainException;
import org.junit.runner.notification.RunNotifier;
import org.junit.runners.ParentRunner;
import org.junit.runners.Suite;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.RunnerBuilder;

import static org.apache.maven.surefire.junitcore.pc.ParallelComputerUtil.resolveConcurrency;
import static org.apache.maven.surefire.junitcore.pc.SchedulingStrategies.createParallelStrategy;
import static org.apache.maven.surefire.junitcore.pc.SchedulingStrategies.createParallelStrategyUnbounded;
import static org.apache.maven.surefire.junitcore.pc.Type.CLASSES;
import static org.apache.maven.surefire.junitcore.pc.Type.METHODS;
import static org.apache.maven.surefire.junitcore.pc.Type.SUITES;

@SuppressWarnings({"javadoc", "checkstyle:javadoctype"})
/*
 * Executing suites, classes and methods with defined concurrency. In this example the threads which completed
 * the suites and classes can be reused in parallel methods.
 * <pre>
 * JUnitCoreParameters parameters = ...;
 * ParallelComputerBuilder builder = new ParallelComputerBuilder(parameters);
 * builder.useOnePool(8).parallelSuites(2).parallelClasses(4).parallelMethods();
 * ParallelComputerBuilder.ParallelComputer computer = builder.buildComputer();
 * Class<?>[] tests = {...};
 * new JUnitCore().run(computer, tests);
 * </pre>
 * Note that the type has always at least one thread even if unspecified. The capacity in
 * {@link ParallelComputerBuilder#useOnePool(int)} must be greater than the number of concurrent suites and classes
 * altogether.
 * <br>
 * The Computer can be stopped in a separate thread. Pending tests will be interrupted if the argument is
 * {@code true}.
 * <pre>
 * computer.describeStopped(true);
 * </pre>
 *
 * @author Tibor Digana (tibor17)
 * @since 2.16
 */
public final class ParallelComputerBuilder {
    private static final ThreadFactory DAEMON_THREAD_FACTORY = DaemonThreadFactory.newDaemonThreadFactory();

    private static final Class<? extends Annotation> JCIP_NOT_THREAD_SAFE = loadNotThreadSafeAnnotations();

    private static final Set<Runner> NULL_SINGLETON = Collections.singleton(null);

    static final int TOTAL_POOL_SIZE_UNDEFINED = 0;

    private final Map<Type, Integer> parallelGroups = new EnumMap<>(Type.class);

    private final ConsoleLogger logger;

    private boolean useSeparatePools;

    private int totalPoolSize;

    private JUnitCoreParameters parameters;

    private boolean optimize;

    private boolean runningInTests;

    /**
     * Calling {@link #useSeparatePools()}.
     * Can be used only in unit tests.
     * Do NOT call this constructor in production.
     */
    ParallelComputerBuilder(ConsoleLogger logger) {
        this.logger = logger;
        runningInTests = true;
        useSeparatePools();
        parallelGroups.put(SUITES, 0);
        parallelGroups.put(CLASSES, 0);
        parallelGroups.put(METHODS, 0);
    }

    public ParallelComputerBuilder(ConsoleLogger logger, JUnitCoreParameters parameters) {
        this(logger);
        runningInTests = false;
        this.parameters = parameters;
    }

    public ParallelComputer buildComputer() {
        return new PC();
    }

    ParallelComputerBuilder useSeparatePools() {
        totalPoolSize = TOTAL_POOL_SIZE_UNDEFINED;
        useSeparatePools = true;
        return this;
    }

    ParallelComputerBuilder useOnePool() {
        totalPoolSize = TOTAL_POOL_SIZE_UNDEFINED;
        useSeparatePools = false;
        return this;
    }

    /**
     * @param totalPoolSize Pool size where suites, classes and methods are executed in parallel.
     *                      If the <code>totalPoolSize</code> is {@link Integer#MAX_VALUE}, the pool capacity is not
     *                      limited.
     * @throws IllegalArgumentException If <code>totalPoolSize</code> is &lt; 1.
     */
    ParallelComputerBuilder useOnePool(int totalPoolSize) {
        if (totalPoolSize < 1) {
            throw new IllegalArgumentException("Size of common pool is less than 1.");
        }
        this.totalPoolSize = totalPoolSize;
        useSeparatePools = false;
        return this;
    }

    boolean isOptimized() {
        return optimize;
    }

    ParallelComputerBuilder optimize(boolean optimize) {
        this.optimize = optimize;
        return this;
    }

    ParallelComputerBuilder parallelSuites() {
        return parallel(SUITES);
    }

    ParallelComputerBuilder parallelSuites(int nThreads) {
        return parallel(nThreads, SUITES);
    }

    ParallelComputerBuilder parallelClasses() {
        return parallel(CLASSES);
    }

    ParallelComputerBuilder parallelClasses(int nThreads) {
        return parallel(nThreads, CLASSES);
    }

    ParallelComputerBuilder parallelMethods() {
        return parallel(METHODS);
    }

    ParallelComputerBuilder parallelMethods(int nThreads) {
        return parallel(nThreads, METHODS);
    }

    private ParallelComputerBuilder parallel(int nThreads, Type parallelType) {
        if (nThreads < 0) {
            throw new IllegalArgumentException("negative nThreads " + nThreads);
        }

        if (parallelType == null) {
            throw new IllegalArgumentException("null parallelType");
        }

        parallelGroups.put(parallelType, nThreads);
        return this;
    }

    private ParallelComputerBuilder parallel(Type parallelType) {
        return parallel(Integer.MAX_VALUE, parallelType);
    }

    private double parallelTestsTimeoutInSeconds() {
        return parameters == null ? 0d : parameters.getParallelTestsTimeoutInSeconds();
    }

    private double parallelTestsTimeoutForcedInSeconds() {
        return parameters == null ? 0d : parameters.getParallelTestsTimeoutForcedInSeconds();
    }

    @SuppressWarnings("unchecked")
    private static Class<? extends Annotation> loadNotThreadSafeAnnotations() {
        try {
            Class c = Class.forName("net.jcip.annotations.NotThreadSafe");
            return c.isAnnotation() ? (Class<? extends Annotation>) c : null;
        } catch (ClassNotFoundException e) {
            return null;
        }
    }

    final class PC extends ParallelComputer {
        private final SingleThreadScheduler notThreadSafeTests =
                new SingleThreadScheduler(ParallelComputerBuilder.this.logger);

        private final Collection<ParentRunner> suites = new LinkedHashSet<>();

        private final Collection<ParentRunner> nestedSuites = new LinkedHashSet<>();

        private final Collection<ParentRunner> classes = new LinkedHashSet<>();

        private final Collection<ParentRunner> nestedClasses = new LinkedHashSet<>();

        private final Collection<Runner> notParallelRunners = new LinkedHashSet<>();

        private int poolCapacity;

        private boolean splitPool;

        private final Map<Type, Integer> allGroups;

        private long nestedClassesChildren;

        private volatile Scheduler master;

        private PC() {
            super(parallelTestsTimeoutInSeconds(), parallelTestsTimeoutForcedInSeconds());
            allGroups = new EnumMap<>(ParallelComputerBuilder.this.parallelGroups);
            poolCapacity = ParallelComputerBuilder.this.totalPoolSize;
            splitPool = ParallelComputerBuilder.this.useSeparatePools;
        }

        Collection<ParentRunner> getSuites() {
            return suites;
        }

        Collection<ParentRunner> getNestedSuites() {
            return nestedSuites;
        }

        Collection<ParentRunner> getClasses() {
            return classes;
        }

        Collection<ParentRunner> getNestedClasses() {
            return nestedClasses;
        }

        Collection<Runner> getNotParallelRunners() {
            return notParallelRunners;
        }

        int getPoolCapacity() {
            return poolCapacity;
        }

        boolean isSplitPool() {
            return splitPool;
        }

        @Override
        protected ShutdownResult describeStopped(boolean shutdownNow) {
            ShutdownResult shutdownResult = notThreadSafeTests.describeStopped(shutdownNow);
            final Scheduler m = master;
            if (m != null) {
                ShutdownResult shutdownResultOfMaster = m.describeStopped(shutdownNow);
                shutdownResult.getTriggeredTests().addAll(shutdownResultOfMaster.getTriggeredTests());
                shutdownResult.getIncompleteTests().addAll(shutdownResultOfMaster.getIncompleteTests());
            }
            return shutdownResult;
        }

        @Override
        protected boolean shutdownThreadPoolsAwaitingKilled() {
            boolean notInterrupted = notThreadSafeTests.shutdownThreadPoolsAwaitingKilled();
            final Scheduler m = master;
            if (m != null) {
                notInterrupted &= m.shutdownThreadPoolsAwaitingKilled();
            }
            return notInterrupted;
        }

        @Override
        public Runner getSuite(RunnerBuilder builder, Class<?>[] cls) throws InitializationError {
            try {
                super.getSuite(builder, cls);
                populateChildrenFromSuites();

                WrappedRunners suiteSuites = wrapRunners(suites);
                WrappedRunners suiteClasses = wrapRunners(classes);

                long suitesCount = suites.size();
                long classesCount = classes.size() + nestedClasses.size();
                long methodsCount = suiteClasses.embeddedChildrenCount + nestedClassesChildren;
                if (!ParallelComputerBuilder.this.runningInTests) {
                    determineThreadCounts(suitesCount, classesCount, methodsCount);
                }

                return setSchedulers(suiteSuites.wrappingSuite, suiteClasses.wrappingSuite);
            } catch (TestSetFailedException e) {
                throw new InitializationError(Collections.<Throwable>singletonList(e));
            }
        }

        @Override
        protected Runner getRunner(RunnerBuilder builder, Class<?> testClass) throws Throwable {
            Runner runner = super.getRunner(builder, testClass);
            if (canSchedule(runner)) {
                if (!isThreadSafe(runner)) {
                    ((ParentRunner) runner).setScheduler(notThreadSafeTests.newRunnerScheduler());
                    notParallelRunners.add(runner);
                } else if (runner instanceof Suite) {
                    suites.add((Suite) runner);
                } else {
                    classes.add((ParentRunner) runner);
                }
            } else {
                notParallelRunners.add(runner);
            }
            return runner;
        }

        private void determineThreadCounts(long suites, long classes, long methods) throws TestSetFailedException {
            RunnerCounter counts =
                    ParallelComputerBuilder.this.optimize ? new RunnerCounter(suites, classes, methods) : null;
            Concurrency concurrency = resolveConcurrency(ParallelComputerBuilder.this.parameters, counts);
            allGroups.put(SUITES, concurrency.suites);
            allGroups.put(CLASSES, concurrency.classes);
            allGroups.put(METHODS, concurrency.methods);
            poolCapacity = concurrency.capacity;
            splitPool &= concurrency.capacity <= 0; // fault if negative; should not happen
        }

        private <T extends Runner> WrappedRunners wrapRunners(Collection<T> runners) throws InitializationError {
            // Do NOT use allGroups here.
            long childrenCounter = 0;
            ArrayList<Runner> runs = new ArrayList<>();
            for (T runner : runners) {
                if (runner != null) {
                    int children = countChildren(runner);
                    childrenCounter += children;
                    runs.add(runner);
                }
            }
            return runs.isEmpty() ? new WrappedRunners() : new WrappedRunners(createSuite(runs), childrenCounter);
        }

        private int countChildren(Runner runner) {
            Description description = runner.getDescription();
            Collection children = description == null ? null : description.getChildren();
            return children == null ? 0 : children.size();
        }

        private ExecutorService createPool(int poolSize) {
            return poolSize < Integer.MAX_VALUE
                    ? Executors.newFixedThreadPool(poolSize, DAEMON_THREAD_FACTORY)
                    : Executors.newCachedThreadPool(DAEMON_THREAD_FACTORY);
        }

        private Scheduler createMaster(ExecutorService pool, int poolSize) {
            // can be 0, 1, 2 or 3
            final int finalRunnersCounter = countFinalRunners();

            final SchedulingStrategy strategy;
            if (finalRunnersCounter <= 1 || poolSize <= 1) {
                strategy = new InvokerStrategy(ParallelComputerBuilder.this.logger);
            } else if (pool != null && poolSize == Integer.MAX_VALUE) {
                strategy = new SharedThreadPoolStrategy(ParallelComputerBuilder.this.logger, pool);
            } else {
                strategy = createParallelStrategy(ParallelComputerBuilder.this.logger, finalRunnersCounter);
            }
            return new Scheduler(ParallelComputerBuilder.this.logger, null, strategy);
        }

        private int countFinalRunners() {
            int counter = notParallelRunners.isEmpty() ? 0 : 1;

            if (!suites.isEmpty() && allGroups.get(SUITES) > 0) {
                ++counter;
            }

            if (!classes.isEmpty() && allGroups.get(CLASSES) > 0) {
                ++counter;
            }

            return counter;
        }

        private void populateChildrenFromSuites() {
            // Do NOT use allGroups here.
            Filter filter = new SuiteFilter();
            for (Iterator<ParentRunner> it = suites.iterator(); it.hasNext(); ) {
                ParentRunner suite = it.next();
                try {
                    suite.filter(filter);
                } catch (NoTestsRemainException e) {
                    it.remove();
                }
            }
        }

        private int totalPoolSize() {
            if (poolCapacity == TOTAL_POOL_SIZE_UNDEFINED) {
                int total = 0;
                for (int nThreads : allGroups.values()) {
                    total += nThreads;
                    if (total < 0) {
                        total = Integer.MAX_VALUE;
                        break;
                    }
                }
                return total;
            } else {
                return poolCapacity;
            }
        }

        private Runner setSchedulers(ParentRunner suiteSuites, ParentRunner suiteClasses) throws InitializationError {
            int parallelSuites = allGroups.get(SUITES);
            int parallelClasses = allGroups.get(CLASSES);
            int parallelMethods = allGroups.get(METHODS);
            int poolSize = totalPoolSize();
            ExecutorService commonPool = splitPool || poolSize == 0 ? null : createPool(poolSize);
            master = createMaster(commonPool, poolSize);

            if (suiteSuites != null) {
                // a scheduler for parallel suites
                if (commonPool != null && parallelSuites > 0) {
                    Balancer balancer = BalancerFactory.createBalancerWithFairness(parallelSuites);
                    suiteSuites.setScheduler(createScheduler(null, commonPool, true, balancer));
                } else {
                    suiteSuites.setScheduler(createScheduler(parallelSuites));
                }
            }

            // schedulers for parallel classes
            ArrayList<ParentRunner> allSuites = new ArrayList<>(suites);
            allSuites.addAll(nestedSuites);
            if (suiteClasses != null) {
                allSuites.add(suiteClasses);
            }
            if (!allSuites.isEmpty()) {
                setSchedulers(allSuites, parallelClasses, commonPool);
            }

            // schedulers for parallel methods
            ArrayList<ParentRunner> allClasses = new ArrayList<>(classes);
            allClasses.addAll(nestedClasses);
            if (!allClasses.isEmpty()) {
                setSchedulers(allClasses, parallelMethods, commonPool);
            }

            // resulting runner for Computer#getSuite() scheduled by master scheduler
            ParentRunner all = createFinalRunner(removeNullRunners(
                    Arrays.<Runner>asList(suiteSuites, suiteClasses, createSuite(notParallelRunners))));
            all.setScheduler(master);
            return all;
        }

        private ParentRunner createFinalRunner(List<Runner> runners) throws InitializationError {
            return new Suite(null, runners) {
                @Override
                public void run(RunNotifier notifier) {
                    try {
                        beforeRunQuietly();
                        super.run(notifier);
                    } finally {
                        afterRunQuietly();
                    }
                }
            };
        }

        private void setSchedulers(Iterable<? extends ParentRunner> runners, int poolSize, ExecutorService commonPool) {
            if (commonPool != null) {
                Balancer concurrencyLimit = BalancerFactory.createBalancerWithFairness(poolSize);
                boolean doParallel = poolSize > 0;
                for (ParentRunner runner : runners) {
                    runner.setScheduler(
                            createScheduler(runner.getDescription(), commonPool, doParallel, concurrencyLimit));
                }
            } else {
                ExecutorService pool = null;
                if (poolSize == Integer.MAX_VALUE) {
                    pool = Executors.newCachedThreadPool(DAEMON_THREAD_FACTORY);
                } else if (poolSize > 0) {
                    pool = Executors.newFixedThreadPool(poolSize, DAEMON_THREAD_FACTORY);
                }
                boolean doParallel = pool != null;
                for (ParentRunner runner : runners) {
                    runner.setScheduler(createScheduler(
                            runner.getDescription(),
                            pool,
                            doParallel,
                            BalancerFactory.createInfinitePermitsBalancer()));
                }
            }
        }

        private Scheduler createScheduler(
                Description desc, ExecutorService pool, boolean doParallel, Balancer concurrency) {
            SchedulingStrategy strategy = doParallel & pool != null
                    ? new SharedThreadPoolStrategy(ParallelComputerBuilder.this.logger, pool)
                    : new InvokerStrategy(ParallelComputerBuilder.this.logger);
            return new Scheduler(ParallelComputerBuilder.this.logger, desc, master, strategy, concurrency);
        }

        private Scheduler createScheduler(int poolSize) {
            final SchedulingStrategy strategy;
            if (poolSize == Integer.MAX_VALUE) {
                strategy = createParallelStrategyUnbounded(ParallelComputerBuilder.this.logger);
            } else if (poolSize == 0) {
                strategy = new InvokerStrategy(ParallelComputerBuilder.this.logger);
            } else {
                strategy = createParallelStrategy(ParallelComputerBuilder.this.logger, poolSize);
            }
            return new Scheduler(ParallelComputerBuilder.this.logger, null, master, strategy);
        }

        private boolean canSchedule(Runner runner) {
            return !(runner instanceof ErrorReportingRunner) && runner instanceof ParentRunner;
        }

        private boolean isThreadSafe(Runner runner) {
            return runner.getDescription().getAnnotation(JCIP_NOT_THREAD_SAFE) == null;
        }

        private class SuiteFilter extends Filter {
            // Do NOT use allGroups in SuiteFilter.

            @Override
            public boolean shouldRun(Description description) {
                return true;
            }

            @Override
            public void apply(Object child) throws NoTestsRemainException {
                super.apply(child);
                if (child instanceof ParentRunner) {
                    ParentRunner runner = (ParentRunner) child;
                    if (!isThreadSafe(runner)) {
                        runner.setScheduler(notThreadSafeTests.newRunnerScheduler());
                    } else if (child instanceof Suite) {
                        nestedSuites.add((Suite) child);
                    } else {
                        ParentRunner parentRunner = (ParentRunner) child;
                        nestedClasses.add(parentRunner);
                        nestedClassesChildren +=
                                parentRunner.getDescription().getChildren().size();
                    }
                }
            }

            @Override
            public String describe() {
                return "";
            }
        }
    }

    private static Suite createSuite(Collection<Runner> runners) throws InitializationError {
        final List<Runner> onlyRunners = removeNullRunners(runners);
        return onlyRunners.isEmpty() ? null : new Suite(null, onlyRunners) {};
    }

    private static List<Runner> removeNullRunners(Collection<Runner> runners) {
        final List<Runner> onlyRunners = new ArrayList<>(runners);
        onlyRunners.removeAll(NULL_SINGLETON);
        return onlyRunners;
    }
}
