gson-compile/gson-compile-processor/src/main/java/io/gitlab/jfronny/gson/compile/processor/GsonCompileProcessor.java

package io.gitlab.jfronny.gson.compile.processor;

import com.squareup.javapoet.*;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.gson.compile.annotations.GSerializable;
import io.gitlab.jfronny.gson.compile.processor.adapter.*;
import io.gitlab.jfronny.gson.compile.processor.util.AbstractProcessor2;
import io.gitlab.jfronny.gson.compile.processor.util.SupportedAnnotationTypes2;
import io.gitlab.jfronny.gson.compile.processor.util.valueprocessor.Properties;
import io.gitlab.jfronny.gson.compile.processor.util.valueprocessor.*;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.tools.Diagnostic;
import java.io.*;
import java.util.*;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({GSerializable.class})
@SupportedOptions({"gsonCompileNoReflect"})
public class GsonCompileProcessor extends AbstractProcessor2 {
    //TODO handle lists, sets, arrays, other common types (-> https://github.com/bluelinelabs/LoganSquare/blob/development/docs/TypeConverters.md)
    private static final List<Adapter> ADAPTERS = List.of(
            new DeclaredAdapter(),
            new PrimitiveAdapter(),
            new StringAdapter(),
            new OtherSerializableAdapter(),
            new ReflectAdapter()
    );

    private Messager message;
    private Filer filer;
    private Set<ClassName> seen;
    private ValueCreator valueCreator;
    private Elements elements;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        message = processingEnv.getMessager();
        filer = processingEnv.getFiler();
        elements = processingEnv.getElementUtils();
        seen = new LinkedHashSet<>();
        valueCreator = new ValueCreator(processingEnv);
        for (Adapter adapter : ADAPTERS) {
            adapter.init(processingEnv);
        }
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnvironment) {
        Set<SerializableClass> toGenerate = new LinkedHashSet<>();
        // Gather all serializable types
        for (TypeElement annotation : annotations) {
            for (Element element : roundEnvironment.getElementsAnnotatedWith(annotation)) {
                for (AnnotationMirror mirror : element.getAnnotationMirrors()) {
                    if (mirror.getAnnotationType().toString().equals(GSerializable.class.getCanonicalName())) {
                        var bld = new Object() {
                            TypeMirror with = null;
                            Boolean generateAdapter = null;
                        };
                        elements.getElementValuesWithDefaults(mirror).forEach((executableElement, value) -> {
                            String name = executableElement.getSimpleName().toString();
                            switch (name) {
                                case "with" -> {
                                    if (bld.with != null) throw new IllegalArgumentException("Duplicate annotation parameter: with");
                                    bld.with = (TypeMirror) value.getValue();
                                }
                                case "generateAdapter" -> {
                                    if (bld.generateAdapter != null) throw new IllegalArgumentException("Duplicate annotation parameter: generateAdapter");
                                    bld.generateAdapter = (Boolean) value.getValue();
                                }
                                default -> throw new IllegalArgumentException("Unexpected annotation parameter: " + name);
                            }
                        });
                        if (bld.with == null) throw new IllegalArgumentException("Missing annotation parameter: with");
                        if (bld.generateAdapter == null) throw new IllegalArgumentException("Missing annotation parameter: generateAdapter");
                        if (bld.with.toString().equals("void")) bld.with = null;
                        if (bld.with != null && bld.generateAdapter) throw new IllegalArgumentException("Adapter for " + element + " already exists, not generating another!");

                        toGenerate.add(SerializableClass.of((TypeElement) element, bld.with, bld.generateAdapter));
                    }
                }
            }
        }
        // Do not allow mutation past this point, especially not from individual process tasks
        toGenerate = Set.copyOf(toGenerate);
        // Generate adapters
        for (SerializableClass toProcess : toGenerate) {
            try {
                process(toProcess, toGenerate);
            } catch (IOException | ElementException e) {
                message.printMessage(Diagnostic.Kind.ERROR, "GsonCompile threw an exception: " + StringFormatter.toString(e), toProcess.classElement());
            }
        }
        return false;
    }

    private void process(SerializableClass toProcess, Set<SerializableClass> other) throws IOException, ElementException {
        ClassName className = toProcess.getClassName();
        if (!seen.add(className)) return; // Don't process the same class more than once

        TypeName classType = toProcess.getTypeName();
        List<TypeVariableName> typeVariables = new ArrayList<>();
        if (classType instanceof ParameterizedTypeName type) {
            for (TypeName argument : type.typeArguments) {
                typeVariables.add(TypeVariableName.get(argument.toString()));
            }
        }

        TypeSpec.Builder spec = TypeSpec.classBuilder(toProcess.generatedClassName().simpleName())
                .addOriginatingElement(toProcess.classElement())
                .addTypeVariables(typeVariables)
                .addModifiers(Modifier.PUBLIC);

        if (toProcess.adapter() != null) {
            generateDelegateToAdapter(spec, classType, toProcess.adapter());
        } else {
            if (toProcess.generateAdapter()) {
                generateDelegatingAdapter(spec, classType, toProcess.generatedClassName());
            }
            generateSerialisation(spec, classType, toProcess.classElement(), typeVariables, other);
        }

        generateAuxiliary(spec, classType);

        JavaFile javaFile = JavaFile.builder(className.packageName(), spec.build())
                .skipJavaLangImports(true)
                .indent("    ")
                .build();
        javaFile.writeTo(filer);
        message.printMessage(Diagnostic.Kind.NOTE, "Processed " + className);
    }

    private static void generateDelegatingAdapter(TypeSpec.Builder spec, TypeName classType, ClassName generatedClassName) {
        spec.addType(
                TypeSpec.classBuilder("Adapter")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .superclass(ParameterizedTypeName.get(Const.TYPE_ADAPTER, classType))
                        .addMethod(MethodSpec.methodBuilder("write")
                                .addAnnotation(Override.class)
                                .addModifiers(Modifier.PUBLIC)
                                .addParameter(Const.GSON_WRITER, "writer")
                                .addParameter(classType, "value")
                                .addException(IOException.class)
                                .addCode(generatedClassName.simpleName() + ".write(writer, value);")
                                .build())
                        .addMethod(MethodSpec.methodBuilder("read")
                                .addAnnotation(Override.class)
                                .addModifiers(Modifier.PUBLIC)
                                .addParameter(Const.GSON_READER, "reader")
                                .addException(IOException.class)
                                .returns(classType)
                                .addCode("return " + generatedClassName.simpleName() + ".read(reader);")
                                .build())
                        .build()
        );
    }

    private static void generateDelegateToAdapter(TypeSpec.Builder spec, TypeName classType, TypeMirror adapter) {
        TypeName adapterType = TypeName.get(adapter);
        spec.addField(
                FieldSpec.builder(adapterType, "ADAPTER", Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                        .initializer("new $T()", adapterType)
                        .build()
        );
        spec.addMethod(
                MethodSpec.methodBuilder(Const.READ)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Const.GSON_READER, "reader")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("return ADAPTER.read(reader);")
                        .build()
        );
        spec.addMethod(
                MethodSpec.methodBuilder(Const.READ)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Const.GSON_READER, "writer")
                        .addParameter(classType, "value")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("ADAPTER.write(reader, value);")
                        .build()
        );
    }

    private static void generateAuxiliary(TypeSpec.Builder spec, TypeName classType) {
        spec.addMethod(
                MethodSpec.methodBuilder(Const.READ)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(TypeName.get(Reader.class), "in")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("""
                                try ($T reader = $T.HOLDER.getGson().newJsonReader(in)) {
                                    return read(reader);
                                }""", Const.GSON_READER, Const.CCORE)
                        .build()
        );

        spec.addMethod(
                MethodSpec.methodBuilder(Const.READ)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(TypeName.get(String.class), "json")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("""
                                try ($1T reader = new $1T(json)) {
                                    return read(reader);
                                }""", StringReader.class)
                        .build()
        );

        spec.addMethod(
                MethodSpec.methodBuilder(Const.READ)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Const.GSON_ELEMENT, "tree")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("""
                                try ($1T reader = new $1T(tree)) {
                                    return read(reader);
                                }""", Const.GSON_TREE_READER)
                        .build()
        );

        spec.addMethod(
                MethodSpec.methodBuilder(Const.WRITE)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Writer.class, "out")
                        .addParameter(classType, "value")
                        .addException(IOException.class)
                        .addCode("""
                                try ($T writer = $T.HOLDER.getGson().newJsonWriter(out)) {
                                    write(writer, value);
                                }""", Const.GSON_WRITER, Const.CCORE)
                        .build()
        );

        spec.addMethod(
                MethodSpec.methodBuilder("toJson")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(classType, "value")
                        .addException(IOException.class)
                        .returns(String.class)
                        .addCode("""
                                try ($1T writer = new $1T()) {
                                    write(writer, value);
                                    return writer.toString();
                                }""", StringWriter.class)
                        .build()
        );

        spec.addMethod(
                MethodSpec.methodBuilder("toJsonTree")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(classType, "value")
                        .addException(IOException.class)
                        .returns(Const.GSON_ELEMENT)
                        .addCode("""
                                try ($1T writer = new $1T()) {
                                    write(writer, value);
                                    return writer.get();
                                }""", Const.GSON_TREE_WRITER)
                        .build()
        );
    }

    private void generateSerialisation(TypeSpec.Builder spec, TypeName classType, TypeElement classElement, List<TypeVariableName> typeVariables, Set<SerializableClass> otherAdapters) throws ElementException {
        Value value = valueCreator.from(classElement);
        ConstructionSource constructionSource = value.getConstructionSource();
        Properties properties = value.getProperties();

        // public static void write(JsonWriter writer, T value) throws IOException
        {
            CodeBlock.Builder code = CodeBlock.builder();
            code.beginControlFlow("if (value == null)")
                    .addStatement("writer.nullValue()")
                    .addStatement("return")
                    .endControlFlow();

            code.addStatement("writer.beginObject()");
            for (Property.Field field : properties.fields) {
                generateComments(field, code);
                code.addStatement("writer.name($S)", getSerializedName(field));
                withAdapter(field, spec, code, typeVariables, otherAdapters, adapter -> adapter.generateWrite(() -> code.add("value.$N", field.getCallableName())));
            }
            for (Property.Getter getter : properties.getters) {
                generateComments(getter, code);
                code.addStatement("writer.name($S)", getSerializedName(getter));
                withAdapter(getter, spec, code, typeVariables, otherAdapters, adapter -> adapter.generateWrite(() -> code.add("value.$N()", getter.getCallableName())));
            }
            code.addStatement("writer.endObject()");

            spec.addMethod(MethodSpec.methodBuilder("write")
                    .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                    .addParameter(Const.GSON_WRITER, "writer")
                    .addParameter(classType, "value")
                    .addException(IOException.class)
                    .addCode(code.build())
                    .build());
        }

        // public static T read(JsonReader reader) throws IOException
        {
            CodeBlock.Builder code = CodeBlock.builder();
            code.beginControlFlow("if (reader.peek() == $T.NULL)", Const.GSON_TOKEN)
                    .addStatement("reader.nextNull()")
                    .addStatement("return null")
                    .endControlFlow();

            boolean isEmpty = true;
            for (Property<?> param : properties.names) {
                isEmpty = false;
                code.addStatement("$T $L = $L", param.getType(), Const.ARG_PREFIX + param.getName(), getDefaultValue(param.getType()));
            }
            if (isEmpty) {
                code.addStatement("reader.skipValue()");
            } else {
                code.addStatement("reader.beginObject()")
                        .beginControlFlow("while (reader.hasNext())")
                        .beginControlFlow("switch (reader.nextName())");
                for (Property<?> param : properties.names) {
                    if (param.getType().getKind().isPrimitive()) {
                        code.add("case $S -> $L = ", getSerializedName(param), Const.ARG_PREFIX + param.getName());
                        withAdapter(param, spec, code, typeVariables, otherAdapters, Adapter::generateRead);
                        code.add(";\n");
                    } else {
                        code.beginControlFlow("case $S ->", getSerializedName(param))
                                .beginControlFlow("if (reader.peek() == $T.NULL)", Const.GSON_TOKEN)
                                .addStatement("reader.nextNull()")
                                .addStatement("$L = null", Const.ARG_PREFIX + param.getName());
                        code.unindent().add("} else $L = ", Const.ARG_PREFIX + param.getName());
                        withAdapter(param, spec, code, typeVariables, otherAdapters, Adapter::generateRead);
                        code.add(";\n")
                                .endControlFlow();
                    }
                }
                code.add("default -> ")
                        .addStatement("reader.skipValue()");

                code.endControlFlow()
                        .endControlFlow()
                        .addStatement("reader.endObject()");
            }

            ClassName creatorName = ClassName.get((TypeElement) constructionSource.getConstructionElement().getEnclosingElement());
            if (constructionSource instanceof ConstructionSource.Builder builder) {
                String args = properties.constructorParams.stream().map(s -> Const.ARG_PREFIX + s.getName()).collect(Collectors.joining(", "));
                if (constructionSource.isConstructor()) {
                    code.add("return new $T($L)", builder.getBuilderClass(), args);
                } else {
                    code.add("return $T.$L($L)", creatorName, classElement.getSimpleName(), args);
                }
                code.add("\n").indent();
                for (Property.BuilderParam param : properties.builderParams) {
                    code.add(".$L($L)\n", param.getCallableName(), Const.ARG_PREFIX + param.getName());
                }
                code.add(".$L();\n", builder.getBuildMethod().getSimpleName()).unindent();
            } else {
                String args = properties.params.stream().map(s -> Const.ARG_PREFIX + s.getName()).collect(Collectors.joining(", "));
                if (constructionSource.isConstructor()) {
                    code.addStatement("return new $T($L)", classType, args);
                } else {
                    code.addStatement("return $T.$L($L)", creatorName, constructionSource.getConstructionElement().getSimpleName(), args);
                }
            }

            spec.addMethod(MethodSpec.methodBuilder("read")
                    .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                    .returns(classType)
                    .addParameter(Const.GSON_READER, "reader")
                    .addException(IOException.class)
                    .addCode(code.build())
                    .build());
        }
    }

    private void generateComments(Property<?> prop, CodeBlock.Builder code) {
        for (AnnotationMirror annotation : prop.getAnnotations()) {
            if (annotation.getAnnotationType().asElement().toString().equals(Const.GCOMMENT.toString())) {
                String comment = (String) annotation.getElementValues().values().iterator().next().getValue();
                code.addStatement("writer.comment($S)", comment);
            }
        }
    }

    private void withAdapter(Property<?> prop, TypeSpec.Builder klazz, CodeBlock.Builder code, List<TypeVariableName> typeVariables, Set<SerializableClass> otherAdapters, Consumer<Adapter> action) {
        for (Adapter adapter : ADAPTERS) {
            adapter.hydrate(prop, klazz, code, typeVariables, otherAdapters);
            if (adapter.applies()) {
                action.accept(adapter);
                adapter.dehydrate();
                return;
            } else adapter.dehydrate();
        }
        message.printMessage(Diagnostic.Kind.ERROR, "Could not find applicable adapter for property " + prop);
    }

    private static String getDefaultValue(TypeMirror type) {
        return switch (type.getKind()) {
            case BYTE, SHORT, INT, LONG, FLOAT, CHAR, DOUBLE -> "0";
            case BOOLEAN -> "false";
            default -> "null";
        };
    }

    private static String getSerializedName(Property<?> property) {
        for (AnnotationMirror annotationMirror : property.getAnnotations()) {
            if (annotationMirror.getAnnotationType().asElement().toString().equals(Const.SERIALIZED_NAME.toString())) {
                return (String) annotationMirror.getElementValues().values().iterator().next().getValue();
            }
        }
        return property.getName();
    }
}