package ru.ifmo.java.serialization.processor;


import com.squareup.javapoet.*;
import ru.ifmo.java.serialization.Letter;
import ru.ifmo.java.serialization.Letterize;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import javax.tools.Diagnostic;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.function.Consumer;
import java.util.function.Function;

@SupportedSourceVersion(SourceVersion.RELEASE_11)
@SupportedAnnotationTypes({"ru.ifmo.java.serialization.Letterize", "ru.ifmo.java.serialization.LetterizeOptional"})
public class AutoSerializationProcessor extends AbstractProcessor {
    private Types typeUtils;
    private Elements elementUtils;
    private Filer filer;
    private Messager messager;

    final TypeName wildcard = WildcardTypeName.subtypeOf(Object.class);
    final TypeName classOfAny = ParameterizedTypeName.get(
            ClassName.get(Class.class), wildcard);

    final TypeName typeMapForSerialized = ParameterizedTypeName.get(
            ClassName.get(HashMap.class), TypeName.get(Class.class), TypeName.get(Consumer.class)
    );

    final TypeName typeMapForDeserialized = ParameterizedTypeName.get(
            ClassName.get(HashMap.class), TypeName.get(Class.class), TypeName.get(Callable.class)
    );

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        typeUtils = processingEnv.getTypeUtils();
        elementUtils = processingEnv.getElementUtils();
        filer = processingEnv.getFiler();
        messager = processingEnv.getMessager();
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        List<MethodSpec> serilalizeMethods = new ArrayList<>();
        List<MethodSpec> deserilalizeMethods = new ArrayList<>();

        for (Element annotatedElement : roundEnv.getElementsAnnotatedWith(Letterize.class)) {
            // Check if a class has been annotated with @Factory
            if (annotatedElement.getKind() != ElementKind.CLASS) {
                error(annotatedElement, "Only classes can be annotated with @%s", Letterize.class.getSimpleName());
                return true; // Exit processing
            }
            TypeElement typeElement = (TypeElement) annotatedElement;
            if (!checkLetterImplementation(typeElement)) {
                error(annotatedElement, "Only classes can be interface with @%s", Letter.class.getSimpleName());
            }

            serilalizeMethods.add(getSerializeMethod(typeElement));
            deserilalizeMethods.add(getDeserializeMethod(typeElement));

        }

        if (!annotations.isEmpty()) {
            generateSerializer(serilalizeMethods);
            generateDeserializer(deserilalizeMethods);
        }

        return false;
    }

    private MethodSpec getDeserializeMethod(TypeElement typeElement) {
        String methodName = "deserialize" + typeElement.getSimpleName();
        MethodSpec.Builder builder = MethodSpec.methodBuilder(methodName)
                .addException(IOException.class);

        TypeMirror typeMirror = typeElement.asType();
        TypeName typeName = TypeName.get(typeMirror);
        builder.returns(typeName);

        builder.addStatement("$N element = new $N()",
                typeElement.getSimpleName(),
                typeElement.getSimpleName());

        List<VariableElement> typeParameters = getTypeParameters(typeElement);
        typeParameters.forEach((x) -> addToMethodDeserializer(builder, x));

        builder.addStatement("return element");
        return builder.build();
    }


    private MethodSpec getSerializeMethod(TypeElement typeElement) {
        String methodName = "serialize" + typeElement.getSimpleName();
        MethodSpec.Builder builder = MethodSpec.methodBuilder(methodName)
                .addException(IOException.class);

        TypeMirror typeMirror = typeElement.asType();
        TypeName typeName = TypeName.get(typeMirror);

        builder.addParameter(typeName, "element");

        List<VariableElement> typeParameters = getTypeParameters(typeElement);
        typeParameters.forEach((x) -> addToMethodSerializer(builder, x));
        return builder.build();
    }

    private List<VariableElement> getTypeParameters(TypeElement typeElement) {
        TypeElement objectType = elementUtils.getTypeElement(Object.class.getCanonicalName());
        List<VariableElement> res = new ArrayList<>();
        while (!typeElement.equals(objectType)) {
            for (Element element : typeElement.getEnclosedElements()) {
                if (element.getKind().equals(ElementKind.FIELD)) {
                    VariableElement variableElement = (VariableElement) element;
                    res.add(variableElement);
                }
            }
            typeElement = (TypeElement) typeUtils.asElement(typeElement.getSuperclass());
        }
        return res;
    }

    private void addToMethodSerializer(MethodSpec.Builder builder, VariableElement variableElement) {
        TypeMirror typeMirror = variableElement.asType();

        Name nameVariable = variableElement.getSimpleName();
        String typeVariable = getTypeVariable(variableElement);
        boolean isString = TypeName.get(typeMirror).equals(TypeName.get(String.class));
        if (typeMirror.getKind().equals(TypeKind.DECLARED) && !isString) {
            builder.addStatement("serialize$N(element.$N)", typeVariable, nameVariable);
        } else{
            builder.addStatement("output.write$N(element.$N)",typeVariable, nameVariable);
        }

    }

    private void addToMethodDeserializer(MethodSpec.Builder builder, VariableElement variableElement) {
        TypeMirror typeMirror = variableElement.asType();

        Name nameVariable = variableElement.getSimpleName();
        String typeVariable = getTypeVariable(variableElement);
        boolean isString = TypeName.get(typeMirror).equals(TypeName.get(String.class));
        if (typeMirror.getKind().equals(TypeKind.DECLARED) && !isString) {
            builder.addStatement("element.$N=deserialize$N()", nameVariable, typeVariable);
        } else{
            builder.addStatement("element.$N=input.read$N()", nameVariable, typeVariable);
        }
    }

    private String getTypeVariable(VariableElement variableElement) {
        TypeMirror typeMirror = variableElement.asType();
        TypeElement typeElement = (TypeElement) typeUtils.asElement(typeMirror);

        if (typeMirror.getKind().equals(TypeKind.DECLARED)) {
            boolean isString = TypeName.get(typeMirror).equals(TypeName.get(String.class));
            if (isString) {
                return "UTF";
            } else {
                return typeElement.getSimpleName().toString();
            }
        } else {
            Function<String, String> firstLetterUp = (String name) -> name.substring(0, 1).toUpperCase() + name.substring(1);
            return firstLetterUp.apply(typeMirror.toString());
        }
    }

    private void generateDeserializer(List<MethodSpec> methods) {

        MethodSpec constructor = getConstructorDeserialized(methods);

        TypeSpec.Builder builder = TypeSpec.classBuilder("Deserializer")
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .addField(DataInputStream.class, "input",
                        Modifier.PRIVATE, Modifier.FINAL)
                .addMethod(constructor);

        builder.addMethods(methods);

        builder.addField(FieldSpec.builder(typeMapForDeserialized, "typeToMethod").initializer("new HashMap<>()")
                .addModifiers(Modifier.FINAL, Modifier.PRIVATE).build());

        String code = "try {\n" +
                "return typeToMethod.get(clazz).call(); \n" +
                "catch (Exception e) {\n" +
                "throw new RuntimeException().addSuspend(e)}\n return null";

        MethodSpec serialize = MethodSpec.methodBuilder("deserialize")
                .addParameter(classOfAny, "clazz")
                .returns(Object.class)
                .addStatement(code).build();

        builder.addMethod(serialize);

        toFiler(builder);
    }

    private void generateSerializer(List<MethodSpec> methods) {

        MethodSpec constructor = getConstructorDeserialized(methods);

        TypeSpec.Builder builder = TypeSpec.classBuilder("Serializer")
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .addField(DataOutputStream.class, "output",
                        Modifier.PRIVATE, Modifier.FINAL)
                .addMethod(constructor);

        builder.addMethods(methods);

        builder.addField(FieldSpec.builder(typeMapForSerialized, "typeToMethod").initializer("new HashMap<>()")
                .addModifiers(Modifier.FINAL, Modifier.PRIVATE).build());

        MethodSpec serialize = MethodSpec.methodBuilder("serialize")
                .addParameter(classOfAny, "clazz")
                .addParameter(Object.class, "object")
                .addStatement("typeToMethod.get(clazz).accept(object)").build();

        builder.addMethod(serialize);


        toFiler(builder);
    }

    private MethodSpec getConstructorDeserialized(List<MethodSpec> methods) {
        MethodSpec.Builder constructorBuilder = MethodSpec.constructorBuilder()
                .addModifiers(Modifier.PUBLIC)
                .addParameter(DataInputStream.class, "input")
                .addStatement("this.$N = $N", "input", "input");

        for (MethodSpec method : methods) {
            String typeString = method.returnType.toString();
            String code = " typeToMethod.put($N.class, (Callable <$N>) () -> {\n" +
                    "try {\n "+
                    "  return this.$N();\n"+
                    "} catch (Exception e) {\n"+
                    " new RuntimeException());} return null;})";

            constructorBuilder.addStatement(code,
                    typeString,
                    typeString,
                    method.name);
        }
        return constructorBuilder.build();
    }


    private MethodSpec getConstructorSerialized(List<MethodSpec> methods) {

        MethodSpec.Builder constructorBuilder = MethodSpec.constructorBuilder()
                .addModifiers(Modifier.PUBLIC)
                .addParameter(DataOutputStream.class, "output")
                .addStatement("this.$N = $N", "output", "output");


        for (MethodSpec method : methods) {
            String typeString = method.parameters.get(0).type.toString();
            String code = " typeToMethod.put($N.class, (Consumer <$N>) ($N x) -> {\n" +
                    "try {\n "+
                    "  this.$N(x);\n"+
                    "} catch (Exception e) {\n"+
                    " e.addSuppressed(new RuntimeException());}})";
            constructorBuilder.addStatement(code,
                    typeString,
                    typeString,
                    typeString,
                    method.name);
        }
        return constructorBuilder.build();
    }


    private void toFiler(TypeSpec.Builder builder) {
        final JavaFile javaFile = JavaFile.builder("ru.ifmo.java.serialization", builder.build()).build();

        try {
            javaFile.writeTo(filer);
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException("Can't write to filer");
        }
    }

    private boolean checkLetterImplementation(TypeElement typeElement) {
        while (true) {
            TypeElement objectType = elementUtils.getTypeElement(Object.class.getCanonicalName());
            TypeElement letterType = elementUtils.getTypeElement(Letter.class.getCanonicalName());
            TypeElement supperType = (TypeElement) typeUtils.asElement(typeElement.getSuperclass());

            boolean contains = typeElement.getInterfaces().contains(letterType.asType());
            if (contains) {
                return true;
            }
            if (supperType.equals(objectType)) {
                return false;
            }

            typeElement = supperType;
        }
    }

    private void error(Element e, String msg, Object... args) {
        messager.printMessage(Diagnostic.Kind.ERROR, String.format(msg, args), e);
   }
}