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.Adapter;
import io.gitlab.jfronny.gson.compile.processor.adapter.Adapters;
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.nio.file.*;
import java.util.*;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;

@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({GSerializable.class})
@SupportedOptions({"gsonCompileNoReflect"})
public class GsonCompileProcessor extends AbstractProcessor2 {
    private Messager message;
    private Filer filer;
    private Map<ClassName, TypeSpec.Builder> seen;
    private ValueCreator valueCreator;
    private Elements elements;
    private boolean hasManifold = false;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        try {
            Class.forName("manifold.ext.Model");
            System.out.println("Detected manifold!");
            hasManifold = true;
        } catch (ClassNotFoundException e) {
            hasManifold = false;
        }
        message = processingEnv.getMessager();
        filer = processingEnv.getFiler();
        elements = processingEnv.getElementUtils();
        seen = new LinkedHashMap<>();
        valueCreator = new ValueCreator(processingEnv);
        for (Adapter adapter : Adapters.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()) {
                    try {
                        if (mirror.getAnnotationType().toString().equals(GSerializable.class.getCanonicalName())) {
                            var bld = new Object() {
                                TypeMirror with = null;
                                TypeMirror builder = null;
                                TypeMirror configure = 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 "builder" -> {
                                        if (bld.builder != null) throw new IllegalArgumentException("Duplicate annotation parameter: builder");
                                        bld.builder = (TypeMirror) value.getValue();
                                    }
                                    case "configure" -> {
                                        if (bld.configure != null) throw new IllegalArgumentException("Duplicate annotation parameter: configure");
                                        bld.configure = (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.builder == null) throw new IllegalArgumentException("Missing annotation parameter: builder");
                            if (bld.configure == null) throw new IllegalArgumentException("Missing annotation parameter: configure");
                            if (bld.generateAdapter == null) throw new IllegalArgumentException("Missing annotation parameter: generateAdapter");

                            toGenerate.add(SerializableClass.of((TypeElement) element, bld.with, bld.builder, bld.configure, bld.generateAdapter, hasManifold));
                        }
                    } catch (ElementException e) {
                        e.printMessage(message);
                    }
                }
            }
        }
        // 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 (ElementException e) {
                e.printMessage(message);
            }
        }
        for (var entry : seen.keySet().stream().collect(Collectors.groupingBy(ClassName::packageName)).entrySet()) {
            Map<List<String>, TypeSpec.Builder> known = entry.getValue().stream().collect(Collectors.toMap(ClassName::simpleNames, seen::get));
            // Generate additional parent classes
            for (List<String> klazz : known.keySet().stream().toList()) {
                List<String> current = new LinkedList<>();
                for (String s : klazz) {
                    current.add(s);
                    TypeSpec.Builder builder = find(known, current);
                    if (builder == null) {
                        builder = TypeSpec.classBuilder(s).addModifiers(Modifier.PUBLIC);
                        if (current.size() == 1 && hasManifold) builder.addAnnotation(Cl.MANIFOLD_EXTENSION);
                        known.put(List.copyOf(current), builder);
                    }
                    if (current.size() > 1) builder.addModifiers(Modifier.STATIC);
                }
            }
            // Add to parent class
            for (var entry1 : known.entrySet()) {
                if (entry1.getKey().size() == 1) continue;
                find(known, entry1.getKey().subList(0, entry1.getKey().size() - 1)).addType(entry1.getValue().build());
            }
            // Print
//            System.out.println("Got " + known.size() + " classes");
//            for (var entry1 : known.entrySet()) {
//                System.out.println("Class " + entry.getKey() + '.' + String.join(".", entry1.getKey()));
//                for (TypeSpec typeSpec : entry1.getValue().typeSpecs) {
//                    System.out.println("- " + typeSpec.name);
//                }
//            }
            // Write top-level classes
            for (var entry1 : known.entrySet()) {
                if (entry1.getKey().size() == 1) {
                    JavaFile javaFile = JavaFile.builder(entry.getKey(), entry1.getValue().build())
                            .skipJavaLangImports(true)
                            .indent("    ")
                            .build();
                    try {
                        javaFile.writeTo(filer);
                    } catch (IOException e) {
                        message.printMessage(Diagnostic.Kind.ERROR, "Could not write source: " + StringFormatter.toString(e));
                    }
                }
            }
        }
        seen.clear();
        return false;
    }

    private <K, V> V find(Map<List<K>, V> map, List<K> key) {
        for (var entry : map.entrySet()) if (entry.getKey().equals(key)) return entry.getValue();
        return null;
    }

    private void process(SerializableClass toProcess, Set<SerializableClass> other) throws ElementException {
        if (seen.containsKey(toProcess.generatedClassName())) 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())
                .addModifiers(Modifier.PUBLIC)
                .addTypeVariables(typeVariables);

        seen.put(toProcess.generatedClassName(), spec);

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

        generateAuxiliary(spec, classType, toProcess.configure());
    }

    private static void generateDelegatingAdapter(TypeSpec.Builder spec, TypeName classType, ClassName generatedClassName) {
        spec.addType(
                TypeSpec.classBuilder("Adapter")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .superclass(ParameterizedTypeName.get(Cl.TYPE_ADAPTER, classType))
                        .addMethod(MethodSpec.methodBuilder("write")
                                .addAnnotation(Override.class)
                                .addModifiers(Modifier.PUBLIC)
                                .addParameter(Cl.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(Cl.GSON_READER, "reader")
                                .addException(IOException.class)
                                .returns(classType)
                                .addCode("return " + generatedClassName.simpleName() + ".read(reader);")
                                .build())
                        .build()
        );
    }

    private void generateDelegateToAdapter(TypeSpec.Builder spec, TypeName classType, TypeMirror adapter) {
        spec.addMethod(
                extension(MethodSpec.methodBuilder("read"))
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Cl.GSON_READER, "reader")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode("return $T.read(reader);", adapter)
                        .build()
        );
        spec.addMethod(
                extension(MethodSpec.methodBuilder("write"), classType, Cl.GSON_WRITER, "writer")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addException(IOException.class)
                        .addCode("$T.write(writer, value);", adapter)
                        .build()
        );
    }

    private void generateAuxiliary(TypeSpec.Builder spec, TypeName classType, TypeMirror configure) {
        final UnaryOperator<CodeBlock.Builder> configureReader = cb -> {
            if (configure != null) cb.addStatement("$T.configure(reader)", configure);
            return cb;
        };
        final UnaryOperator<CodeBlock.Builder> configureWriter = cb -> {
            if (configure != null) cb.addStatement("$T.configure(writer)", configure);
            return cb;
        };

        final String readStatement = "return read(reader)";
        final String writeStatement = hasManifold ? "write(value, writer)" : "write(writer, value)";

        spec.addMethod(
                extension(MethodSpec.methodBuilder("read"))
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(TypeName.get(Reader.class), "in")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode(configureReader.apply(CodeBlock.builder().beginControlFlow("try ($1T reader = new $1T(in))", Cl.GSON_READER))
                                .addStatement(readStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("read"))
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(TypeName.get(String.class), "json")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode(CodeBlock.builder().beginControlFlow("try ($1T reader = new $1T(json))", StringReader.class)
                                .addStatement(readStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("read"))
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Cl.GSON_ELEMENT, "tree")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode(configureReader.apply(CodeBlock.builder().beginControlFlow("try ($1T reader = new $1T(tree))", Cl.GSON_TREE_READER))
                                .addStatement(readStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("read"))
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addParameter(Path.class, "path")
                        .addException(IOException.class)
                        .returns(classType)
                        .addCode(CodeBlock.builder().beginControlFlow("try ($T reader = $T.newBufferedReader(path))", BufferedReader.class, Files.class)
                                .addStatement(readStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("write"), classType, Writer.class, "out")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addException(IOException.class)
                        .addCode(configureWriter.apply(CodeBlock.builder().beginControlFlow("try ($1T writer = new $1T(out))", Cl.GSON_WRITER))
                                .addStatement(writeStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("write"), classType, Path.class, "path")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addException(IOException.class)
                        .addCode(CodeBlock.builder().beginControlFlow("try ($1T writer = $2T.newBufferedWriter(path, $3T.CREATE, $3T.WRITE, $3T.TRUNCATE_EXISTING))", BufferedWriter.class, Files.class, StandardOpenOption.class)
                                .addStatement(writeStatement)
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("toJson"), classType)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addException(IOException.class)
                        .returns(String.class)
                        .addCode(CodeBlock.builder().beginControlFlow("try ($1T writer = new $1T())", StringWriter.class)
                                .addStatement(writeStatement)
                                .addStatement("return writer.toString()")
                                .endControlFlow()
                                .build())
                        .build()
        );

        spec.addMethod(
                extension(MethodSpec.methodBuilder("toJsonTree"), classType)
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .addException(IOException.class)
                        .returns(Cl.GSON_ELEMENT)
                        .addCode(configureWriter.apply(CodeBlock.builder().beginControlFlow("try ($1T writer = new $1T())", Cl.GSON_TREE_WRITER))
                                .addStatement(writeStatement)
                                .addStatement("return writer.get()")
                                .endControlFlow()
                                .build())
                        .build()
        );
    }

    private void generateSerialisation(TypeSpec.Builder spec, SerializableClass self, List<TypeVariableName> typeVariables, Set<SerializableClass> otherAdapters) throws ElementException {
        Value value = self.builder() == null ? valueCreator.from(self.classElement(), false) : valueCreator.from(TypeHelper.asDeclaredType(self.builder()).asElement(), true);
        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 param : properties.fields) {
                if (Properties.containsName(properties.getters, param)) continue;
                Runnable writeGet = () -> code.add("value.$N", param.getCallableName());
                if (param.getType().getKind().isPrimitive()) {
                    generateComments(param, code);
                    code.addStatement("writer.name($S)", getSerializedName(param));
                    Adapters.generateWrite(param, spec, code, typeVariables, otherAdapters, message, writeGet);
                } else {
                    code.beginControlFlow("if (value.$N != null || writer.getSerializeNulls())", param.getCallableName());
                    generateComments(param, code);
                    code.addStatement("writer.name($S)", getSerializedName(param));
                    code.addStatement("if (value.$N == null) writer.nullValue()", param.getCallableName());
                    code.beginControlFlow("else");
                    Adapters.generateWrite(param, spec, code, typeVariables, otherAdapters, message, writeGet);
                    code.endControlFlow();
                    code.endControlFlow();
                }
            }
            for (Property.Getter param : properties.getters) {
                if (param.getType().getKind().isPrimitive()) {
                    generateComments(param, code);
                    code.addStatement("writer.name($S)", getSerializedName(param));
                    Adapters.generateWrite(param, spec, code, typeVariables, otherAdapters, message, () -> code.add("value.$N()", param.getCallableName()));
                } else {
                    code.addStatement("$T $L$N = value.$N()", param.getType(), "$", param.getCallableName(), param.getCallableName());
                    code.beginControlFlow("if ($L$N != null || writer.getSerializeNulls())", "$", param.getCallableName());
                    generateComments(param, code);
                    code.addStatement("writer.name($S)", getSerializedName(param));
                    code.addStatement("if ($L$N == null) writer.nullValue()", "$", param.getCallableName());
                    code.beginControlFlow("else");
                    Adapters.generateWrite(param, spec, code, typeVariables, otherAdapters, message, () -> code.add("$L$N", "$", param.getCallableName()));
                    code.endControlFlow();
                    code.endControlFlow();
                }
            }
            code.addStatement("writer.endObject()");

            spec.addMethod(extension(MethodSpec.methodBuilder("write"), self.getTypeName(), Cl.GSON_WRITER, "writer")
                    .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                    .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)", Cl.GSON_TOKEN)
                    .addStatement("reader.nextNull()")
                    .addStatement("return null")
                    .endControlFlow();

            boolean isEmpty = true;
            for (Property<?> param : properties.names) {
                isEmpty = false;
                code.addStatement("$T _$N = $L", param.getType(), param.getName(), TypeHelper.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 -> _$N = ", getSerializedName(param), param.getName());
                        Adapters.generateRead(param, spec, code, typeVariables, otherAdapters, message);
                        code.add(";\n");
                    } else {
                        code.beginControlFlow("case $S ->", getSerializedName(param))
                                .beginControlFlow("if (reader.peek() == $T.NULL)", Cl.GSON_TOKEN)
                                .addStatement("reader.nextNull()")
                                .addStatement("_$N = null", param.getName());
                        code.unindent().add("} else _$N = ", param.getName());
                        Adapters.generateRead(param, spec, code, typeVariables, otherAdapters, message);
                        code.add(";\n")
                                .endControlFlow();
                    }
                }
                code.add("default -> ")
                        .addStatement("reader.skipValue()");

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

            code.addStatement("$T result", self.getTypeName());
            ClassName creatorName = ClassName.get((TypeElement) constructionSource.getConstructionElement().getEnclosingElement());
            if (constructionSource instanceof ConstructionSource.Builder builder) {
                StringBuilder args = new StringBuilder();
                for (Property.ConstructorParam param : properties.constructorParams) {
                    args.append(", _").append(param.getName());
                }
                code.add("$T builder = ", builder.getBuilderClass());
                if (constructionSource.isConstructor()) {
                    code.add("new $T($L)", builder.getBuilderClass(), args.length() > 0 ? args.substring(2) : "");
                } else {
                    code.add("$T.$N($L)", creatorName, self.classElement().getSimpleName(), args.length() > 0 ? args.substring(2) : "");
                }
                code.add(";\n");
                for (Property.Setter param : properties.builderParams) {
                    code.addStatement("builder.$N(_$N)", param.getCallableName(), param.getName());
                }
                code.addStatement("result = builder.$N()", builder.getBuildMethod().getSimpleName());
            } else {
                StringBuilder args = new StringBuilder();
                for (Property.Param param : properties.params) {
                    args.append(", _").append(param.getName());
                }
                if (constructionSource.isConstructor()) {
                    code.addStatement("result = new $T($L)", self.getTypeName(), args.length() > 0 ? args.substring(2) : "");
                } else {
                    code.addStatement("result = $T.$N($L)", creatorName, constructionSource.getConstructionElement().getSimpleName(), args.length() > 0 ? args.substring(2) : "");
                }
            }
            for (Property.Setter setter : properties.setters) {
                code.addStatement("result.$N(_$N)", setter.getCallableName(), setter.getName());
            }
            for (Property.Field field : properties.fields) {
                if (Properties.containsName(properties.setters, field)) continue;
                if (Properties.containsName(properties.params, field)) continue;
                code.addStatement("result.$N = _$N", field.getName(), field.getCallableName());
            }
            code.addStatement("return result");

            spec.addMethod(extension(MethodSpec.methodBuilder("read"))
                    .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                    .returns(self.getTypeName())
                    .addParameter(Cl.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(Cl.GCOMMENT.toString())) {
                String comment = (String) annotation.getElementValues().values().iterator().next().getValue();
                code.addStatement("if (writer.isLenient()) writer.comment($S)", comment);
            }
        }
    }

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

    private MethodSpec.Builder extension(MethodSpec.Builder method) {
        if (hasManifold) method.addAnnotation(Cl.MANIFOLD_EXTENSION);
        return method;
    }

    private MethodSpec.Builder extension(MethodSpec.Builder method, TypeName thizName) {
        if (hasManifold) {
            method.addAnnotation(Cl.MANIFOLD_EXTENSION);
            method.addParameter(ParameterSpec.builder(thizName, "value").addAnnotation(Cl.MANIFOLD_THIS).build());
        }
        else {
            method.addParameter(thizName, "value");
        }
        return method;
    }

    private MethodSpec.Builder extension(MethodSpec.Builder method, TypeName thizName, TypeName otherName, String other) {
        if (hasManifold) {
            method.addAnnotation(Cl.MANIFOLD_EXTENSION);
            method.addParameter(ParameterSpec.builder(thizName, "value").addAnnotation(Cl.MANIFOLD_THIS).build())
                    .addParameter(otherName, other);
        }
        else {
            method.addParameter(otherName, other)
                    .addParameter(thizName, "value");
        }
        return method;
    }

    private MethodSpec.Builder extension(MethodSpec.Builder method, TypeName thizName, Class<?> otherName, String other) {
        return extension(method, thizName, TypeName.get(otherName), other);
    }
}