package com.intellij.bash;

import com.intellij.bash.psi.BashHeredoc;
import com.intellij.bash.psi.BashString;
import com.intellij.bash.psi.BashVariable;
import com.intellij.lang.annotation.Annotation;
import com.intellij.lang.annotation.AnnotationHolder;
import com.intellij.lang.annotation.Annotator;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiRecursiveElementWalkingVisitor;
import org.jetbrains.annotations.NotNull;

import static com.intellij.bash.BashSyntaxHighlighter.STRING;
import static com.intellij.bash.BashSyntaxHighlighter.VAR_USE;

public class BashAnnotator implements Annotator {
  @Override
  public void annotate(@NotNull PsiElement o, @NotNull AnnotationHolder holder) {
    if (o instanceof BashString) {
      holder.createInfoAnnotation(o, null).setTextAttributes(STRING);
      highlightVariables(o, holder);
    }
    else if (o instanceof BashHeredoc) {
      annotateHeredoc((BashHeredoc) o, holder);
    }
  }

  private void annotateHeredoc(BashHeredoc o, AnnotationHolder holder) {
    Annotation annotation = holder.createInfoAnnotation(o, null);
    annotation.setTextAttributes(BashSyntaxHighlighter.HERE_DOC);
    highlightVariables(o, holder);
  }

  private void highlightVariables(@NotNull PsiElement container, @NotNull AnnotationHolder holder) {
    new PsiRecursiveElementWalkingVisitor() {
      @Override
      public void visitElement(PsiElement element) {
        if (element instanceof BashVariable) {
          holder.createInfoAnnotation(element, null).setTextAttributes(VAR_USE);
        }
        super.visitElement(element);
      }
    }.visitElement(container);
  }
}

@@ -151,56 +151,8 @@
 le;%0A
-import com.intellij.lang.annotation.Annotation;%0A
 impo
@@ -244,24 +244,85 @@
 .Annotator;%0A
+import com.intellij.openapi.editor.colors.TextAttributesKey;%0A
 import com.i
@@ -450,70 +450,51 @@
 ll;%0A
-%0A
 import 
-static com.intellij.bash.BashSyntaxHighlighter.STRING;
+org.jetbrains.annotations.Nullable;%0A
 %0Aimp
@@ -544,23 +544,17 @@
 lighter.
-VAR_USE
+*
 ;%0A%0Apubli
@@ -737,63 +737,24 @@
     
-holder.createInfoAnnotation(o, null).setTextAttributes(
+mark(o, holder, 
 STRI
@@ -932,16 +932,25 @@
 Heredoc(
+@NotNull 
 BashHere
@@ -955,16 +955,25 @@
 redoc o,
+ @NotNull
  Annotat
@@ -1000,133 +1000,146 @@
     
-Annotation annotation = holder.createInfoAnnotation(o, null);%0A    annotation.setTextAttributes(BashSyntaxHighlighter.
+mark(o, holder, HERE_DOC);%0A    mark(o.getHeredocMarkerStart(), holder, HERE_DOC_START);%0A    mark(o.getHeredocMarkerEnd(), holder, 
 HERE_DOC
 );%0A 
@@ -1130,27 +1130,32 @@
 er, HERE_DOC
+_END
 );%0A
+%0A
     highligh
@@ -1459,69 +1459,30 @@
     
-holder.createInfoAnnotation(element, null).setTextAttributes(
+mark(element, holder, 
 VAR_
@@ -1576,11 +1576,215 @@
 er);%0A  %7D
+%0A%0A  private static void mark(@Nullable PsiElement o, @NotNull AnnotationHolder holder, @NotNull TextAttributesKey key) %7B%0A    if (o != null) holder.createInfoAnnotation(o, null).setTextAttributes(key);%0A  %7D
 %0A%7D%0A

mark heredoc start/end