

package org.snt.inmemantlr.listener;

import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.tree.ErrorNode;
import org.antlr.v4.runtime.tree.TerminalNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.snt.inmemantlr.tree.ParseTree;
import org.snt.inmemantlr.tree.ParseTreeNode;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Predicate;

/**
 * default tree listener implementation
 */
public class DefaultTreeListener extends DefaultListener {

    private static final Logger LOGGER = LoggerFactory.getLogger(DefaultTreeListener.class);

    private static final long serialVersionUID = 5637734678821255670L;

    protected Deque<String> sctx = new ArrayDeque<>();
    protected StringBuilder glob = new StringBuilder();
    protected ParseTree parseTree = null;
    protected ParseTreeNode nodeptr = null;
    protected Predicate<String> filter = null;
    protected boolean includeTerminals = false;

    /**
     * constructor
     */
    public DefaultTreeListener() {
        this(false);
    }

    /**
     * constructor
     * @param includeTerminals include terminal nodes in resulting tree
     */
    public DefaultTreeListener(boolean includeTerminals) {
        this(x -> !x.isEmpty());
        this.includeTerminals = includeTerminals;
    }

    /**
     * constructor
     *
     * @param filter condition that has to hold for every node
     */
    public DefaultTreeListener(Predicate<String> filter) {
        sctx.add("S");
        parseTree = new ParseTree("root", "root");
        nodeptr = parseTree.getRoot();
        this.filter = filter;
    }

    @Override
    public void visitTerminal(TerminalNode terminalNode) {
        if(includeTerminals) {
            Token token = terminalNode.getSymbol();
            ParseTreeNode n = parseTree.newNode(nodeptr,"",
                    terminalNode.toString(),
                    token.getStartIndex(),
                    token.getStopIndex(),
                    token.getLine(),
                    token.getCharPositionInLine());
            nodeptr.addChild(n);
        }
    }

    @Override
    public void visitErrorNode(ErrorNode errorNode) {
    }

    @Override
    public void enterEveryRule(ParserRuleContext ctx) {
        String rule = getRuleByKey(ctx.getRuleIndex());
        if (filter.test(rule)) {
            Token s = ctx.getStart();
            Token e = ctx.getStop();
            ParseTreeNode n = parseTree.newNode(nodeptr, rule, ctx.getText(),
                    s != null ? s.getStartIndex() : 0,
                    e != null ? e.getStopIndex() : 0,
                    s != null ? s.getLine() : 0,
                    s != null ? s.getCharPositionInLine() : 0);
            nodeptr.addChild(n);
            nodeptr = n;
        }
    }

    @Override
    public void exitEveryRule(ParserRuleContext ctx) {
        String rule = getRuleByKey(ctx.getRuleIndex());
        if (filter.test(rule)) {
            nodeptr = nodeptr.getParent();
        }
    }

    @Override
    public void reset() {
        super.reset();
        sctx.clear();
        sctx.add("S");
        parseTree = new ParseTree("root", "root");
        nodeptr = parseTree.getRoot();
        glob.delete(0, glob.length());
    }

    /**
     * get ast
     *
     * @return ast
     */
    public ParseTree getParseTree() {
        return parseTree;
    }

    /**
     * get ast nodes
     *
     * @return set of ast nodes
     */
    public Set<ParseTreeNode> getNodes() {
        return new HashSet<>(parseTree.getNodes());
    }

    @Override
    public String toString() {
        return glob.toString();
    }
}
