/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */

package org.elasticsearch.painless.action;

import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.Token;
import org.elasticsearch.painless.ScriptTestCase;
import org.elasticsearch.painless.antlr.EnhancedSuggestLexer;
import org.elasticsearch.painless.antlr.SuggestLexer;
import org.elasticsearch.painless.spi.PainlessTestScript;

import java.util.List;

public class SuggestTests extends ScriptTestCase {

    private List<? extends Token> getSuggestTokens(String source) {
        ANTLRInputStream stream = new ANTLRInputStream(source);
        SuggestLexer lexer = new EnhancedSuggestLexer(stream, scriptEngine.getContextsToLookups().get(PainlessTestScript.CONTEXT));
        lexer.removeErrorListeners();
        return lexer.getAllTokens();
    }

    private void compareTokens(List<? extends Token> tokens, String... expected) {
        assertEquals(expected.length % 2, 0);
        assertEquals(tokens.size(), expected.length / 2);

        int index = 0;
        for (Token token : tokens) {
            assertEquals(SuggestLexer.VOCABULARY.getDisplayName(token.getType()), expected[index++]);
            assertEquals(token.getText(), expected[index++]);
        }
    }

    public void testSuggestLexer() {
        compareTokens(getSuggestTokens("test"), SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID), "test");

        compareTokens(
            getSuggestTokens("int test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "int",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("ArrayList test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "ArrayList",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("def test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "def",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("int[] test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ATYPE),
            "int[]",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("ArrayList[] test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ATYPE),
            "ArrayList[]",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("def[] test;"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ATYPE),
            "def[]",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );

        compareTokens(
            getSuggestTokens("List test = new ArrayList(); test."),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "List",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ASSIGN),
            "=",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.NEW),
            "new",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "ArrayList",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "(",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RP),
            ")",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.DOT),
            "."
        );

        compareTokens(
            getSuggestTokens("List test = new ArrayList(); test.add"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "List",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ASSIGN),
            "=",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.NEW),
            "new",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "ArrayList",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "(",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RP),
            ")",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.DOT),
            ".",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.DOTID),
            "add"
        );

        compareTokens(
            getSuggestTokens("List test = new ArrayList(); test.add("),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "List",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ASSIGN),
            "=",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.NEW),
            "new",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "ArrayList",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "(",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RP),
            ")",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.DOT),
            ".",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.DOTID),
            "add",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "("
        );

        compareTokens(
            getSuggestTokens("def test(int param) {return param;} test(2);"),
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "def",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "(",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.TYPE),
            "int",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "param",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RP),
            ")",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LBRACK),
            "{",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RETURN),
            "return",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "param",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RBRACK),
            "}",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.ID),
            "test",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.LP),
            "(",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.INTEGER),
            "2",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.RP),
            ")",
            SuggestLexer.VOCABULARY.getDisplayName(SuggestLexer.SEMICOLON),
            ";"
        );
    }
}
