package com.boothen.jsonedit.model;

import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CommonToken;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.Token;
import org.eclipse.jface.text.BadLocationException;
import org.eclipse.jface.text.IDocument;
import org.eclipse.jface.text.TextAttribute;
import org.eclipse.jface.text.rules.IToken;
import org.eclipse.jface.text.rules.ITokenScanner;

import com.boothen.jsonedit.antlr.JSONLexer;
import com.boothen.jsonedit.core.JsonLog;
import com.boothen.jsonedit.core.preferences.TokenStyle;

/**
 * Use the ANTLR lexer to extract tokens from a document.
 */
public class AntlrTokenScanner implements ITokenScanner {

    private int offset;
    private Lexer lexer;
    private CommonToken previous;
    private CommonToken current;
    private TokenStyler tokenStyler;

    /**
     * Uses the given JSONLexer
     * @param lexer the lexer to use
     */
    public AntlrTokenScanner(Lexer lexer) {
        this.lexer = lexer;
        this.tokenStyler = new TokenStyler() {
            @Override
            public TextAttribute apply(TokenStyle style) {
                return new TextAttribute(null);
            }
        };
    }

    /**
     * @param lexer the lexer to use
     * @param mapping the mapping from token type to result
     */
    public AntlrTokenScanner(Lexer lexer, TokenStyler mapping) {
        this.lexer = lexer;
        this.tokenStyler = mapping;
    }

    @Override
    public void setRange(IDocument document, int offset, int length) {
        this.offset = offset;
        try {
            String text = document.get(offset, length);
            lexer.setInputStream(new ANTLRInputStream(text));
        } catch (BadLocationException e) {
            JsonLog.logError("Attempting to access a non-existing position", e);
        }
    }

    @Override
    public IToken nextToken() {
        Token token = lexer.nextToken();

        if (!(token instanceof CommonToken)) {
            // this doesn't seem to happen in practice
            // if it ever does: don't update current/previous, just bail out
            return org.eclipse.jface.text.rules.Token.UNDEFINED;
        }

        // update only if token is relevant (no WS)
        if (current == null || current.getType() != JSONLexer.WS) {
            previous = current;
        }

        current = (CommonToken) token;

        if (token.getType() == Token.EOF) {
            return org.eclipse.jface.text.rules.Token.EOF;
        }

        if (token.getType() == JSONLexer.WS) {
            return org.eclipse.jface.text.rules.Token.WHITESPACE;
        }

        int currentType = current.getType();
        int previousType = previous != null ? previous.getType() : Token.EOF;
        TokenStyle style = getStyle(currentType, previousType);
        TextAttribute data = tokenStyler.apply(style);

        return new org.eclipse.jface.text.rules.Token(data);
    }

    @Override
    public int getTokenOffset() {
        return offset + current.getStartIndex();
    }

    @Override
    public int getTokenLength() {
        return current.getStopIndex() - current.getStartIndex() + 1;
    }

    private TokenStyle getStyle(int currentTokenType, int previousTokenType) {
        switch (currentTokenType) {
        case JSONLexer.STRING:
            if (previousTokenType == JSONLexer.COLON) {
                return TokenStyle.TEXT;
            } else {
                return TokenStyle.KEY;
            }

        case JSONLexer.NUMBER:
            return TokenStyle.NUMBER;

        case JSONLexer.TRUE:
        case JSONLexer.FALSE:
            return TokenStyle.BOOLEAN;

        case JSONLexer.NULL:
            return TokenStyle.NULL;

        default:
            return TokenStyle.DEFAULT;
        }
    }
}

@@ -604,16 +604,21 @@
 extract 
+JSON 
 tokens f
@@ -633,16 +633,137 @@
 cument.%0A
+ * It uses a 2-step lookahead to identify KEY text strings. It therefore does not require to%0A * construct a syntax tree.%0A
  */%0Apubl
@@ -890,24 +890,53 @@
 onToken 
-previous
+current;%0A    private CommonToken next
 ;%0A    pr
@@ -945,38 +945,40 @@
 ate CommonToken 
-curren
+afterNex
 t;%0A    private T
@@ -1863,16 +1863,125 @@
 text));%0A
+            next = (CommonToken) lexer.nextToken();%0A            afterNext = (CommonToken) lexer.nextToken();%0A
         
@@ -2468,64 +2468,8 @@
     
-// update only if token is relevant (no WS)%0A        if (
 curr
@@ -2477,112 +2477,57 @@
 nt =
-=
  n
-ull %7C%7C current.getType() != JSONLexer.WS) %7B%0A            previous = current;%0A        %7D%0A%0A        curren
+ext;%0A        next = afterNext;%0A        afterNex
 t = 
@@ -2552,37 +2552,39 @@
 n;%0A%0A        if (
-tok
+curr
 en
+t
 .getType() == To
@@ -2676,21 +2676,23 @@
     if (
-tok
+curr
 en
+t
 .getType
@@ -2851,72 +2851,82 @@
 int 
-previousType = previous != null ? previous.getType() : Token.EOF
+nextType = next.getType();%0A        int afterNextType = afterNext.getType()
 ;%0A  
@@ -2972,24 +2972,35 @@
 ntType, 
-previous
+nextType, afterNext
 Type);%0A 
@@ -3398,35 +3398,40 @@
 entT
-okenType, int previousToken
+ype, int nextType, int afterNext
 Type
@@ -3458,21 +3458,16 @@
 currentT
-okenT
 ype) %7B%0A 
@@ -3512,48 +3512,246 @@
     
-if (previousTokenType == JSONLexer.COLON
+// There could be whitespace between current and the next (real) token%0A            // we therefore need to look two tokens ahead%0A            if (nextType == JSONLexer.COLON %7C%7C (nextType == JSONLexer.WS && afterNextType == JSONLexer.COLON)
 ) %7B%0A
@@ -3788,12 +3788,11 @@
 yle.
-TEXT
+KEY
 ;%0A  
@@ -3836,35 +3836,36 @@
 turn TokenStyle.
-KEY
+TEXT
 ;%0A            %7D%0A

Use two-step lookahead to corrently identify keys