/*
 * Copyright (c) 2014, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

package com.sun.javafx;

import com.sun.javafx.stage.StageHelper;
import javafx.animation.AnimationTimer;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.stage.Stage;
import junit.framework.Assert;
import org.junit.internal.runners.statements.Fail;
import org.junit.internal.runners.statements.InvokeMethod;
import org.junit.rules.MethodRule;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.Statement;

import java.lang.ref.Reference;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class FXUnit implements MethodRule {

    private static CountDownLatch startupLatch = new CountDownLatch(1);

    public static class TestApplication extends Application {
        public TestApplication() { }
        @Override
        public void start(Stage primaryStage) throws Exception {
            Platform.setImplicitExit(false);
            primaryStage.close();
            startupLatch.countDown();
        }
    }

    public static class AsyncException extends RuntimeException {
        private AsyncException(Exception e) {
            super();
            initCause(e);
        }
    }

    private class InvokeAndWaitStatement extends Statement {
        private Statement statement;
        private InvokeAndWaitStatement(Statement statement) {
            this.statement = statement;
        }

        @Override
        public void evaluate() throws Throwable {
            invokeAndWait(statement);
        }
    }

    public static interface TestRunnable {
        public void run() throws Exception;
    }

    public FXUnit() {
        synchronized (FXUnit.class) {
            startToolkit();
        }
    }


    @Override
    public Statement apply(Statement statement, FrameworkMethod frameworkMethod,
                           Object test) {
        List<Statement> befores = new ArrayList<>();
        List<Statement> afters = new ArrayList<>();
        try {
            Method[] methods = test.getClass().getMethods();
            for (Method m : methods) {
                FXBefore before = m.getAnnotation(FXBefore.class);
                if (before != null) {
                    InvokeMethod invoker = new InvokeMethod(
                            new FrameworkMethod(m), test);
                    if (before.value().equals(FXThread.MAIN)) {
                        befores.add(invoker);
                    } else {
                        befores.add(new InvokeAndWaitStatement(invoker));
                    }
                }
                FXAfter after = m.getAnnotation(FXAfter.class);
                if (after != null) {
                    InvokeMethod invoker = new InvokeMethod(
                            new FrameworkMethod(m), test);
                    if (after.value().equals(FXThread.MAIN)) {
                        afters.add(invoker);
                    } else {
                        afters.add(new InvokeAndWaitStatement(invoker));
                    }
                }
            }
        } catch (Exception e) {
            return new Fail(e);
        }
        FXTestThread threadAnnotation = frameworkMethod.getAnnotation(FXTestThread.class);
        FXThread testThread = threadAnnotation == null
                              ? FXThread.APPLICATION
                              : threadAnnotation.value();
        Statement newStatement = testThread.equals(FXThread.MAIN)
                ? statement
                : new InvokeAndWaitStatement(statement);
        return new Statement() {
            @Override
            public void evaluate() throws Throwable {
                for (Statement s : befores) {
                    s.evaluate();
                }
                newStatement.evaluate();
                for (Statement s : afters) {
                    s.evaluate();
                }
            }
        };
    }

    public synchronized void startToolkit() {
        if (startupLatch.getCount() != 0) {
            if (Boolean.getBoolean("headless")) {
                System.getProperty("glass.platform", "Monocle");
                System.getProperty("monocle.platform", "Headless");
                System.getProperty("prism.order", "sw");
            }
            new Thread(() -> Application.launch(TestApplication.class)).start();
            try {
                startupLatch.await();
            } catch (InterruptedException e) {
                Assert.fail("Interrupted when initializing toolkit");
            }
        } else {
            invokeAndWait(() -> {
                List<Stage> stages = new ArrayList<>();
                stages.addAll(StageHelper.getStages());
                stages.forEach((stage) -> stage.close());
            });
            waitForNextPulse();
        }
    }

    public void waitForNextPulse() {
        CountDownLatch pulseLatch = new CountDownLatch(1);
        Platform.runLater(() -> new AnimationTimer() {
            @Override
            public void handle(long now) {
                pulseLatch.countDown();
                stop();
            }
        }.start());
        try {
            pulseLatch.await();
        } catch (InterruptedException e) {
            Assert.fail("Interrupted when waiting for the next pulse");
        }
    }

    public void invokeAndWait(TestRunnable runnable) {
        invokeAndWait(new Statement() {
            @Override
            public void evaluate() throws Throwable {
                runnable.run();
            }
        });
    }

    private void invokeAndWait(Statement statement) {
        if (Platform.isFxApplicationThread()) {
            try {
                statement.evaluate();
            } catch (RuntimeException | Error e) {
                throw e;
            } catch (Throwable t) {
                throw new AsyncException((Exception) t);
            }
        } else {
            final CountDownLatch runLatch = new CountDownLatch(1);
            final AtomicReference<Throwable> throwable = new AtomicReference<>();
            Platform.runLater(() -> {
                try {
                    try {
                        statement.evaluate();
                    } catch (Throwable t) {
                        throwable.set(t);
                    }
                } finally {
                    runLatch.countDown();
                }
            });
            try {
                runLatch.await();
            } catch (InterruptedException e) {
                throw new AsyncException(e);
            }
            Throwable t = throwable.get();
            if (t != null) {
                if (t instanceof Error) {
                    throw (Error) t;
                } else if (t instanceof RuntimeException) {
                    throw (RuntimeException) t;
                } else if (t instanceof Exception) {
                    throw new AsyncException((Exception) t);
                }
            }
        }
    }

}