package org.kravemir.lightvalue;

import com.squareup.javapoet.*;
import org.kravemir.lightvalue.model.ClassDef;
import org.kravemir.lightvalue.model.ModelFile;
import org.kravemir.lightvalue.model.PropertyDef;
import org.kravemir.lightvalue.naming.DefaultNameProvider;
import org.kravemir.lightvalue.naming.NameProvider;

import javax.lang.model.element.Modifier;
import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

public class ModelGenerator {

    private final NameProvider nameProvider;

    public ModelGenerator() {
        this(new DefaultNameProvider());
    }

    public ModelGenerator(NameProvider nameProvider) {
        this.nameProvider = nameProvider;
    }

    public void generate(ModelFile modelFile, File outputDir) throws IOException {
        ModelGenerationContext context = new ModelGenerationContext();
        context.setPackageName(modelFile.getFilePackage());

        for(Map.Entry<String, String> entry : modelFile.getImports().entrySet()) {
            context.registerType(entry.getKey(), ClassName.bestGuess(entry.getValue()));
        }

        for(ClassDef classDef : modelFile.getClassDefs()) {
            JavaFile javaFile = generateTopClassDef(context, classDef);
            context.registerType(classDef.getName(), ClassName.get(context.getPackageName(), classDef.getName()));
            javaFile.writeTo(outputDir);
        }
    }

    private JavaFile generateTopClassDef(ModelGenerationContext context, ClassDef classDef) {
        ClassName classTypeName = ClassName.get(context.getPackageName(), classDef.getName());
        ClassModelGenerationContext classContext = new ClassModelGenerationContext(
                context,
                classDef,
                classTypeName,
                classTypeName.nestedClass("Builder")
        );

        Collection<TypeSpec> subModelsSpecs = generateSubModels(classContext);

        TypeSpec typeSpec = TypeSpec
                .classBuilder(classDef.getName())
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .addFields(generateFields(classContext, Modifier.PRIVATE, Modifier.FINAL))
                .addMethod(generateConstructor(classContext))
                .addMethods(generateGetters(classContext))
                .addMethod(generateHashCode(classDef))
                .addMethod(generateEquals(classTypeName, classDef))
                .addMethod(generateToBuilderMethod(classContext))
                .addMethods(generateNewBuilderMethods(classContext))
                .addType(generateBuilder(classContext))
                .addTypes(subModelsSpecs)
                .build();


        return JavaFile.builder(context.getPackageName(), typeSpec).build();
    }

    private Collection<TypeSpec> generateSubModels(ClassModelGenerationContext parentContext) {
        Collection<ClassDef> classDefs = parentContext.getClassDef().getSubClassDefs();

        if(classDefs == null || classDefs.isEmpty()) {
            return Collections.emptyList();
        }

        List<TypeSpec> typeSpecs = new ArrayList<>();
        for(ClassDef classDef : classDefs) {
            TypeSpec subModelType = generateSubModel(parentContext, classDef);
            typeSpecs.add(subModelType);
        }

        return typeSpecs;
    }

    private TypeSpec generateSubModel(ClassModelGenerationContext parentContext, ClassDef classDef) {
        ClassName classTypeName = parentContext.getClassTypeName().nestedClass(classDef.getName());
        ClassModelGenerationContext classContext = new ClassModelGenerationContext(
                parentContext,
                classDef,
                classTypeName,
                classTypeName.nestedClass("Builder")
        );

        parentContext.registerType(classDef.getName(), classTypeName);

        Collection<TypeSpec> subModelsSpecs = generateSubModels(classContext);

        TypeSpec typeSpec = TypeSpec
                .classBuilder(classDef.getName())
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .addFields(generateFields(classContext, Modifier.PRIVATE, Modifier.FINAL))
                .addMethod(generateConstructor(classContext))
                .addMethods(generateGetters(classContext))
                .addMethod(generateHashCode(classDef))
                .addMethod(generateEquals(classTypeName, classDef))
                .addMethod(generateToBuilderMethod(classContext))
                .addMethods(generateNewBuilderMethods(classContext))
                .addType(generateBuilder(classContext))
                .addTypes(subModelsSpecs)
                .build();

        return typeSpec;
    }

    private MethodSpec generateConstructor(ClassModelGenerationContext context) {
        MethodSpec.Builder builder = MethodSpec
                .constructorBuilder()
                .addModifiers(Modifier.PRIVATE)
                .addParameter(
                        ParameterSpec.builder(
                                context.getBuilderTypeName(),
                                "builder",
                                Modifier.FINAL
                        ).build()
                );

        for(PropertyDef property : context.getClassDef().getProperties()) {
            builder.addStatement("this.$L = builder.$L", nameProvider.getVariableName(property), nameProvider.getVariableName(property));
        }

        return builder.build();
    }

    private Iterable<FieldSpec> generateFields(ClassModelGenerationContext context, Modifier... modifiers) {
        Collection<PropertyDef> properties = context.getClassDef().getProperties();
        if(properties == null || properties.isEmpty()) {
            return Collections.emptyList();
        }

        List<FieldSpec> typeSpecs = new ArrayList<>();
        for(PropertyDef propertyDef : properties) {
            typeSpecs.add(toFieldSpec(context, propertyDef, modifiers));
        }

        return typeSpecs;
    }

    private FieldSpec toFieldSpec(ModelGenerationContext context, PropertyDef propertyDef, Modifier... modifiers) {
        return FieldSpec.builder(
                getTypeForProperty(context, propertyDef),
                nameProvider.getVariableName(propertyDef),
                modifiers
        ).build();
    }

    private Iterable<MethodSpec> generateGetters(ClassModelGenerationContext context) {
        Collection<PropertyDef> properties = context.getClassDef().getProperties();
        if(properties == null || properties.isEmpty()) {
            return Collections.emptyList();
        }

        List<MethodSpec> methodSpecs = new ArrayList<>();
        for(PropertyDef propertyDef : properties) {
            methodSpecs.add(generateGetter(context, propertyDef));
        }

        return methodSpecs;
    }

    private MethodSpec generateGetter(ModelGenerationContext context, PropertyDef propertyDef) {
        return MethodSpec
                .methodBuilder("get" + nameProvider.getAccessorName(propertyDef))
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .returns(getTypeForProperty(context, propertyDef))
                .addCode(CodeBlock.builder()
                        .addStatement("return this.$L", nameProvider.getVariableName(propertyDef))
                        .build()
                )
                .build();
    }

    private MethodSpec generateEquals(ClassName classTypeName, ClassDef classDef) {
        return MethodSpec
                .methodBuilder("equals")
                .addAnnotation(Override.class)
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .addParameter(Object.class, "o")
                .returns(TypeName.BOOLEAN)
                .addCode("if (this == o) return true;\n")
                .addCode("if (o == null || getClass() != o.getClass()) return false;\n")
                .addCode("$T other = ($T) o;\n", classTypeName, classTypeName)
                .addStatement(
                        genReturnEqualsStatement(classDef),
                        genReturnEqualsSubstitutions(classDef)
                )
                .build();
    }

    private String genReturnEqualsStatement(ClassDef classDef) {
        if(classDef.getProperties().isEmpty()) {
            return "return true";
        }

        return classDef.getProperties().stream()
                .map(this::toComparisonExpression)
                .collect(Collectors.joining(" && \n", "return ", ""));
    }

    private String toComparisonExpression(PropertyDef propertyDef) {
        String variableName = nameProvider.getVariableName(propertyDef);

        if(isPrimitiveType(propertyDef)) {
            return String.format("this.%s == other.%s", variableName, variableName);
        } else {
            return String.format("$T.equals(this.%s,other.%s)", variableName, variableName);
        }
    }

    private Object[] genReturnEqualsSubstitutions(ClassDef classDef) {
        if(classDef.getProperties().isEmpty()) {
            return new Object[0];
        }

        return classDef.getProperties().stream()
                .filter(p -> !isPrimitiveType(p))
                .map(p -> Objects.class)
                .toArray();
    }

    private boolean isPrimitiveType(PropertyDef propertyDef) {
        if(propertyDef.isMap() || propertyDef.isRepeated()) {
            return false;
        }

        switch (propertyDef.getType1()) {
            case "bool":
            case "int32":
            case "double":
                return true;
            case "string":
            default:
                return false;
        }
    }

    private MethodSpec generateHashCode(ClassDef classDef) {
        return MethodSpec
                .methodBuilder("hashCode")
                .addAnnotation(Override.class)
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .returns(TypeName.INT)
                .addStatement(
                        genReturnHashCodeStatement(classDef),
                        genReturnHashCodeSubstitutions(classDef)
                )
                .build();
    }

    private String genReturnHashCodeStatement(ClassDef classDef) {
        return classDef.getProperties().stream()
                .map(p -> "$L")
                .collect(Collectors.joining(",", "return $T.hash(", ")"));
    }

    private Object[] genReturnHashCodeSubstitutions(ClassDef classDef) {
        List<Object> substitutions = new ArrayList<>(1 + classDef.getProperties().size());
        substitutions.add(TypeName.get(Objects.class));
        for (PropertyDef property : classDef.getProperties()) {
            substitutions.add(nameProvider.getVariableName(property));
        }
        return substitutions.toArray();
    }

    private MethodSpec generateToBuilderMethod(ClassModelGenerationContext context) {
        return MethodSpec
                .methodBuilder("toBuilder")
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .returns(context.getBuilderTypeName())
                .addStatement("return new $T(this)", context.getBuilderTypeName())
                .build();
    }

    private Iterable<MethodSpec> generateNewBuilderMethods(ClassModelGenerationContext context) {
        return Collections.singletonList(
                MethodSpec
                        .methodBuilder("newBuilder")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                        .returns(context.getBuilderTypeName())
                        .addCode("return new $T();\n", context.getBuilderTypeName())
                        .build()
        );
    }

    private TypeSpec generateBuilder(ClassModelGenerationContext context) {
        return TypeSpec
                .classBuilder("Builder")
                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                .addFields(generateFields(context, Modifier.PRIVATE))
                .addMethod(MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC).build())
                .addMethod(generateCopyConstructor(context))
                .addMethods(generateSetters(context))
                .addMethod(generateBuildMethod(context.getClassTypeName()))
                .build();
    }

    private MethodSpec generateCopyConstructor(ClassModelGenerationContext context) {
        MethodSpec.Builder builder = MethodSpec
                .constructorBuilder()
                .addModifiers(Modifier.PRIVATE)
                .addParameter(
                        ParameterSpec.builder(
                                context.getClassTypeName(),
                                "original",
                                Modifier.FINAL
                        ).build()
                );

        for(PropertyDef property : context.getClassDef().getProperties()) {
            builder.addStatement("this.$L = original.$L", nameProvider.getVariableName(property), nameProvider.getVariableName(property));
        }

        return builder.build();
    }

    private Iterable<MethodSpec> generateSetters(ClassModelGenerationContext context) {
        Collection<PropertyDef> properties = context.getClassDef().getProperties();
        if(properties == null || properties.isEmpty()) {
            return Collections.emptyList();
        }

        List<MethodSpec> methodSpecs = new ArrayList<>();
        for(PropertyDef propertyDef : properties) {

            if(propertyDef.isRepeated()) {
                methodSpecs.addAll(generateAddMethods(context, propertyDef));
            } else {
                methodSpecs.add(generateSetter(context, propertyDef));
            }
        }

        return methodSpecs;
    }

    private MethodSpec generateSetter(ClassModelGenerationContext context, PropertyDef propertyDef) {
        return MethodSpec
                .methodBuilder("set" + nameProvider.getAccessorName(propertyDef))
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .returns(context.getBuilderTypeName())
                .addParameter(getTypeForProperty(context, propertyDef), "value", Modifier.FINAL)
                .addStatement("this.$L = value", nameProvider.getVariableName(propertyDef))
                .addStatement("return this")
                .build();
    }

    private Collection<MethodSpec> generateAddMethods(ClassModelGenerationContext context, PropertyDef propertyDef) {
        return Arrays.asList(
                MethodSpec
                        .methodBuilder("set" + nameProvider.getAccessorName(propertyDef))
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .returns(context.getBuilderTypeName())
                        .addParameter(getTypeForProperty(context, propertyDef), "value", Modifier.FINAL)
                        .addStatement("this.$L = new $T<>(value)", nameProvider.getVariableName(propertyDef), ClassName.get(ArrayList.class))
                        .addStatement("return this")
                        .build(),
                MethodSpec
                        .methodBuilder("add" + nameProvider.getSingularAccessorName(propertyDef))
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .returns(context.getBuilderTypeName())
                        .addParameter(context.getType(propertyDef.getType1()), "value", Modifier.FINAL)
                        .addStatement(
                                "if(this.$L == null) { this.$L = new $T<>(); }",
                                nameProvider.getVariableName(propertyDef),
                                nameProvider.getVariableName(propertyDef),
                                ClassName.get(ArrayList.class)
                        )
                        .addStatement("this.$L.add(value)", nameProvider.getVariableName(propertyDef))
                        .addStatement("return this")
                        .build(),
                MethodSpec
                        .methodBuilder("addAll" + nameProvider.getAccessorName(propertyDef))
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .returns(context.getBuilderTypeName())
                        .addParameter(getTypeForProperty(context, propertyDef), "values", Modifier.FINAL)
                        .addStatement(
                                "if(this.$L == null) { this.$L = new $T<>(); }",
                                nameProvider.getVariableName(propertyDef),
                                nameProvider.getVariableName(propertyDef),
                                ClassName.get(ArrayList.class)
                        )
                        .addStatement("this.$L.addAll(values)", nameProvider.getVariableName(propertyDef))
                        .addStatement("return this")
                        .build()
        );
    }

    private MethodSpec generateBuildMethod(ClassName classTypeName) {
        return MethodSpec
                .methodBuilder("build")
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .returns(classTypeName)
                .addStatement("return new $T(this)", classTypeName)
                .build();
    }

    private TypeName getTypeForProperty(ModelGenerationContext context, PropertyDef propertyDef) {
        if(propertyDef.isRepeated()) {
            return ParameterizedTypeName.get(ClassName.get(List.class), context.getType(propertyDef.getType1()));
        } else if(propertyDef.isMap()) {
            return ParameterizedTypeName.get(
                    ClassName.get(Map.class),
                    context.getType(propertyDef.getType1()),
                    context.getType(propertyDef.getType2())
            );
        } else {
            return context.getType(propertyDef.getType1());
        }
    }
}
