/*

 * Copyright (c) 2014, 2015, 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.

 *

 * 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.

 */

import java.io.ByteArrayInputStream;

import java.io.ByteArrayOutputStream;

import java.io.IOException;

import java.io.InputStream;

import java.io.PrintStream;

import java.io.StringWriter;

import java.lang.reflect.Method;

import java.nio.file.Path;

import java.util.ArrayList;

import java.util.Arrays;

import java.util.Collection;

import java.util.Collections;

import java.util.HashMap;

import java.util.LinkedHashMap;

import java.util.LinkedHashSet;

import java.util.List;

import java.util.Map;

import java.util.Set;

import java.util.TreeMap;

import java.util.function.Consumer;

import java.util.function.Predicate;

import java.util.function.Supplier;

import java.util.stream.Collectors;

import java.util.stream.Stream;

import javax.tools.Diagnostic;

import jdk.jshell.EvalException;

import jdk.jshell.JShell;

import jdk.jshell.JShell.Subscription;

import jdk.jshell.Snippet;

import jdk.jshell.DeclarationSnippet;

import jdk.jshell.ExpressionSnippet;

import jdk.jshell.ImportSnippet;

import jdk.jshell.Snippet.Kind;

import jdk.jshell.MethodSnippet;

import jdk.jshell.Snippet.Status;

import jdk.jshell.Snippet.SubKind;

import jdk.jshell.TypeDeclSnippet;

import jdk.jshell.VarSnippet;

import jdk.jshell.SnippetEvent;

import jdk.jshell.SourceCodeAnalysis;

import jdk.jshell.SourceCodeAnalysis.CompletionInfo;

import jdk.jshell.SourceCodeAnalysis.Completeness;

import jdk.jshell.SourceCodeAnalysis.QualifiedNames;

import jdk.jshell.SourceCodeAnalysis.Suggestion;

import jdk.jshell.UnresolvedReferenceException;

import org.testng.annotations.AfterMethod;

import org.testng.annotations.BeforeMethod;

import jdk.jshell.Diag;

import static java.util.stream.Collectors.toList;

import static java.util.stream.Collectors.toSet;

import static jdk.jshell.Snippet.Status.*;

import static org.testng.Assert.*;

import static jdk.jshell.Snippet.SubKind.METHOD_SUBKIND;

public class KullaTesting {

    public static final String IGNORE_VALUE = "<ignore-value>";

    public static final Class<? extends Throwable> IGNORE_EXCEPTION = (new Throwable() {}).getClass();

    public static final Snippet MAIN_SNIPPET;

    private SourceCodeAnalysis analysis = null;

    private JShell state = null;

    private InputStream inStream = null;

    private ByteArrayOutputStream outStream = null;

    private ByteArrayOutputStream errStream = null;

    private Map<String, Snippet> idToSnippet = new LinkedHashMap<>();

    private Set<Snippet> allSnippets = new LinkedHashSet<>();

    private List<String> classpath;

    static {

        JShell js = JShell.create();

        MAIN_SNIPPET = js.eval("MAIN_SNIPPET").get(0).snippet();

        js.close();

        assertTrue(MAIN_SNIPPET != null, "Bad MAIN_SNIPPET set-up -- must not be null");

    }

    public enum DiagCheck {

        DIAG_OK,

        DIAG_WARNING,

        DIAG_ERROR,

        DIAG_IGNORE

    }

    public void setInput(String s) {

        setInput(new ByteArrayInputStream(s.getBytes()));

    }

    public void setInput(InputStream in) {

        inStream = in;

    }

    public String getOutput() {

        String s = outStream.toString();

        outStream.reset();

        return s;

    }

    public String getErrorOutput() {

        String s = errStream.toString();

        errStream.reset();

        return s;

    }

    /**

     * @return the analysis

     */

    public SourceCodeAnalysis getAnalysis() {

        if (analysis == null) {

            analysis = state.sourceCodeAnalysis();

        }

        return analysis;

    }

    /**

     * @return the state

     */

    public JShell getState() {

        return state;

    }

    public List<Snippet> getActiveKeys() {

        return allSnippets.stream()

                .filter(k -> getState().status(k).isActive())

                .collect(Collectors.toList());

    }

    public void addToClasspath(String path) {

        classpath.add(path);

        getState().addToClasspath(path);

    }

    public void addToClasspath(Path path) {

        addToClasspath(path.toString());

    }

    @BeforeMethod

    public void setUp() {

        setUp(b -> {});

    }

    public void setUp(Consumer<JShell.Builder> bc) {

        InputStream in = new InputStream() {

            @Override

            public int read() throws IOException {

                assertNotNull(inStream);

                return inStream.read();

            }

            @Override

            public int read(byte[] b) throws IOException {

                assertNotNull(inStream);

                return inStream.read(b);

            }

            @Override

            public int read(byte[] b, int off, int len) throws IOException {

                assertNotNull(inStream);

                return inStream.read(b, off, len);

            }

        };

        outStream = new ByteArrayOutputStream();

        errStream = new ByteArrayOutputStream();

        JShell.Builder builder = JShell.builder()

                .in(in)

                .out(new PrintStream(outStream))

                .err(new PrintStream(errStream));

        bc.accept(builder);

        state = builder.build();

        allSnippets = new LinkedHashSet<>();

        idToSnippet = new LinkedHashMap<>();

        classpath = new ArrayList<>();

    }

    @AfterMethod

    public void tearDown() {

        if (state != null) state.close();

        state = null;

        analysis = null;

        allSnippets = null;

        idToSnippet = null;

        classpath = null;

    }

    public List<String> assertUnresolvedDependencies(DeclarationSnippet key, int unresolvedSize) {

        List<String> unresolved = getState().unresolvedDependencies(key).collect(toList());

        assertEquals(unresolved.size(), unresolvedSize, "Input: " + key.source() + ", checking unresolved: ");

        return unresolved;

    }

    public DeclarationSnippet assertUnresolvedDependencies1(DeclarationSnippet key, Status status, String name) {

        List<String> unresolved = assertUnresolvedDependencies(key, 1);

        String input = key.source();

        assertEquals(unresolved.size(), 1, "Given input: " + input + ", checking unresolved");

        assertEquals(unresolved.get(0), name, "Given input: " + input + ", checking unresolved: ");

        assertEquals(getState().status(key), status, "Given input: " + input + ", checking status: ");

        return key;

    }

    public DeclarationSnippet assertEvalUnresolvedException(String input, String name, int unresolvedSize, int diagnosticsSize) {

        List<SnippetEvent> events = assertEval(input, null, UnresolvedReferenceException.class, DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, null);

        SnippetEvent ste = events.get(0);

        DeclarationSnippet sn = ((UnresolvedReferenceException) ste.exception()).getSnippet();

        assertEquals(sn.name(), name, "Given input: " + input + ", checking name");

        assertEquals(getState().unresolvedDependencies(sn).count(), unresolvedSize, "Given input: " + input + ", checking unresolved");

        assertEquals(getState().diagnostics(sn).count(), (long) diagnosticsSize, "Given input: " + input + ", checking diagnostics");

        return sn;

    }

    public Snippet assertKeyMatch(String input, boolean isExecutable, SubKind expectedSubKind, STEInfo mainInfo, STEInfo... updates) {

        Snippet key = key(assertEval(input, IGNORE_VALUE, mainInfo, updates));

        String source = key.source();

        assertEquals(source, input, "Key \"" + input + "\" source mismatch, got: " + source + ", expected: " + input);

        SubKind subkind = key.subKind();

        assertEquals(subkind, expectedSubKind, "Key \"" + input + "\" subkind mismatch, got: "

                + subkind + ", expected: " + expectedSubKind);

        assertEquals(subkind.isExecutable(), isExecutable, "Key \"" + input + "\", expected isExecutable: "

                + isExecutable + ", got: " + subkind.isExecutable());

        Snippet.Kind expectedKind = getKind(key);

        assertEquals(key.kind(), expectedKind, "Checking kind: ");

        assertEquals(expectedSubKind.kind(), expectedKind, "Checking kind: ");

        return key;

    }

    private Kind getKind(Snippet key) {

        SubKind expectedSubKind = key.subKind();

        Kind expectedKind;

        switch (expectedSubKind) {

            case SINGLE_TYPE_IMPORT_SUBKIND:

            case SINGLE_STATIC_IMPORT_SUBKIND:

            case TYPE_IMPORT_ON_DEMAND_SUBKIND:

            case STATIC_IMPORT_ON_DEMAND_SUBKIND:

                expectedKind = Kind.IMPORT;

                break;

            case CLASS_SUBKIND:

            case INTERFACE_SUBKIND:

            case ENUM_SUBKIND:

            case ANNOTATION_TYPE_SUBKIND:

                expectedKind = Kind.TYPE_DECL;

                break;

            case METHOD_SUBKIND:

                expectedKind = Kind.METHOD;

                break;

            case VAR_DECLARATION_SUBKIND:

            case TEMP_VAR_EXPRESSION_SUBKIND:

            case VAR_DECLARATION_WITH_INITIALIZER_SUBKIND:

                expectedKind = Kind.VAR;

                break;

            case VAR_VALUE_SUBKIND:

            case ASSIGNMENT_SUBKIND:

                expectedKind = Kind.EXPRESSION;

                break;

            case STATEMENT_SUBKIND:

                expectedKind = Kind.STATEMENT;

                break;

            case UNKNOWN_SUBKIND:

                expectedKind = Kind.ERRONEOUS;

                break;

            default:

                throw new AssertionError("Unsupported key: " + key.getClass().getCanonicalName());

        }

        return expectedKind;

    }

    public ImportSnippet assertImportKeyMatch(String input, String name, SubKind subkind, STEInfo mainInfo, STEInfo... updates) {

        Snippet key = assertKeyMatch(input, false, subkind, mainInfo, updates);

        assertTrue(key instanceof ImportSnippet, "Expected an ImportKey, got: " + key.getClass().getName());

        ImportSnippet importKey = (ImportSnippet) key;

        assertEquals(importKey.name(), name, "Input \"" + input +

                "\" name mismatch, got: " + importKey.name() + ", expected: " + name);

        assertEquals(importKey.kind(), Kind.IMPORT, "Checking kind: ");

        return importKey;

    }

    public DeclarationSnippet assertDeclarationKeyMatch(String input, boolean isExecutable, String name, SubKind subkind, STEInfo mainInfo, STEInfo... updates) {

        Snippet key = assertKeyMatch(input, isExecutable, subkind, mainInfo, updates);

        assertTrue(key instanceof DeclarationSnippet, "Expected a DeclarationKey, got: " + key.getClass().getName());

        DeclarationSnippet declKey = (DeclarationSnippet) key;

        assertEquals(declKey.name(), name, "Input \"" + input +

                "\" name mismatch, got: " + declKey.name() + ", expected: " + name);

        return declKey;

    }

    public VarSnippet assertVarKeyMatch(String input, boolean isExecutable, String name, SubKind kind, String typeName, STEInfo mainInfo, STEInfo... updates) {

        Snippet sn = assertDeclarationKeyMatch(input, isExecutable, name, kind, mainInfo, updates);

        assertTrue(sn instanceof VarSnippet, "Expected a VarKey, got: " + sn.getClass().getName());

        VarSnippet variableKey = (VarSnippet) sn;

        String signature = variableKey.typeName();

        assertEquals(signature, typeName, "Key \"" + input +

                "\" typeName mismatch, got: " + signature + ", expected: " + typeName);

        assertEquals(variableKey.kind(), Kind.VAR, "Checking kind: ");

        return variableKey;

    }

    public void assertExpressionKeyMatch(String input, String name, SubKind kind, String typeName) {

        Snippet key = assertKeyMatch(input, true, kind, added(VALID));

        assertTrue(key instanceof ExpressionSnippet, "Expected a ExpressionKey, got: " + key.getClass().getName());

        ExpressionSnippet exprKey = (ExpressionSnippet) key;

        assertEquals(exprKey.name(), name, "Input \"" + input +

                "\" name mismatch, got: " + exprKey.name() + ", expected: " + name);

        assertEquals(exprKey.typeName(), typeName, "Key \"" + input +

                "\" typeName mismatch, got: " + exprKey.typeName() + ", expected: " + typeName);

        assertEquals(exprKey.kind(), Kind.EXPRESSION, "Checking kind: ");

    }

    // For expressions throwing an EvalException

    public SnippetEvent assertEvalException(String input) {

        List<SnippetEvent> events = assertEval(input, null, EvalException.class,

                DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, null);

        return events.get(0);

    }

    public List<SnippetEvent> assertEvalFail(String input) {

        return assertEval(input, null, null,

                DiagCheck.DIAG_ERROR, DiagCheck.DIAG_IGNORE, added(REJECTED));

    }

    public List<SnippetEvent> assertEval(String input) {

        return assertEval(input, IGNORE_VALUE, null, DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, added(VALID));

    }

    public List<SnippetEvent> assertEval(String input, String value) {

        return assertEval(input, value, null, DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, added(VALID));

    }

    public List<SnippetEvent> assertEval(String input, STEInfo mainInfo, STEInfo... updates) {

        return assertEval(input, IGNORE_VALUE, null, DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, mainInfo, updates);

    }

    public List<SnippetEvent> assertEval(String input, String value,

            STEInfo mainInfo, STEInfo... updates) {

        return assertEval(input, value, null, DiagCheck.DIAG_OK, DiagCheck.DIAG_OK, mainInfo, updates);

    }

    public List<SnippetEvent> assertEval(String input, DiagCheck diagMain, DiagCheck diagUpdates) {

        return assertEval(input, IGNORE_VALUE, null, diagMain, diagUpdates, added(VALID));

    }

    public List<SnippetEvent> assertEval(String input, DiagCheck diagMain, DiagCheck diagUpdates,

            STEInfo mainInfo, STEInfo... updates) {

        return assertEval(input, IGNORE_VALUE, null, diagMain, diagUpdates, mainInfo, updates);

    }

    public List<SnippetEvent> assertEval(String input,

            String value, Class<? extends Throwable> exceptionClass,

            DiagCheck diagMain, DiagCheck diagUpdates,

            STEInfo mainInfo, STEInfo... updates) {

        return assertEval(input, diagMain, diagUpdates, new EventChain(mainInfo, value, exceptionClass, updates));

    }

    // Use this directly or usually indirectly for all non-empty calls to eval()

    public List<SnippetEvent> assertEval(String input,

           DiagCheck diagMain, DiagCheck diagUpdates, EventChain... eventChains) {

        return checkEvents(() -> getState().eval(input), "eval(" + input + ")", diagMain, diagUpdates, eventChains);

    }

    <T> void assertStreamMatch(Stream<T> result, T... expected) {

        Set<T> sns = result.collect(toSet());

        Set<T> exp = Stream.of(expected).collect(toSet());

        assertEquals(sns, exp);

    }

    private Map<Snippet, Snippet> closure(List<SnippetEvent> events) {

        Map<Snippet, Snippet> transitions = new HashMap<>();

        for (SnippetEvent event : events) {

            transitions.put(event.snippet(), event.causeSnippet());

        }

        Map<Snippet, Snippet> causeSnippets = new HashMap<>();

        for (Map.Entry<Snippet, Snippet> entry : transitions.entrySet()) {

            Snippet snippet = entry.getKey();

            Snippet cause = getInitialCause(transitions, entry.getValue());

            causeSnippets.put(snippet, cause);

        }

        return causeSnippets;

    }

    private Snippet getInitialCause(Map<Snippet, Snippet> transitions, Snippet snippet) {

        Snippet result;

        while ((result = transitions.get(snippet)) != null) {

            snippet = result;

        }

        return snippet;

    }

    private Map<Snippet, List<SnippetEvent>> groupByCauseSnippet(List<SnippetEvent> events) {

        Map<Snippet, List<SnippetEvent>> map = new TreeMap<>((a, b) -> a.id().compareTo(b.id()));

        for (SnippetEvent event : events) {

            if (event == null) {

                throw new InternalError("null event found in " + events);

            }

            if (event.snippet() == null) {

                throw new InternalError("null event Snippet found in " + events);

            }

            if (event.snippet().id() == null) {

                throw new InternalError("null event Snippet id() found in " + events);

            }

        }

        for (SnippetEvent event : events) {

            if (event.causeSnippet() == null) {

                map.computeIfAbsent(event.snippet(), ($) -> new ArrayList<>()).add(event);

            }

        }

        Map<Snippet, Snippet> causeSnippets = closure(events);

        for (SnippetEvent event : events) {

            Snippet causeSnippet = causeSnippets.get(event.snippet());

            if (causeSnippet != null) {

                map.get(causeSnippet).add(event);

            }

        }

        for (Map.Entry<Snippet, List<SnippetEvent>> entry : map.entrySet()) {

            Collections.sort(entry.getValue(),

                    (a, b) -> a.causeSnippet() == null

                            ? -1 : b.causeSnippet() == null

                            ? 1 : a.snippet().id().compareTo(b.snippet().id()));

        }

        return map;

    }

    private List<STEInfo> getInfos(EventChain... eventChains) {

        List<STEInfo> list = new ArrayList<>();

        for (EventChain i : eventChains) {

            list.add(i.mainInfo);

            Collections.addAll(list, i.updates);

        }

        return list;

    }

    private List<SnippetEvent> checkEvents(Supplier<List<SnippetEvent>> toTest,

             String descriptor,

             DiagCheck diagMain, DiagCheck diagUpdates,

             EventChain... eventChains) {

        List<SnippetEvent> dispatched = new ArrayList<>();

        Subscription token = getState().onSnippetEvent(kse -> {

            if (dispatched.size() > 0 && dispatched.get(dispatched.size() - 1) == null) {

                throw new RuntimeException("dispatch event after done");

            }

            dispatched.add(kse);

        });

        List<SnippetEvent> events = toTest.get();

        getState().unsubscribe(token);

        assertEquals(dispatched.size(), events.size(), "dispatched event size not the same as event size");

        for (int i = events.size() - 1; i >= 0; --i) {

            assertEquals(dispatched.get(i), events.get(i), "Event element " + i + " does not match");

        }

        dispatched.add(null); // mark end of dispatchs

        for (SnippetEvent evt : events) {

            assertTrue(evt.snippet() != null, "key must never be null, but it was for: " + descriptor);

            assertTrue(evt.previousStatus() != null, "previousStatus must never be null, but it was for: " + descriptor);

            assertTrue(evt.status() != null, "status must never be null, but it was for: " + descriptor);

            assertTrue(evt.status() != NONEXISTENT, "status must not be NONEXISTENT: " + descriptor);

            if (evt.previousStatus() != NONEXISTENT) {

                Snippet old = idToSnippet.get(evt.snippet().id());

                if (old != null) {

                    switch (evt.status()) {

                        case DROPPED:

                            assertEquals(old, evt.snippet(),

                                    "Drop: Old snippet must be what is dropped -- input: " + descriptor);

                            break;

                        case OVERWRITTEN:

                            assertEquals(old, evt.snippet(),

                                    "Overwrite: Old snippet (" + old

                                    + ") must be what is overwritten -- input: "

                                    + descriptor + " -- " + evt);

                            break;

                        default:

                            if (evt.causeSnippet() == null) {

                                // New source

                                assertNotEquals(old, evt.snippet(),

                                        "New source: Old snippet must be different from the replacing -- input: "

                                        + descriptor);

                            } else {

                                // An update (key Overwrite??)

                                assertEquals(old, evt.snippet(),

                                        "Update: Old snippet must be equal to the replacing -- input: "

                                        + descriptor);

                            }

                            break;

                    }

                }

            }

        }

        for (SnippetEvent evt : events) {

            if (evt.causeSnippet() == null && evt.status() != DROPPED) {

                allSnippets.add(evt.snippet());

                idToSnippet.put(evt.snippet().id(), evt.snippet());

            }

        }

        assertTrue(events.size() >= 1, "Expected at least one event, got none.");

        List<STEInfo> all = getInfos(eventChains);

        if (events.size() != all.size()) {

            StringBuilder sb = new StringBuilder();

            sb.append("Got events --\n");

            for (SnippetEvent evt : events) {

                sb.append("  key: ").append(evt.snippet());