package jnpf.elnFormula;

import jnpf.elnFormula.parser.ElnFormulaBaseVisitor;
import jnpf.elnFormula.parser.ElnFormulaLexer;
import jnpf.elnFormula.parser.ElnFormulaParser.AdditiveContext;
import jnpf.elnFormula.parser.ElnFormulaParser.ArgumentsContext;
import jnpf.elnFormula.parser.ElnFormulaParser.ComparisonContext;
import jnpf.elnFormula.parser.ElnFormulaParser.EqualityContext;
import jnpf.elnFormula.parser.ElnFormulaParser.ExpressionContext;
import jnpf.elnFormula.parser.ElnFormulaParser.LogicalAndContext;
import jnpf.elnFormula.parser.ElnFormulaParser.LogicalOrContext;
import jnpf.elnFormula.parser.ElnFormulaParser.MultiplicativeContext;
import jnpf.elnFormula.parser.ElnFormulaParser.PowerContext;
import jnpf.elnFormula.parser.ElnFormulaParser.PrimaryContext;
import jnpf.elnFormula.parser.ElnFormulaParser.UnaryContext;
import org.antlr.v4.runtime.BaseErrorListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.LexerNoViableAltException;
import org.antlr.v4.runtime.RecognitionException;
import org.antlr.v4.runtime.Recognizer;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.tree.TerminalNode;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public final class FormulaParser {

    private static final Set<String> FORBIDDEN_IDENTIFIERS;

    static {
        Set<String> values = new HashSet<>();
        Collections.addAll(values, "new", "class", "import", "package", "this", "super");
        FORBIDDEN_IDENTIFIERS = Collections.unmodifiableSet(values);
    }

    private final FormulaLimits limits;

    public FormulaParser() {
        this(FormulaLimits.DEFAULTS);
    }

    public FormulaParser(FormulaLimits limits) {
        if (limits == null) {
            throw new IllegalArgumentException("limits must not be null");
        }
        this.limits = limits;
    }

    public static boolean isReservedParameterCode(String code) {
        return "true".equals(code) || "false".equals(code) || "null".equals(code)
                || FORBIDDEN_IDENTIFIERS.contains(code);
    }

    public FormulaAst.Node parse(String expression) {
        String source = expression == null ? "" : expression;
        if (source.length() > limits.getMaxExpressionLength()) {
            throw error("EXPRESSION_TOO_LONG", "公式长度超过限制", 0, source.length());
        }
        SourceOffsets offsets = new SourceOffsets(source);

        CollectingErrorListener lexerErrors = new CollectingErrorListener("LEXICAL_ERROR", offsets);
        ElnFormulaLexer lexer = new ElnFormulaLexer(CharStreams.fromString(source));
        lexer.removeErrorListeners();
        lexer.addErrorListener(lexerErrors);
        CommonTokenStream tokens = new CommonTokenStream(lexer);
        tokens.fill();
        lexerErrors.throwFirst();

        int tokenCount = 0;
        for (Token token : tokens.getTokens()) {
            if (token.getType() == Token.EOF) {
                continue;
            }
            tokenCount++;
            if ((token.getType() == ElnFormulaLexer.IDENTIFIER
                    || token.getType() == ElnFormulaLexer.FUNCTION)
                    && token.getText().length() > 80) {
                throw error("IDENTIFIER_TOO_LONG", "标识符长度超过限制",
                        offsets.toUtf16(token.getStartIndex()),
                        offsets.toUtf16(token.getStopIndex() + 1));
            }
            if (token.getType() == ElnFormulaLexer.IDENTIFIER
                    && FORBIDDEN_IDENTIFIERS.contains(token.getText())) {
                throw error("FORBIDDEN_IDENTIFIER", "标识符不允许使用",
                        offsets.toUtf16(token.getStartIndex()),
                        offsets.toUtf16(token.getStopIndex() + 1));
            }
        }
        if (tokenCount > limits.getMaxTokens()) {
            throw error("TOKEN_LIMIT_EXCEEDED", "公式 token 数量超过限制", 0, source.length());
        }

        tokens.seek(0);
        CollectingErrorListener parserErrors = new CollectingErrorListener("SYNTAX_ERROR", offsets);
        jnpf.elnFormula.parser.ElnFormulaParser generated =
                new jnpf.elnFormula.parser.ElnFormulaParser(tokens);
        generated.removeErrorListeners();
        generated.addErrorListener(parserErrors);
        jnpf.elnFormula.parser.ElnFormulaParser.ParseContext tree = generated.parse();
        parserErrors.throwFirst();

        return new AstBuilder(limits, offsets).visitExpression(tree.expression()).node;
    }

    private static ParseException error(String code, String message, int start, int end) {
        int safeStart = Math.max(0, start);
        return new ParseException(code, message, safeStart, Math.max(safeStart, end));
    }

    public static final class ParseException extends RuntimeException {
        private final String code;
        private final int startOffset;
        private final int endOffset;

        private ParseException(String code, String message, int startOffset, int endOffset) {
            super(message);
            this.code = code;
            this.startOffset = startOffset;
            this.endOffset = endOffset;
        }

        public String getCode() {
            return code;
        }

        public int getStartOffset() {
            return startOffset;
        }

        public int getEndOffset() {
            return endOffset;
        }
    }

    private static final class CollectingErrorListener extends BaseErrorListener {
        private final String code;
        private final SourceOffsets offsets;
        private ParseException first;

        private CollectingErrorListener(String code, SourceOffsets offsets) {
            this.code = code;
            this.offsets = offsets;
        }

        @Override
        public void syntaxError(Recognizer<?, ?> recognizer, Object offendingSymbol,
                                int line, int charPositionInLine, String message,
                                RecognitionException exception) {
            if (first != null) {
                return;
            }
            int start = offsets.absoluteCodePoint(line, charPositionInLine);
            int end = start + 1;
            if (exception instanceof LexerNoViableAltException) {
                start = ((LexerNoViableAltException) exception).getStartIndex();
                end = start + 1;
            } else if (offendingSymbol instanceof Token) {
                Token token = (Token) offendingSymbol;
                start = Math.max(0, token.getStartIndex());
                end = token.getType() == Token.EOF ? start : token.getStopIndex() + 1;
            }
            String safeMessage = "LEXICAL_ERROR".equals(code)
                    ? "公式包含非法字符" : "公式语法错误";
            first = error(code, safeMessage, offsets.toUtf16(start), offsets.toUtf16(end));
        }

        private void throwFirst() {
            if (first != null) {
                throw first;
            }
        }
    }

    private static final class SourceOffsets {
        private final int[] codePointToUtf16;
        private final List<Integer> lineStartCodePoints;

        private SourceOffsets(String source) {
            int codePointCount = source.codePointCount(0, source.length());
            this.codePointToUtf16 = new int[codePointCount + 1];
            this.lineStartCodePoints = new ArrayList<>();
            lineStartCodePoints.add(0);

            int utf16Offset = 0;
            for (int codePointOffset = 0; codePointOffset < codePointCount; codePointOffset++) {
                codePointToUtf16[codePointOffset] = utf16Offset;
                int codePoint = source.codePointAt(utf16Offset);
                utf16Offset += Character.charCount(codePoint);
                if (codePoint == '\n') {
                    lineStartCodePoints.add(codePointOffset + 1);
                }
            }
            codePointToUtf16[codePointCount] = source.length();
        }

        private int toUtf16(int codePointOffset) {
            int safeOffset = Math.max(0, Math.min(codePointOffset, codePointToUtf16.length - 1));
            return codePointToUtf16[safeOffset];
        }

        private int absoluteCodePoint(int line, int charPositionInLine) {
            int lineIndex = Math.max(0, Math.min(line - 1, lineStartCodePoints.size() - 1));
            return lineStartCodePoints.get(lineIndex) + Math.max(0, charPositionInLine);
        }
    }

    private static final class AstValue {
        private final FormulaAst.Node node;
        private final int depth;

        private AstValue(FormulaAst.Node node, int depth) {
            this.node = node;
            this.depth = depth;
        }
    }

    private static final class AstBuilder extends ElnFormulaBaseVisitor<AstValue> {
        private final FormulaLimits limits;
        private final SourceOffsets offsets;

        private AstBuilder(FormulaLimits limits, SourceOffsets offsets) {
            this.limits = limits;
            this.offsets = offsets;
        }

        @Override
        public AstValue visitExpression(ExpressionContext context) {
            return visit(context.logicalOr());
        }

        @Override
        public AstValue visitLogicalOr(LogicalOrContext context) {
            return fold(context.logicalAnd(), context.OR());
        }

        @Override
        public AstValue visitLogicalAnd(LogicalAndContext context) {
            return fold(context.equality(), context.AND());
        }

        @Override
        public AstValue visitEquality(EqualityContext context) {
            return fold(context.comparison(), terminals(context.EQ(), context.NE()));
        }

        @Override
        public AstValue visitComparison(ComparisonContext context) {
            return fold(context.additive(),
                    terminals(context.GT(), context.GE(), context.LT(), context.LE()));
        }

        @Override
        public AstValue visitAdditive(AdditiveContext context) {
            return fold(context.multiplicative(), terminals(context.PLUS(), context.MINUS()));
        }

        @Override
        public AstValue visitMultiplicative(MultiplicativeContext context) {
            return fold(context.unary(), terminals(context.STAR(), context.SLASH(), context.PERCENT()));
        }

        @Override
        public AstValue visitUnary(UnaryContext context) {
            if (context.unary() == null) {
                return visit(context.power());
            }
            AstValue operand = visit(context.unary());
            Token operator = context.getStart();
            return checked(new FormulaAst.Unary(start(operator), operand.node.getEndOffset(),
                    operator.getText(), operand.node), operand.depth + 1);
        }

        @Override
        public AstValue visitPower(PowerContext context) {
            AstValue left = visit(context.primary());
            if (context.unary() == null) {
                return left;
            }
            return binary(left, context.CARET().getSymbol(), visit(context.unary()));
        }

        @Override
        public AstValue visitPrimary(PrimaryContext context) {
            if (context.DECIMAL() != null) {
                Token token = context.DECIMAL().getSymbol();
                return leaf(new FormulaAst.Literal(start(token), end(token),
                        FormulaAst.LiteralKind.DECIMAL, new BigDecimal(token.getText())));
            }
            if (context.STRING() != null) {
                Token token = context.STRING().getSymbol();
                return leaf(new FormulaAst.Literal(start(token), end(token),
                        FormulaAst.LiteralKind.TEXT, unescape(token.getText())));
            }
            if (context.TRUE() != null || context.FALSE() != null) {
                Token token = context.TRUE() != null
                        ? context.TRUE().getSymbol() : context.FALSE().getSymbol();
                return leaf(new FormulaAst.Literal(start(token), end(token),
                        FormulaAst.LiteralKind.BOOLEAN, Boolean.valueOf(token.getText())));
            }
            if (context.NULL() != null) {
                Token token = context.NULL().getSymbol();
                return leaf(new FormulaAst.Literal(start(token), end(token),
                        FormulaAst.LiteralKind.NULL, null));
            }
            if (context.IDENTIFIER() != null) {
                Token token = context.IDENTIFIER().getSymbol();
                return leaf(new FormulaAst.Parameter(start(token), end(token), token.getText()));
            }
            if (context.FUNCTION() != null) {
                return function(context);
            }
            AstValue grouped = visitExpression(context.expression());
            return checked(grouped.node, grouped.depth + 1);
        }

        private AstValue function(PrimaryContext context) {
            Token function = context.FUNCTION().getSymbol();
            List<FormulaAst.Node> arguments = new ArrayList<>();
            int depth = 1;
            ArgumentsContext argumentsContext = context.arguments();
            if (argumentsContext != null) {
                if (argumentsContext.expression().size() > limits.getMaxFunctionArguments()) {
                    throw error("FUNCTION_ARGUMENT_LIMIT_EXCEEDED", "函数参数数量超过限制",
                            start(function), end(context.getStop()));
                }
                for (ExpressionContext item : argumentsContext.expression()) {
                    AstValue value = visitExpression(item);
                    arguments.add(value.node);
                    depth = Math.max(depth, value.depth + 1);
                }
            }
            return checked(new FormulaAst.FunctionCall(start(function), end(context.getStop()),
                    function.getText(), arguments), depth);
        }

        private <T extends org.antlr.v4.runtime.ParserRuleContext> AstValue fold(
                List<T> operands, List<TerminalNode> operators) {
            AstValue result = visit(operands.get(0));
            for (int i = 1; i < operands.size(); i++) {
                result = binary(result, operators.get(i - 1).getSymbol(), visit(operands.get(i)));
            }
            return result;
        }

        @SafeVarargs
        private final List<TerminalNode> terminals(List<TerminalNode>... groups) {
            List<TerminalNode> result = new ArrayList<>();
            for (List<TerminalNode> group : groups) {
                result.addAll(group);
            }
            result.sort((left, right) -> Integer.compare(
                    left.getSymbol().getTokenIndex(), right.getSymbol().getTokenIndex()));
            return result;
        }

        private AstValue binary(AstValue left, Token operator, AstValue right) {
            return checked(new FormulaAst.Binary(left.node.getStartOffset(), right.node.getEndOffset(),
                    operator.getText(), left.node, right.node),
                    Math.max(left.depth, right.depth) + 1);
        }

        private AstValue leaf(FormulaAst.Node node) {
            return checked(node, 1);
        }

        private AstValue checked(FormulaAst.Node node, int depth) {
            if (depth > limits.getMaxAstDepth()) {
                throw error("AST_DEPTH_EXCEEDED", "公式嵌套深度超过限制",
                        node.getStartOffset(), node.getEndOffset());
            }
            return new AstValue(node, depth);
        }

        private int start(Token token) {
            return offsets.toUtf16(token.getStartIndex());
        }

        private int end(Token token) {
            return offsets.toUtf16(token.getStopIndex() + 1);
        }

        private static String unescape(String tokenText) {
            StringBuilder result = new StringBuilder(tokenText.length() - 2);
            for (int i = 1; i < tokenText.length() - 1; i++) {
                char current = tokenText.charAt(i);
                if (current != '\\') {
                    result.append(current);
                    continue;
                }
                char escaped = tokenText.charAt(++i);
                if (escaped == 'n') {
                    result.append('\n');
                } else if (escaped == 'r') {
                    result.append('\r');
                } else if (escaped == 't') {
                    result.append('\t');
                } else {
                    result.append(escaped);
                }
            }
            return result.toString();
        }
    }
}