package net.ranides.assira.grammar.impl;

import lombok.Data;
import net.ranides.assira.collection.maps.CrossMap;
import net.ranides.assira.collection.maps.OpenCrossMap;
import net.ranides.assira.collection.sets.OpenSet;
import net.ranides.assira.generic.Ref;
import net.ranides.assira.grammar.G4ErrorHandler;
import net.ranides.assira.grammar.G4Exception;
import net.ranides.assira.grammar.G4Node;
import net.ranides.assira.grammar.G4Parser;
import net.ranides.assira.grammar.G4ParserCoverage;
import net.ranides.assira.grammar.impl.CoverageSupport.SimpleBranchCoverage;
import net.ranides.assira.grammar.impl.CoverageSupport.SimpleRuleCoverage;
import net.ranides.assira.text.StringUtils;
import org.antlr.runtime.ANTLRStringStream;
import org.antlr.runtime.RecognitionException;
import org.antlr.runtime.tree.CommonTree;
import org.antlr.runtime.tree.TreeVisitor;
import org.antlr.runtime.tree.TreeVisitorAction;
import org.antlr.v4.Tool;
import org.antlr.v4.analysis.LeftRecursiveRuleAltInfo;
import org.antlr.v4.parse.GrammarASTAdaptor;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.LexerInterpreter;
import org.antlr.v4.runtime.ParserInterpreter;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.Vocabulary;
import org.antlr.v4.runtime.misc.Pair;
import org.antlr.v4.runtime.tree.RuleNode;
import org.antlr.v4.tool.ANTLRMessage;
import org.antlr.v4.tool.ANTLRToolListener;
import org.antlr.v4.tool.Alternative;
import org.antlr.v4.tool.Grammar;
import org.antlr.v4.tool.LabelElementPair;
import org.antlr.v4.tool.Rule;
import org.antlr.v4.tool.ast.AltAST;
import org.antlr.v4.tool.ast.GrammarAST;
import org.antlr.v4.tool.ast.GrammarRootAST;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.TreeMap;

public class G4ParserInterpreter implements G4Parser {

    private final Grammar grammar;

    private final G4ErrorHandler handler;

    private final CrossMap<Integer, Integer, String> branchLabels;

    private final Set<String> branchValidator;

    private final CrossMap<String, String, Integer> expressionLabels;

    public G4ParserInterpreter(String grammar, String imports, G4ErrorHandler handler) {
        try {
            if(imports == null) {
                this.grammar = new Grammar(grammar, new ANTLRToolListenerAdapter(handler));
            } else {
                this.grammar = newGrammar(grammar, imports, handler);
            }

            this.handler = handler;
            this.branchLabels = deduceBranchLabels(this.grammar);
            this.expressionLabels = deduceExpressionLabels(this.grammar);

            this.branchValidator = new OpenSet<>();
            branchLabels.forEach((ruleIndex, alt, label)-> {
                branchValidator.add(NodeSupport.QUALIFIER_NAME + label);
                branchValidator.add(ruleName(ruleIndex) + NodeSupport.QUALIFIER_NAME + label);
            });
        } catch (RecognitionException e) {
            throw new G4Exception(e);
        }
    }

    private Grammar newGrammar(String grammarText, String imports, G4ErrorHandler handler) throws RecognitionException {
        Tool tool = new Tool();
        tool.addListener(new ANTLRToolListenerAdapter(handler));

        ANTLRStringStream in = new ANTLRStringStream(grammarText);
        in.name = "((string))";

        GrammarRootAST ast = tool.parse(in.name, in);
        if ( ast==null ) {
            throw new UnsupportedOperationException();
        }

        if (ast.tokenStream == null) {
            throw new IllegalStateException("expected ast to have a token stream");
        }

        Grammar out = new Grammar(tool, ast);

        in.name = out.fileName = out.name + ".g4";
        tool.libDirectory = imports;

        new TreeVisitor(new GrammarASTAdaptor()).visit(ast, new TreeVisitorAction() {
            @Override
            public Object pre(Object t) { ((GrammarAST)t).g = out; return t; }
            @Override
            public Object post(Object t) { return t; }
        });

        tool.process(out, false);

        return out;
    }

    /**
     * Tricky method: we deduce alternative labels for every rule
     */
    private static CrossMap<Integer, Integer, String> deduceBranchLabels(Grammar grammar) {
        CrossMap<Integer, Integer, String> labels = new OpenCrossMap<>();

        for(Rule currentRule : grammar.rules.values()) {
            int ruleIndex = currentRule.index;
            Map<String, List<Pair<Integer, AltAST>>> altLabels = currentRule.getAltLabels();
            if(altLabels == null) {
                continue;
            }
            for(List<Pair<Integer, AltAST>> list : altLabels.values()) {
                for(Pair<Integer, AltAST> pair : list) {
                    Integer altIndex = pair.a;
                    GrammarAST altLabel = pair.b.altLabel;
                    labels.put(ruleIndex, altIndex, altLabel.toString());
                }
            }
        }

        return labels;
    }

    /**
     * Tricky method: we deduce node positions for labels used inside each rule
     */
    private static CrossMap<String, String, Integer> deduceExpressionLabels(Grammar grammar) {
        CrossMap<String, String, Integer> expressionLabels = new OpenCrossMap<>();

        for (Rule rule : grammar.rules.values()) {
            String ruleName = rule.name;
            for (Alternative alternative : rule.alt) {
                if(alternative == null) {
                    continue;
                }
                for (Map.Entry<String, List<LabelElementPair>> entry : alternative.labelDefs.entrySet()) {
                    String labelName = entry.getKey();
                    for (LabelElementPair pair : entry.getValue()) {
                        int labelIndex = deduceLabelIndex(ruleName, labelName, pair.element);
                        String labelQualifier = decudeLabelQualifier(ruleName, labelName, pair.element);
                        expressionLabels.put(ruleName + labelQualifier, labelName, labelIndex);
                    }
                }
            }
        }

        return expressionLabels;
    }

    private static int deduceLabelIndex(String ruleName, String labelName, GrammarAST element) {
        return element.parent.childIndex;
    }

    private static String decudeLabelQualifier(String ruleName, String labelName, GrammarAST element) {
        for (CommonTree now = element; now!=null; now = now.parent) {
            if (!(now instanceof AltAST)) {
                continue;
            }
            AltAST alt = (AltAST) now;

            LeftRecursiveRuleAltInfo left = alt.leftRecursiveAltInfo;

            if (left != null) {
                if (left.altLabel != null) {
                    return NodeSupport.QUALIFIER_NAME + left.altLabel;
                } else {
                    return NodeSupport.QUALIFIER_INDEX + left.altNum;
                }
            }

            if (alt.altLabel != null) {
                return NodeSupport.QUALIFIER_NAME + alt.altLabel;
            }

            if (alt.alt != null) {
                return NodeSupport.QUALIFIER_INDEX + alt.alt.altNum;
            }

            if (alt.parent == null) {
                return NodeSupport.QUALIFIER_INDEX + (now.childIndex + 1);
            }
        }

        throw new G4Exception("Awkward label '" + labelName + "' inside rule: " + ruleName);
    }

    @Override
    public boolean validateRuleName(String ruleName) {
        if(ruleName.contains("#")) {
            return branchValidator.contains(ruleName);
        } else {
            return grammar.getRule(ruleName) != null;
        }
    }

    @Override
    public boolean validateTokenName(String tokenName) {
        return grammar.getTokenType(tokenName) != Token.INVALID_TYPE;
    }

    @Override
    public String ruleName(int ruleIndex) {
        return grammar.getRule(ruleIndex).name;
    }

    @Override
    public String ruleName(Class<?> ruleType) {
        return StringUtils.uncapitalize(ruleType.getSimpleName().replaceAll("Context$", ""));
    }

    @Override
    public String tokenName(int tokenIndex) {
        return grammar.getVocabulary().getSymbolicName(tokenIndex);
    }

    @Override
    public G4Node parse(String rule, String text) {
        return new MyRoot(rule, text).tree();
    }

    @Override
    public G4ParserCoverage rules() {
        return new MyCoverage();
    }

    private class MyRoot extends RootAbstract {

        private final NodeForRule tree;

        public MyRoot(String ruleName, String text) {
            LexerInterpreter lexer = G4ParserInterpreter.this.grammar.createLexerInterpreter(CharStreams.fromString(text));

            lexer.removeErrorListeners();
            lexer.addErrorListener(new NodeSupport.DelegateListener(handler));

            CommonTokenStream tokens = new CommonTokenStream(lexer);
            ParserInterpreter interpreter = G4ParserInterpreter.this.grammar.createGrammarParserInterpreter(tokens);

            interpreter.removeErrorListeners();
            interpreter.addErrorListener(new NodeSupport.DelegateListener(handler));

            Rule rule = G4ParserInterpreter.this.grammar.getRule(ruleName);
            if (rule == null) {
                throw new G4Exception("Unknown rule: " + ruleName);
            }

            this.tree = new NodeForRule(this, null, interpreter.parse(rule.index));
        }

        @Override
        public G4Parser parser() {
            return G4ParserInterpreter.this;
        }

        @Override
        public G4Node tree() {
            return tree;
        }

        @Override
        public String branchName(G4Node g4node) {
            RuleNode node = (RuleNode) g4node.unwrap();
            int ruleIndex = node.getRuleContext().getRuleIndex();
            int ruleAlt = node.getRuleContext().getAltNumber();
            String label = branchLabels.get(ruleIndex, ruleAlt);

            return label != null ?
                NodeSupport.QUALIFIER_NAME + label :
                NodeSupport.QUALIFIER_INDEX + ruleAlt;
        }

        @Override
        public Ref<G4Node> findLabel(G4Node g4node, String labelName) {
            Integer index = expressionLabels.get(g4node.name(), labelName);
            if(index != null) {
                return g4node.children().at(index);
            } else {
                return Ref.empty(() -> new NoSuchElementException("Unknown label '" + labelName + "' inside '" + g4node + "'"));
            }
        }
    }

    private class MyCoverage extends G4ParserCoverage {

        private final Map<String, RuleCoverage> rules;
        private final Map<String, BranchCoverage> branches;
        private final Map<String, String> aliases;

        public MyCoverage() {
            rules = new TreeMap<>();
            branches = new TreeMap<>();
            aliases = new TreeMap<>();

            Vocabulary v = grammar.getVocabulary();

            for(int i=1, n = v.getMaxTokenType(); i<=n; i++) {
                String token = v.getSymbolicName(i);
                rules.put(token, new SimpleRuleCoverage(token, Ref.empty()));
            }

            for(Rule  rule : grammar.rules.values()) {
                getRuleAliases(rule);
                String ruleName = rule.name;
                Ref<String> alias = Ref.ofNullable(aliases.get(ruleName));
                rules.put(ruleName, new SimpleRuleCoverage(ruleName, alias));
            }

            for(String fqn : branchValidator) {
                if(!fqn.startsWith(NodeSupport.QUALIFIER_NAME)) {
                    String[] m = fqn.split(NodeSupport.QUALIFIER_NAME);
                    String ruleName = m[0];
                    String qualifier = NodeSupport.QUALIFIER_NAME + m[1];
                    Ref<String> alias = Ref.ofNullable(aliases.get(ruleName + qualifier));
                    branches.put(qualifier, new SimpleBranchCoverage(ruleName, qualifier, alias));
                }
            }

            branches.values().forEach(branch -> {
                rules.get(branch.rule()).branches().add(branch);
            });
        }

        private void getRuleAliases(Rule rule) {
            if(rule.alt != null) {
                for (Alternative alt : rule.alt) {
                    if (alt != null) {
                        String alias = null;
                        if(alt.ruleRefs.size() == 0 && alt.tokenRefs.size() == 1) {
                            alias = alt.tokenRefs.keySet().iterator().next();
                        }
                        if(alt.ruleRefs.size() == 1 && alt.tokenRefs.size() == 0) {
                            alias = alt.ruleRefs.keySet().iterator().next();
                        }
                        if(alias != null && alt.ast.getChildren().size() == 1) {
                            int ri = rule.index;
                            int ai = alt.altNum;
                            String altName = branchLabels.get(ri, ai);

                            if(altName != null) {
                                aliases.put(rule.name + NodeSupport.QUALIFIER_NAME+ altName, alias);
                            }
                            else if(rule.alt.length == 2) {
                                aliases.put(rule.name, alias);
                            }
                        }
                    }
                }
            }
        }

        @Override
        public Collection<G4ParserCoverage.RuleCoverage> rules() {
            return rules.values();
        }

        @Override
        public Collection<BranchCoverage> branches() {
            return branches.values();
        }

        @Override
        public Ref<G4ParserCoverage.RuleCoverage> rule(String ruleName) {
            return Ref.ofNullable(rules.get(ruleName));
        }

        @Override
        public Ref<G4ParserCoverage.BranchCoverage> branch(String ruleName) {
            return Ref.ofNullable(branches.get(ruleName));
        }
    }

    @Data
    private static final class ANTLRToolListenerAdapter implements ANTLRToolListener {

        private final G4ErrorHandler handler;

        @Override
        public void info(String msg) {
            handler.onError(msg, new G4Exception(msg));
        }

        @Override
        public void error(ANTLRMessage msg) {
            handler.onError(msg.toString(), msg.getCause());
        }

        @Override
        public void warning(ANTLRMessage msg) {
            handler.onError(msg.toString(), msg.getCause());
        }
    }
}
