package org.kravemir.lightvalue;

import com.squareup.javapoet.*;
import org.kravemir.lightvalue.model.ClassDef;
import org.kravemir.lightvalue.model.ModelFile;

import javax.lang.model.element.Modifier;
import java.io.File;
import java.io.IOException;
import java.util.Collection;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class JacksonMixInsGenerator {

    private final static ClassName OBJECT_MAPPER_NAME = ClassName.get(
            "com.fasterxml.jackson.databind", "ObjectMapper"
    );

    private final static ClassName JSON_DESERIALIZE_NAME = ClassName.get(
            "com.fasterxml.jackson.databind.annotation", "JsonDeserialize"
    );
    private final static ClassName JSON_POJO_BUILDER_NAME = ClassName.get(
            "com.fasterxml.jackson.databind.annotation", "JsonPOJOBuilder"
    );

    public void generate(ModelFile modelFile, File outputDir) throws IOException {
        TypeSpec topClassTypeSpec = TypeSpec
                .classBuilder("JacksonMixIns")
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .addTypes(generateMixInClasses(modelFile))
                .addMethod(generateRegisterMixInsMethod(modelFile))
                .build();

        JavaFile javaFile = JavaFile.builder(modelFile.getFilePackage(), topClassTypeSpec).build();
        javaFile.writeTo(outputDir);
    }

    private Iterable<TypeSpec> generateMixInClasses(ModelFile modelFile) {
        return generateMixInClasses(
                name -> ClassName.get(modelFile.getFilePackage(), name),
                name -> ClassName.get(modelFile.getFilePackage(), "JacksonMixIns", name),
                modelFile.getClassDefs()
        );
    }

    private Iterable<TypeSpec> generateMixInClasses(
            final Function<String, ClassName> targetNameGenerator,
            final Function<String, ClassName> mixInSourceGenerator,
            final Collection<ClassDef> classDefs
    ) {
        return classDefs.stream().map(classDef -> {
            final ClassName targetName = targetNameGenerator.apply(classDef.getName());
            final ClassName mixInSource = mixInSourceGenerator.apply(classDef.getName() + "MixIn");

            return TypeSpec
                    .classBuilder(mixInSource)
                    .addAnnotation(AnnotationSpec
                            .builder(JSON_DESERIALIZE_NAME)
                            .addMember(
                                    "builder",
                                    CodeBlock.of("$T.Builder.class", targetName)
                            )
                            .build()
                    )
                    .addModifiers(Modifier.STATIC, Modifier.ABSTRACT, Modifier.PUBLIC)
                    .addType(generateBuilder())
                    .addTypes(generateMixInClasses(
                            targetName::nestedClass,
                            mixInSource::nestedClass,
                            classDef.getSubClassDefs()
                    ))
                    .build();
        }).collect(Collectors.toList());
    }

    private TypeSpec generateBuilder() {
        return TypeSpec
                .classBuilder("Builder")
                .addAnnotation(AnnotationSpec
                        .builder(JSON_POJO_BUILDER_NAME)
                        .addMember("withPrefix", CodeBlock.of("\"set\""))
                        .build()
                )
                .addModifiers(Modifier.STATIC, Modifier.ABSTRACT, Modifier.PUBLIC)
                .build();
    }

    private MethodSpec generateRegisterMixInsMethod(ModelFile modelFile) {
        return MethodSpec
                .methodBuilder("registerMixIns")
                .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                .addParameter(OBJECT_MAPPER_NAME, "objectMapper")
                .addCode(
                        generateRegisterCodeBlocks(
                                name -> ClassName.get(modelFile.getFilePackage(), name),
                                name -> ClassName.get(modelFile.getFilePackage(), "JacksonMixIns", name),
                                modelFile.getClassDefs()
                        ).collect(CodeBlock.joining("\n"))
                )
                .build();
    }

    private Stream<CodeBlock> generateRegisterCodeBlocks(
            final Function<String, ClassName> targetNameGenerator,
            final Function<String, ClassName> mixInSourceGenerator,
            final Collection<ClassDef> classDefs
    ) {
        return classDefs.stream().flatMap(classDef -> {
            final ClassName targetName = targetNameGenerator.apply(classDef.getName());
            final ClassName mixInSource = mixInSourceGenerator.apply(classDef.getName() + "MixIn");

            return Stream.concat(
                    Stream.of(
                            generateRegisterValueClass(targetName, mixInSource),
                            generateRegisterBuilder(targetName, mixInSource)
                    ),
                    generateRegisterCodeBlocks(
                            targetName::nestedClass,
                            mixInSource::nestedClass,
                            classDef.getSubClassDefs()
                    )
            );
        });
    }

    private CodeBlock generateRegisterValueClass(ClassName targetName, ClassName mixInSource) {
        return CodeBlock.of("objectMapper.addMixIn($T.class, $T.class);", targetName, mixInSource);
    }

    private CodeBlock generateRegisterBuilder(ClassName targetName, ClassName mixInSource) {
        return CodeBlock.of("objectMapper.addMixIn($T.Builder.class, $T.Builder.class);\n", targetName, mixInSource);
    }
}
