package me.xdrop.jrand.annotation.processing;

import com.github.javaparser.ast.CompilationUnit;
import me.xdrop.jrand.annotation.PropertyFlag;

import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.tools.Diagnostic;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Set;

@SupportedAnnotationTypes("me.xdrop.jrand.annotation.PropertyFlag")
@SupportedSourceVersion(SourceVersion.RELEASE_7)
public class PropertyFlagProcessor extends BaseProcessor {
    private PropertyMethodGenerator propertyMethodGenerator;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        propertyMethodGenerator = new PropertyMethodGenerator(processingEnv, getRepository());
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        for (Element element : roundEnv.getElementsAnnotatedWith(PropertyFlag.class)) {
            if (element.getKind() != ElementKind.FIELD) {
                getMessager().printMessage(Diagnostic.Kind.ERROR, "Annotation [PropertyFlag] should only be used on fields.");
                return true;
            }
            VariableElement symbol = (VariableElement) element;
            PropertyFlag annotation = symbol.getAnnotation(PropertyFlag.class);
            PackageElement pkg = processingEnv.getElementUtils().getPackageOf(symbol);
            String className = symbol.getEnclosingElement().getSimpleName().toString();
            CompilationUnit compilationUnit = getRepository().getCU(pkg.toString(), className);
            String fullPackage = pkg.getQualifiedName().toString();
            CompilationUnit newCU = propertyMethodGenerator.buildMethod(compilationUnit, symbol,
                    className, fullPackage,annotation.value());
            // Get the last part of the package eg. me.xdrop.generators.text would be text
            String lastPackage = getLastPackageName(pkg);
            Path path = Paths.get(getOutputPathGenerators().toString(), lastPackage);

            try {
                getRepository().writeTo(path, className, newCU);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return false;
    }
}

@@ -90,16 +90,124 @@
 onUnit;%0A
+import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;%0Aimport com.squareup.javapoet.MethodSpec;%0A
 import m
@@ -682,19 +682,17 @@
 va.util.
-Set
+*
 ;%0A%0A@Supp
@@ -922,16 +922,101 @@
 nerator;
+%0A    private Map%3CClassIdentifier, List%3CMethodSpec%3E%3E methodPerClass = new HashMap%3C%3E();
 %0A%0A    @O
@@ -1339,24 +1339,101 @@
 roundEnv) %7B%0A
+%0A        if (roundEnv.processingOver())%7B%0A            return true;%0A        %7D%0A%0A
         for 
@@ -1879,16 +1879,72 @@
 .class);
+%0A%0A            System.out.println(annotation.toString());
 %0A       
@@ -2140,30 +2140,17 @@
 onUnit c
-ompilationUnit
+u
  = getRe
@@ -2194,16 +2194,17 @@
 sName);%0A
+%0A
         
@@ -2211,64 +2211,173 @@
     
-String fullPackage = pkg.getQualifiedName().toString(
+ClassIdentifier id = new ClassIdentifier(className, pkg);%0A            List%3CMethodSpec%3E methodsInClass = methodPerClass.computeIfAbsent(id, k -%3E new ArrayList%3C%3E()
 );%0A
+%0A
     
@@ -2388,32 +2388,30 @@
     
-CompilationUnit newCU = 
+methodsInClass.addAll(
 prop
@@ -2445,24 +2445,12 @@
 thod
+s
 (c
-ompilationUnit
+u
 , sy
@@ -2479,49 +2479,202 @@
     
-className, fullPackage,annotation.value()
+annotation.value()));%0A        %7D%0A%0A%0A        methodPerClass.forEach((k,v) -%3E %7B%0A            PackageElement pkg = k.getPackageName();%0A            String fullPackage = pkg.getQualifiedName().toString(
 );%0A 
@@ -2903,24 +2903,344 @@
 stPackage);%0A
+            String className = k.getClassName();%0A            CompilationUnit cu = getRepository().getCU(pkg.toString(), className);%0A            List%3CMethodSpec%3E methodsToAdd = v;%0A            System.out.println(%22Writing...%22  + className);%0A%0A            getRepository().addMethods(cu, className, fullPackage, methodsToAdd);
 %0A           
@@ -3307,13 +3307,10 @@
 me, 
-newCU
+cu
 );%0A 
@@ -3398,32 +3398,38 @@
      %7D%0A        %7D
+);%0A%0A%0A%0A
 %0A        return 

Change the way Property processor works