package com.intellij.plugins.thrift.annotator;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import com.intellij.lang.annotation.Annotation;
import com.intellij.lang.annotation.AnnotationHolder;
import com.intellij.lang.annotation.Annotator;
import com.intellij.openapi.editor.DefaultLanguageHighlighterColors;
import com.intellij.openapi.editor.colors.TextAttributesKey;
import com.intellij.plugins.thrift.highlight.ThriftSyntaxHighlighterColors;
import com.intellij.plugins.thrift.lang.lexer.ThriftTokenTypes;
import com.intellij.plugins.thrift.lang.psi.ThriftDefinitionName;
import com.intellij.plugins.thrift.lang.psi.ThriftService;
import com.intellij.plugins.thrift.lang.psi.ThriftStruct;
import com.intellij.plugins.thrift.lang.psi.ThriftTypedef;
import com.intellij.plugins.thrift.lang.psi.ThriftUnion;
import com.intellij.plugins.thrift.lang.psi.ThriftVisitor;
import com.intellij.plugins.thrift.util.ThriftUtils;
import com.intellij.psi.PsiElement;
import com.intellij.psi.impl.source.tree.LeafPsiElement;
import com.intellij.psi.tree.IElementType;

/**
 * Created by fkorotkov.
 */
public class ThriftColorAnnotator extends ThriftVisitor implements Annotator
{
	private AnnotationHolder myAnnotationHolder;

	@Override
	public void annotate(@NotNull PsiElement element, @NotNull AnnotationHolder holder)
	{
		myAnnotationHolder = holder;
		element.accept(this);
	}

	@Override
	public void visitElement(PsiElement element)
	{
		if(element instanceof LeafPsiElement)
		{
			IElementType tokenType = ((LeafPsiElement) element).getElementType();
			if(tokenType == ThriftTokenTypes.IDENTIFIER && ThriftUtils.getKeywords().contains(element.getText()))
			{
				annotateKeyword(element);
			}
		}
	}

	@Override
	public void visitStruct(@NotNull ThriftStruct o)
	{
		highlightName(o.getDefinitionName(), DefaultLanguageHighlighterColors.CLASS_NAME);
	}

	@Override
	public void visitUnion(@NotNull ThriftUnion o)
	{
		highlightName(o.getDefinitionName(), DefaultLanguageHighlighterColors.CLASS_NAME);
	}

	@Override
	public void visitTypedef(@NotNull ThriftTypedef o)
	{
		highlightName(o.getDefinitionName(), DefaultLanguageHighlighterColors.TYPE_ALIAS_NAME);
	}

	@Override
	public void visitService(@NotNull ThriftService o)
	{
		highlightName(o.getDefinitionName(), DefaultLanguageHighlighterColors.CLASS_NAME);
	}

	private void highlightName(@Nullable ThriftDefinitionName definitionName, TextAttributesKey key)
	{
		if(definitionName == null)
		{
			return;
		}
		Annotation annotation = myAnnotationHolder.createInfoAnnotation(definitionName.getNameIdentifier(), null);
		annotation.setTextAttributes(key);
	}

	private void annotateKeyword(@NotNull PsiElement element)
	{
		final Annotation annotation = myAnnotationHolder.createInfoAnnotation(element, null);
		annotation.setTextAttributes(TextAttributesKey.find(ThriftSyntaxHighlighterColors.THRIFT_KEYWORD));
	}
}

@@ -611,16 +611,77 @@
 onName;%0A
+import com.intellij.plugins.thrift.lang.psi.ThriftException;%0A
 import c
@@ -1873,32 +1873,35 @@
 %09%09highlightName(
+o, 
 o.getDefinitionN
@@ -1909,53 +1909,8 @@
 me()
-, DefaultLanguageHighlighterColors.CLASS_NAME
 );%0A%09
@@ -1982,32 +1982,35 @@
 %09%09highlightName(
+o, 
 o.getDefinitionN
@@ -2018,53 +2018,8 @@
 me()
-, DefaultLanguageHighlighterColors.CLASS_NAME
 );%0A%09
@@ -2095,32 +2095,35 @@
 %09%09highlightName(
+o, 
 o.getDefinitionN
@@ -2131,58 +2131,121 @@
 me()
-, DefaultLanguageHighlighterColors.TYPE_ALIAS_NAME
+);%0A%09%7D%0A%0A%09@Override%0A%09public void visitService(@NotNull ThriftService o)%0A%09%7B%0A%09%09highlightName(o, o.getDefinitionName()
 );%0A%09
@@ -2268,39 +2268,41 @@
 ublic void visit
-Service
+Exception
 (@NotNull Thrift
@@ -2293,39 +2293,41 @@
 (@NotNull Thrift
-Service
+Exception
  o)%0A%09%7B%0A%09%09highlig
@@ -2333,16 +2333,19 @@
 ghtName(
+o, 
 o.getDef
@@ -2361,54 +2361,423 @@
 me()
-, DefaultLanguageHighlighterColors.CLASS_NAME)
+);%0A%09%7D%0A%0A%09private static TextAttributesKey getAttributesKey(PsiElement element)%0A%09%7B%0A%09%09if(element instanceof ThriftService %7C%7C element instanceof ThriftUnion %7C%7C element instanceof ThriftStruct %7C%7C element instanceof ThriftException)%0A%09%09%7B%0A%09%09%09return DefaultLanguageHighlighterColors.CLASS_NAME;%0A%09%09%7D%0A%09%09else if(element instanceof ThriftTypedef)%0A%09%09%7B%0A%09%09%09return DefaultLanguageHighlighterColors.TYPE_ALIAS_NAME;%0A%09%09%7D%0A%09%09return null
 ;%0A%09%7D
@@ -2806,16 +2806,44 @@
 ghtName(
+@NotNull PsiElement parent, 
 @Nullabl
@@ -2883,31 +2883,8 @@
 Name
-, TextAttributesKey key
 )%0A%09%7B
@@ -2903,16 +2903,126 @@
 tionName
+ == null)%0A%09%09%7B%0A%09%09%09return;%0A%09%09%7D%0A%0A%09%09TextAttributesKey attributesKey = getAttributesKey(parent);%0A%09%09if(attributesKey
  == null
@@ -3182,17 +3182,27 @@
 ributes(
-k
+attributesK
 ey);%0A%09%7D%0A

rewrite highlight