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.util.AbstractProcessor2;
import io.gitlab.jfronny.gson.compile.processor.util.SupportedAnnotationTypes2;
import io.gitlab.jfronny.gson.compile.processor.util.valueprocessor.*;
import io.gitlab.jfronny.gson.compile.processor.util.valueprocessor.Properties;
import org.jetbrains.annotations.Nullable;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.*;
import javax.lang.model.util.*;
import javax.tools.Diagnostic;
import java.io.*;
import java.lang.reflect.ParameterizedType;
import java.util.*;
import java.util.stream.Collectors;

//TODO support for records
@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({GSerializable.class})
public class GsonCompileProcessor extends AbstractProcessor2 {
    private Messager message;
    private Types typeUtils;
    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();
        typeUtils = processingEnv.getTypeUtils();
        elements = processingEnv.getElementUtils();
        seen = new LinkedHashSet<>();
        valueCreator = new ValueCreator(processingEnv);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnvironment) {
        Set<ToProcess> toProcesses = new LinkedHashSet<>();
        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!");
                        toProcesses.add(new ToProcess((TypeElement) element, bld.with, bld.generateAdapter));
                    }
                }
            }
        }
        for (ToProcess toProcess : toProcesses) {
            try {
                process(toProcess.element, toProcess.adapter, toProcess.generateAdapter);
            } catch (IOException | ElementException e) {
                message.printMessage(Diagnostic.Kind.ERROR, "GsonCompile threw an exception: " + StringFormatter.toString(e), toProcess.element);
            }
        }
        return false;
    }

    record ToProcess(TypeElement element, @Nullable TypeMirror adapter, boolean generateAdapter) {}

    private void process(TypeElement classElement, @Nullable TypeMirror adapter, boolean generateAdapter) throws IOException, ElementException {
        ClassName className = ClassName.get(classElement);
        if (!seen.add(className)) return; // Don't process the same class more than once

        ClassName generatedClassName = ClassName.get(className.packageName(), Const.PREFIX + String.join("_", className.simpleNames()));
        TypeName classType = TypeName.get(classElement.asType());
        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(generatedClassName.simpleName())
                .addOriginatingElement(classElement)
                .addTypeVariables(typeVariables)
                .addModifiers(Modifier.PUBLIC);

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

        generateAuxiliary(spec, classType);

        //TODO register adapter as available and use in generated adapters

        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 ($1T reader = new $1T(in)) {
                                    return read(reader);
                                }""", Const.GSON_READER)
                        .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 ($1T writer = new $1T(out)) {
                                    write(writer, value);
                                }""", Const.GSON_WRITER)
                        .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) 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) {
                code.addStatement("writer.name($S)", getSerializedName(field));
                generateWrite(field, spec, code, "writer", "value.$N", typeVariables);
            }
            for (Property.Getter getter : properties.getters) {
                code.addStatement("writer.name($S)", getSerializedName(getter));
                generateWrite(getter, spec, code, "writer", "value.$N()", typeVariables);
            }
            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) {
                    String read = generateRead(param, spec, "reader", typeVariables);
                    if (param.getType().getKind().isPrimitive()) {
                        code.add("case $S -> ", getSerializedName(param));
                        code.addStatement("$L = $L", Const.ARG_PREFIX + param.getName(), read);
                    } 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())
                                .endControlFlow("else $L = $L", Const.ARG_PREFIX + param.getName(), read)
                                .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 generateWrite(Property<?> prop, TypeSpec.Builder klazz, CodeBlock.Builder code, String writer, String get, List<TypeVariableName> typeVariables) {
        TypeMirror unboxed = unbox(prop.getType());
        if (unbox(prop.getType()).getKind().isPrimitive() || prop.getType().toString().equals(String.class.getCanonicalName())) {
            if (prop.getType().equals(unboxed)) {
                code.addStatement("$L.value(" + get + ")", writer, prop.getCallableName());
            } else {
                String pName = Const.ARG_PREFIX + prop.getName();
                code.addStatement("$T $L = " + get, prop.getType(), pName, prop.getCallableName())
                        .beginControlFlow("if ($L == null)", pName)
                        .addStatement("$L.nullValue()", writer)
                        .endControlFlow("else $L.value($L)", writer, pName);
            }
            return;
        }
        //TODO handle lists, sets, arrays, other common types (-> https://github.com/bluelinelabs/LoganSquare/blob/development/docs/TypeConverters.md)
        //TODO support comment annotations
        code.addStatement(getAdapter(prop, klazz, typeVariables) + ".write(" + writer + ", " + get + ")", prop.getCallableName());
    }

    private String generateRead(Property<?> prop, TypeSpec.Builder klazz, String reader, List<TypeVariableName> typeVariables) {
        TypeMirror unboxed = unbox(prop.getType());
        if (unboxed.getKind().isPrimitive()) {
            return switch (unboxed.getKind()) {
                case BOOLEAN -> reader + ".nextBoolean()";
                case BYTE -> "(byte) " + reader + ".nextInt()";
                case SHORT -> "(short) " + reader + ".nextInt()";
                case INT -> reader + ".nextInt()";
                case LONG -> reader + ".nextLong()";
                case CHAR -> "(char) " + reader + ".nextInt()";
                case FLOAT -> "(float) " + reader + ".nextDouble()";
                case DOUBLE -> reader + ".nextDouble()";
                default -> throw new IllegalArgumentException("Unsupported primitive: " + unboxed.getKind());
            };
        }
        if (prop.getType().toString().equals(String.class.getCanonicalName())) {
            return reader + ".nextString()";
        }
        //TODO handle lists, sets, arrays, other common types (-> https://github.com/bluelinelabs/LoganSquare/blob/development/docs/TypeConverters.md)
        return getAdapter(prop, klazz, typeVariables) + ".read(" + reader + ")";
    }

    private String getAdapter(Property<?> prop, TypeSpec.Builder klazz, List<TypeVariableName> typeVariables) {
        String typeAdapterName = Const.ADAPTER_PREFIX + prop.getName();
        for (FieldSpec spec : klazz.fieldSpecs) {
            if (spec.name.equals(typeAdapterName)) return typeAdapterName;
        }
        DeclaredType typeAdapterClass = findTypeAdapterClass(prop.getAnnotations());
        if (typeAdapterClass != null) {
            if (isInstance(typeAdapterClass, Const.TYPE_ADAPTER.toString())) {
                klazz.addField(
                        FieldSpec.builder(TypeName.get(typeAdapterClass), typeAdapterName)
                                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                                .initializer("new $T()", typeAdapterClass)
                                .build()
                );
            } else if (isInstance(typeAdapterClass, Const.TYPE_ADAPTER_FACTORY.toString())) {
                TypeName typeAdapterType = ParameterizedTypeName.get(Const.TYPE_ADAPTER, TypeName.get(prop.getType()).box());
                CodeBlock.Builder block = CodeBlock.builder();
                block.add("new $T().create($T.HOLDER.getGson(), ", typeAdapterClass, Const.CCORE);
                appendFieldTypeToken(block, prop, typeVariables, false);
                block.add(")");
                klazz.addField(
                        FieldSpec.builder(typeAdapterType, typeAdapterName)
                                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                                .initializer(block.build())
                                .build()
                );
            } else message.printMessage(Diagnostic.Kind.ERROR, "@JsonAdapter value must by TypeAdapter or TypeAdapterFactory reference.", prop.getElement());
        } else {
            //TODO handle known custom type adapters and return proper static class
            TypeName typeAdapterType = ParameterizedTypeName.get(Const.TYPE_ADAPTER, TypeName.get(prop.getType()).box());
            CodeBlock.Builder block = CodeBlock.builder();
            block.add("$T.HOLDER.getGson().getAdapter(", Const.CCORE);
            appendFieldTypeToken(block, prop, typeVariables, true);
            block.add(")");
            klazz.addField(
                    FieldSpec.builder(typeAdapterType, typeAdapterName)
                            .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                            .initializer(block.build())
                            .build()
            );
        }
        return typeAdapterName;
    }

    private void appendFieldTypeToken(CodeBlock.Builder block, Property<?> property, List<TypeVariableName> typeVariables, boolean allowClassType) {
        TypeMirror type = property.getType();
        TypeName typeName = TypeName.get(type);

        if (isComplexType(type)) {
            TypeName typeTokenType = ParameterizedTypeName.get(Const.TYPE_TOKEN, typeName);
            List<? extends TypeMirror> typeParams = getGenericTypes(type);
            if (typeParams.isEmpty()) {
                block.add("new $T() {}", typeTokenType);
            } else {
                block.add("($T) $T.getParameterized($T.class, ", typeTokenType, Const.TYPE_TOKEN, typeUtils.erasure(type));
                for (Iterator<? extends TypeMirror> iterator = typeParams.iterator(); iterator.hasNext(); ) {
                    TypeMirror typeParam = iterator.next();
                    int typeIndex = typeVariables.indexOf(TypeVariableName.get(typeParam.toString()));
                    block.add("(($T)typeToken.getType()).getActualTypeArguments()[$L]", ParameterizedType.class, typeIndex);
                    if (iterator.hasNext()) {
                        block.add(", ");
                    }
                }
                block.add(")");
            }
        } else if (isGenericType(type)) {
            TypeName typeTokenType = ParameterizedTypeName.get(Const.TYPE_TOKEN, typeName);
            int typeIndex = typeVariables.indexOf(TypeVariableName.get(property.getType().toString()));
            block.add("($T) $T.get((($T)typeToken.getType()).getActualTypeArguments()[$L])",
                    typeTokenType, Const.TYPE_TOKEN, ParameterizedType.class, typeIndex);
        } else {
            if (allowClassType) {
                block.add("$T.class", typeName);
            } else {
                block.add("TypeToken.get($T.class)", typeName);
            }
        }
    }

    private boolean isComplexType(TypeMirror type) {
        Element element = typeUtils.asElement(type);
        if (!(element instanceof TypeElement typeElement)) return false;
        return !typeElement.getTypeParameters().isEmpty();
    }

    private boolean isGenericType(TypeMirror type) {
        return type.getKind() == TypeKind.TYPEVAR;
    }

    private List<? extends TypeMirror> getGenericTypes(TypeMirror type) {
        DeclaredType declaredType = asDeclaredType(type);
        if (declaredType == null) {
            return Collections.emptyList();
        }
        ArrayList<TypeMirror> result = new ArrayList<>();
        for (TypeMirror argType : declaredType.getTypeArguments()) {
            if (argType.getKind() == TypeKind.TYPEVAR) {
                result.add(argType);
            }
        }
        return result;
    }

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

    private TypeMirror unbox(TypeMirror type) {
        try {
            return typeUtils.unboxedType(type);
        } catch (IllegalArgumentException e) {
            return type;
        }
    }

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

    private DeclaredType findTypeAdapterClass(List<? extends AnnotationMirror> annotations) {
        for (AnnotationMirror annotation : annotations) {
            String typeName = annotation.getAnnotationType().toString();
            if (typeName.equals(Const.JSON_ADAPTER.toString())) {
                Map<? extends ExecutableElement, ? extends AnnotationValue> elements = annotation.getElementValues();
                if (!elements.isEmpty()) {
                    AnnotationValue value = elements.values().iterator().next();
                    return (DeclaredType) value.getValue();
                }
            }
        }
        return null;
    }

    private boolean isInstance(DeclaredType type, String parentClassName) {
        if (type == null) return false;
        TypeElement element = (TypeElement) type.asElement();
        for (TypeMirror interfaceType : element.getInterfaces()) {
            if (typeUtils.erasure(interfaceType).toString().equals(parentClassName)) return true;
        }
        TypeMirror superclassType = element.getSuperclass();
        if (superclassType != null) {
            if (typeUtils.erasure(superclassType).toString().equals(parentClassName)) {
                return true;
            } else {
                return isInstance(asDeclaredType(superclassType), parentClassName);
            }
        }
        return false;
    }

    private static DeclaredType asDeclaredType(TypeMirror type) {
        return type.accept(new SimpleTypeVisitor14<>() {
            @Override
            public DeclaredType visitDeclared(DeclaredType t, Object o) {
                return t;
            }
        }, null);
    }
}