package net.wizardsoflua.annotation.processor;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Messager;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.tools.Diagnostic;

import com.google.common.collect.Maps;

public abstract class ExceptionHandlingProcessor extends AbstractProcessor {
  private final Map<Entry<TypeElement, Element>, Exception> retryNextRound = new HashMap<>();

  @Override
  public final boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
    if (!roundEnv.processingOver()) {
      Set<Entry<TypeElement, Element>> retryNow = new HashSet<>(retryNextRound.keySet());
      retryNextRound.clear();
      for (TypeElement annotation : annotations) {
        Set<? extends Element> annotatedElements = roundEnv.getElementsAnnotatedWith(annotation);
        for (Element annotatedElement : annotatedElements) {
          process(annotation, annotatedElement, roundEnv);
        }
      }
      for (Entry<TypeElement, Element> entry : retryNow) {
        TypeElement annotation = entry.getKey();
        Element annotatedElement = entry.getValue();
        process(annotation, annotatedElement, roundEnv);
      }
    } else {
      processingOver();
      logExceptions();
    }
    return areAnnotationsClaimed(annotations);
  }

  protected boolean areAnnotationsClaimed(Set<? extends TypeElement> annotations) {
    return false;
  }

  protected void processingOver() {}

  private void logExceptions() {
    for (Entry<Entry<TypeElement, Element>, Exception> entry : retryNextRound.entrySet()) {
      Exception ex = entry.getValue();
      Messager messager = processingEnv.getMessager();
      if (ex instanceof ProcessingException) {
        ((ProcessingException) ex).printTo(messager);
      } else if (ex instanceof MultipleProcessingExceptions) {
        for (ProcessingException pex : ((MultipleProcessingExceptions) ex).getExceptions()) {
          pex.printTo(messager);
        }
      } else {
        CharSequence message = ex.getMessage();
        message = message != null ? message : ex.toString();
        Element annotatedElement = entry.getKey().getValue();
        messager.printMessage(Diagnostic.Kind.ERROR, message, annotatedElement);
        ex.printStackTrace();
      }
    }
  }

  private void process(TypeElement annotation, Element annotatedElement,
      RoundEnvironment roundEnv) {
    try {
      doProcess(annotation, annotatedElement, roundEnv);
    } catch (Exception ex) {
      Entry<TypeElement, Element> entry = Maps.immutableEntry(annotation, annotatedElement);
      retryNextRound.put(entry, ex);
    }
  }

  protected abstract void doProcess(TypeElement annotation, Element annotatedElement,
      RoundEnvironment roundEnv) throws Exception;
}

@@ -41,16 +41,66 @@
 essor;%0A%0A
+import static javax.tools.Diagnostic.Kind.ERROR;%0A%0A
 import j
@@ -458,39 +458,50 @@
 nt;%0A
+%0A
 import 
-javax.tools.Diagnostic;%0A
+com.google.common.base.Throwables;
 %0Aimp
@@ -2350,26 +2350,20 @@
     
-message = 
+if (
 message 
 != n
@@ -2362,151 +2362,287 @@
 age 
-!
+=
 = null
- ? message : ex.toString();%0A        Element annotatedElement = entry.getKey().getValue();%0A        messager.printMessage(Diagnostic.Kind.
+) %7B%0A          message = Throwables.getStackTraceAsString(ex);%0A        %7D%0A        Element annotatedElement = entry.getKey().getValue();%0A        if (annotatedElement == null) %7B%0A          messager.printMessage(ERROR, message);%0A        %7D else %7B%0A          messager.printMessage(
 ERRO
@@ -2668,24 +2668,34 @@
 edElement);%0A
+        %7D%0A
         ex.p

Improve error handling in annotation processor