package jadx.gui.ui.codearea;

import javax.swing.text.Segment;

import org.fife.ui.rsyntaxtextarea.Token;
import org.fife.ui.rsyntaxtextarea.TokenImpl;
import org.fife.ui.rsyntaxtextarea.TokenTypes;
import org.fife.ui.rsyntaxtextarea.modes.JavaTokenMaker;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import jadx.api.JavaClass;
import jadx.api.JavaNode;

public final class JadxTokenMaker extends JavaTokenMaker {
	private static final Logger LOG = LoggerFactory.getLogger(JadxTokenMaker.class);

	private final CodeArea codeArea;

	public JadxTokenMaker(CodeArea codeArea) {
		this.codeArea = codeArea;
	}

	@Override
	public Token getTokenList(Segment text, int initialTokenType, int startOffset) {
		Token tokens = super.getTokenList(text, initialTokenType, startOffset);
		if (tokens.getType() != TokenTypes.NULL) {
			try {
				processTokens(tokens);
			} catch (Exception e) {
				LOG.error("Process tokens failed for text: {}", text, e);
			}
		}
		return tokens;
	}

	private void processTokens(Token tokens) {
		Token prev = null;
		Token current = tokens;
		while (current != null) {
			if (prev != null) {
				switch (current.getType()) {
					case TokenTypes.IDENTIFIER:
						current = mergeLongClassNames(prev, current, false);
						break;

					case TokenTypes.ANNOTATION:
						current = mergeLongClassNames(prev, current, true);
						break;
				}
			}
			prev = current;
			current = current.getNextToken();
		}
	}

	@NotNull
	private Token mergeLongClassNames(Token prev, Token current, boolean annotation) {
		int offset = current.getTextOffset();
		if (annotation) {
			offset++;
		}
		JavaNode javaNode = codeArea.getJavaNodeAtOffset(offset);
		if (javaNode instanceof JavaClass) {
			String name = javaNode.getName();
			String lexeme = current.getLexeme();
			if (annotation && lexeme.length() > 1) {
				lexeme = lexeme.substring(1);
			}
			if (!lexeme.equals(name) && isClassNameStart(javaNode, lexeme)) {
				// try to replace long class name with one token
				Token replace = concatTokensUntil(current, name);
				if (replace != null && prev instanceof TokenImpl) {
					TokenImpl impl = ((TokenImpl) prev);
					impl.setNextToken(replace);
					current = replace;
				}
			}
		}
		return current;
	}

	private boolean isClassNameStart(JavaNode javaNode, String lexeme) {
		if (javaNode.getFullName().startsWith(lexeme)) {
			// full class name
			return true;
		}
		if (javaNode.getTopParentClass().getName().startsWith(lexeme)) {
			// inner class references from parent class
			return true;
		}
		return false;
	}

	@Nullable
	private Token concatTokensUntil(Token start, String endText) {
		StringBuilder sb = new StringBuilder();
		Token current = start;
		while (current != null && current.getType() != TokenTypes.NULL) {
			String text = current.getLexeme();
			if (text != null) {
				sb.append(text);
				if (text.equals(endText)) {
					char[] line = sb.toString().toCharArray();
					TokenImpl token = new TokenImpl(line, 0, line.length - 1, start.getOffset(),
							start.getType(), start.getLanguageIndex());
					token.setNextToken(current.getNextToken());
					return token;
				}
			}
			current = current.getNextToken();
		}
		return null;
	}
}

@@ -24,16 +24,39 @@
 earea;%0A%0A
+import java.util.Set;%0A%0A
 import j
@@ -80,16 +80,16 @@
 egment;%0A
-
 %0Aimport 
@@ -470,16 +470,70 @@
 aNode;%0A%0A
+import static jadx.api.plugins.utils.Utils.constSet;%0A%0A
 public f
@@ -1315,24 +1315,111 @@
 etType()) %7B%0A
+%09%09%09%09%09case TokenTypes.RESERVED_WORD:%0A%09%09%09%09%09%09fixContextualKeyword(current);%0A%09%09%09%09%09%09break;%0A%0A
 %09%09%09%09%09case To
@@ -1646,16 +1646,16 @@
 urrent;%0A
-
 %09%09%09curre
@@ -1691,16 +1691,511 @@
 %09%09%7D%0A%09%7D%0A%0A
+%09private static final Set%3CString%3E CONTEXTUAL_KEYWORDS = constSet(%0A%09%09%09%22exports%22, %22module%22, %22non-sealed%22, %22open%22, %22opens%22, %22permits%22, %22provides%22, %22record%22,%0A%09%09%09%22requires%22, %22sealed%22, %22to%22, %22transitive%22, %22uses%22, %22var%22, %22with%22, %22yield%22);%0A%0A%09private static void fixContextualKeyword(Token token) %7B%0A%09%09String lexeme = token.getLexeme(); // TODO: create new string every call, better to avoid%0A%09%09if (lexeme != null && CONTEXTUAL_KEYWORDS.contains(lexeme)) %7B%0A%09%09%09token.setType(TokenTypes.IDENTIFIER);%0A%09%09%7D%0A%09%7D%0A%0A
 %09@NotNul

revert contextual keywords to identifiers (#1394)