package io.trygvis.container.compiler;

import io.trygvis.container.compiler.model.TypeRef;
import io.trygvis.persistence.EntityMirror;
import io.trygvis.persistence.FieldMirror;
import io.trygvis.persistence.GeneratorConfiguration;
import io.trygvis.persistence.SequenceMirror;
import io.trygvis.persistence.SqlEntity;
import io.trygvis.persistence.SqlEntitySet;
import io.trygvis.persistence.TypeHandler;
import io.trygvis.persistence.generators.DaoGenerator;

import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.AnnotationValue;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.PackageElement;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.persistence.Id;
import javax.persistence.SequenceGenerator;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import static io.trygvis.container.compiler.Utils.toFieldName;
import static io.trygvis.container.compiler.Utils.writeFile;
import static io.trygvis.persistence.FieldMirror.FieldType.PRIMITIVE;
import static io.trygvis.persistence.FieldMirror.FieldType.REFERENCE;
import static io.trygvis.persistence.FieldMirror.GetterType;
import static io.trygvis.persistence.FieldMirror.SetterType;
import static io.trygvis.persistence.TypeHandler.StringEnumTypeHandler;
import static java.lang.Character.isUpperCase;
import static java.lang.String.format;
import static javax.lang.model.util.ElementFilter.fieldsIn;
import static javax.lang.model.util.ElementFilter.methodsIn;
import static org.apache.commons.lang.StringUtils.join;

public class EntityHandler extends AbstractHandler {
    private GeneratorConfiguration generatorConfiguration = new GeneratorConfiguration();
    private SqlUnitModel sqlUnit;

    public EntityHandler(ProcessingEnvironment processingEnv, SqlUnitModel sqlUnit) {
        super(processingEnv);
        this.sqlUnit = sqlUnit;
    }

    public void phase1(Set<TypeElement> sqlEntities, Set<PackageElement> packages) throws Exception {
        for (TypeElement entity : sqlEntities) {
            AnnotationMirror sqlEntity = getAnnotation(SqlEntity.class, entity);
            for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> v : sqlEntity.getElementValues().entrySet()) {
                String field = v.getKey().getSimpleName().toString();
                switch (field) {
                    case "value":
                        Class<?> typeHandlerClass = getClass().getClassLoader().loadClass(v.getValue().getValue().toString());

                        TypeHandler typeHandler = (TypeHandler) typeHandlerClass.newInstance();
                        String type = entity.asType().toString();
                        generatorConfiguration.addTypeHandler(new TypeRef(type), typeHandler);
                        System.out.println("Loaded TypeHandler for " + type + " through " + typeHandlerClass.getCanonicalName());
                        break;
                    default:
                        throw new InternalErrorException("Unknown @SqlEntity field: " + field);
                }
            }
        }

        if (packages.size() == 0) {
            throw new CompilerException("There has to be exactly one @SqlEntitySet annotated package.");
        }
        PackageElement packageElement = packages.iterator().next();

        if (packages.size() != 1) {
            throw new CompilerException(packageElement, "There can only be one @SqlEntitySet annotated package.");
        }
        AnnotationMirror sqlEntitySet = getAnnotation(SqlEntitySet.class, packageElement);
        String name = null;
        for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> v : sqlEntitySet.getElementValues().entrySet()) {
            String field = v.getKey().getSimpleName().toString();
            String value = v.getValue().getValue().toString();
            switch (field) {
                case "name":
                    name = value;
            }
        }
        sqlUnit.setPackageName(packageElement.getQualifiedName().toString());
        sqlUnit.setName(name);
    }

    Set<TypeElement> mappedSuperclasses = new HashSet<>();
    Set<TypeElement> entities = new HashSet<>();

    public void record(TypeElement element, boolean superclass) throws Exception {
        (superclass ? mappedSuperclasses : entities).add(element);
    }

    private EntityMirror getEntity(TypeElement element, boolean superclass) throws Exception {
        EntityMirror entity = sqlUnit.entity(element);
        if (entity != null) {
            return entity;
        }

        entity = buildMirror(element, superclass);
        sqlUnit.add(entity, element);
        return entity;
    }

    public EntityMirror buildMirror(TypeElement element, boolean concrete) throws Exception {
        EntityMirror superEntity = null;

        Set<TypeMirror> ancestry = new TreeSet<>(new TypeMirrorComparator());
        ancestry.add(element.asType());
        TypeMirror superType = element.getSuperclass();
        while (!superType.toString().equals("java.lang.Object")) {
            TypeElement superTypeMirror = elements.getTypeElement(superType.toString());
            if (superTypeMirror == null) {
                throw new CompilerException(element, "Couldn't find element for super class: " + element.getSuperclass());
            }

            if (mappedSuperclasses.contains(superTypeMirror)) {
                superEntity = getEntity(superTypeMirror, false);
                break;
            }
            ancestry.add(superType);
            superType = superTypeMirror.getSuperclass();
        }

        EntityMirror entity = new EntityMirror(
                generatorConfiguration,
                new TypeRef(types.getDeclaredType(element)),
                concrete,
                superEntity,
                concrete ? Utils.sqlName(element.getSimpleName().toString()) : null);

        Map<String, ExecutableElement> getters = new TreeMap<>();
        Map<String, ExecutableElement> setters = new TreeMap<>();
        for (ExecutableElement m : methodsIn(elements.getAllMembers(element))) {
            if (!ancestry.contains(m.getEnclosingElement().asType())) {
                continue;
            }

            String name = m.getSimpleName().toString();
            if (name.length() < 4 || !isUpperCase(name.charAt(3))) {
                continue;
            }

            boolean isVoid = m.getReturnType().getKind().equals(TypeKind.VOID);

            if (name.startsWith("get") && m.getParameters().size() == 0 && !isVoid) {
                getters.put(toFieldName(name), m);
            }

            if (name.startsWith("set") && m.getParameters().size() == 1 && isVoid) {
                setters.put(toFieldName(name), m);
            }
        }

        for (VariableElement f : fieldsIn(elements.getAllMembers(element))) {
            System.out.println("f.getEnclosingElement() = " + f.getEnclosingElement());
            if (!ancestry.contains(f.getEnclosingElement().asType())) {
                continue;
            }

            String name = f.getSimpleName().toString();
            ExecutableElement getter = getters.remove(name);
            ExecutableElement setter = setters.remove(name);
            FieldMirror field = fromElement(generatorConfiguration, f, getter, setter);
            if (field == null) {
                continue;
            }
            entity.add(field);
            System.out.println("Field: " + field);
        }

        for (Map.Entry<String, ExecutableElement> e : getters.entrySet()) {
            String name = e.getKey();
            ExecutableElement getter = e.getValue();
            ExecutableElement setter = setters.remove(name);

            if (!ancestry.contains(getter.getEnclosingElement().asType())) {
                continue;
            }

            FieldMirror field = fromElement(generatorConfiguration, null, getter, setter);
            if (field == null) {
                continue;
            }
            entity.add(field);
            System.out.println("Field: " + field);
        }

        Iterator<ExecutableElement> it = setters.values().iterator();
        while (it.hasNext()) {
            ExecutableElement setter = it.next();
            if (setter.getModifiers().contains(Modifier.STATIC)) {
                it.remove();
            }
        }

        if (!setters.isEmpty()) {
            throw new CompilerException(element, "Missing getters for these setters: " + join(setters.keySet(), ", "));
        }

        // -----------------------------------------------------------------------
        // Process any extra annotations
        // -----------------------------------------------------------------------

        processSequenceGenerators(element);

        return entity;
    }

    public void validate(Element element, EntityMirror entity) {
        if (!entity.concrete) {
            return;
        }

        List<FieldMirror> idFields = new ArrayList<>();
        for (FieldMirror field : entity.getFields()) {
            if (field.id) {
                idFields.add(field);
            }
        }

        if (idFields.size() == 0) {
            throw new CompilerException(element, "An @Entity is required to have at least one @Id field.");
        }
        if (idFields.size() != 1) {
            throw new CompilerException(element, "This implementation only support a single @Id annotated field.");
        }

        entity.setIdType(idFields.get(0).type);
    }

    private void processSequenceGenerators(Element element) {
        AnnotationMirror sequenceGenerator = findAnnotation(SequenceGenerator.class, element);

        if (sequenceGenerator != null) {
            String name = null;
            String sequenceName = null;
            int initialValue = 0;
            int allocationSize = 0;

            for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> v : sequenceGenerator.getElementValues().entrySet()) {
                String field = v.getKey().getSimpleName().toString();
                String value = v.getValue().getValue().toString();
                switch (field) {
                    case "name":
                        name = value;
                        break;
                    case "sequenceName":
                        sequenceName = value;
                        break;
                    case "initialValue":
                        initialValue = Integer.valueOf(value);
                        break;
                    case "allocationSize":
                        allocationSize = Integer.valueOf(value);
                        break;
                    default:
                        throw new InternalErrorException("Unsupported field on @SequenceGenerator: " + field);
                }
            }
            if (name != null) {
                sequenceName = sequenceName == null ? Utils.sqlName(name) : sequenceName;
                sqlUnit.add(new SequenceMirror(name, sequenceName, initialValue, allocationSize), element);
            }
        }
    }

    public FieldMirror fromElement(GeneratorConfiguration generatorConfiguration, VariableElement var,
                                   ExecutableElement getter, ExecutableElement setter) {
        // TODO: check the setter for annotations too
        // TODO: check for transient and @Transient
        final SetterType setterType;
        final GetterType getterType;
        TypeRef type;
        Element element;
        String javaName;
        String sqlName;
        boolean id;

        if (var != null) {
            if (var.getModifiers().contains(Modifier.STATIC)) {
                return null;
            }
            boolean isFinal = var.getModifiers().contains(Modifier.FINAL);
            if (isFinal) {
                setterType = SetterType.CONSTRUCTOR;
            } else {
                setterType = setter != null ? SetterType.METHOD : SetterType.FIELD;
            }
            getterType = getter != null ? GetterType.METHOD : GetterType.FIELD;
            type = new TypeRef(var.asType());
            element = types.asElement(var.asType());
            javaName = var.getSimpleName().toString();
            id = isId(var);
            processSequenceGenerators(var);
        } else {
            if (getter.getModifiers().contains(Modifier.STATIC)) {
                return null;
            }
            if (setter == null) {
                // Skipping fields is closer to what hibernate does.
                return null;
            }
            setterType = SetterType.METHOD;
            getterType = GetterType.METHOD;
            type = new TypeRef(getter.getReturnType());
            element = types.asElement(getter.getReturnType());
            id = isId(getter);
            processSequenceGenerators(getter);
            // TODO: this might be relaxed, just find the common type and use that.
            if (!types.isSameType(getter.getReturnType(), setter.getParameters().get(0).asType())) {
                throw new CompilerException(format("The setter and getter %s/%s  must access the same types.",
                        setter.getSimpleName(), getter.getSimpleName()));
            }
            javaName = toFieldName(getter.getSimpleName().toString());
        }
        sqlName = Utils.sqlName(javaName);

        // Register type handler if this is an enum
        if (element != null && element.getKind() == ElementKind.ENUM) {
            if (!generatorConfiguration.hasTypeHandler(type)) {
                generatorConfiguration.addTypeHandler(type, new StringEnumTypeHandler(type));
            }
        }

        boolean notNull = false;
        boolean unique = false;
        boolean primitive = generatorConfiguration.hasTypeHandler(type);
        if (id && !primitive) {
            throw new CompilerException(var, "A @Id field has to be a primitive or embedded.");
        }

        FieldMirror field;
        if (primitive) {
            TypeHandler typeHandler = generatorConfiguration.typeHandler(type);
            // TODO: check for configuration conflict
            notNull = !typeHandler.nullable;
            field = new FieldMirror(PRIMITIVE, setterType, getterType, type, javaName, sqlName, id, notNull, unique);
        } else if (generatorConfiguration.hasTypeHandler(type)) {
            throw new CompilerException(var, "Missing type handler for type: " + type.fqName);
        } else {
            field = new FieldMirror(REFERENCE, setterType, getterType, type, javaName, sqlName, id, notNull, unique);
        }
        return field;
    }

    public static boolean isId(Element var) {
        return var.getAnnotation(Id.class) != null;
    }

    public void phase3(boolean errorRaised) throws Exception {
        for (TypeElement mappedSuperclass : mappedSuperclasses) {
            getEntity(mappedSuperclass, false);
        }

        for (TypeElement entity : entities) {
            getEntity(entity, true);
        }

        for (EntityMirror entity : sqlUnit.getEntities().values()) {
            try {
                validate(sqlUnit.element(entity), entity);

                if (!entity.concrete) {
                    continue;
                }

                System.out.println("entity.type = " + entity.type);
                DaoGenerator daoGenerator = new DaoGenerator(generatorConfiguration, sqlUnit, entity);
                writeFile(processingEnv, daoGenerator.generate(), sqlUnit.element(entity));
                writeFile(processingEnv, daoGenerator.generateRow(), sqlUnit.element(entity));
            } catch (CompilerException | InternalErrorException e) {
                // Ignore any exceptions if we had an error from before
                if (errorRaised) {
                    continue;
                }
                throw e;
            }
        }
    }

    private AnnotationMirror findAnnotation(Class<?> c, Element type) {
        TypeMirror annotationType = elements.getTypeElement(c.getCanonicalName()).asType();
        for (AnnotationMirror a : type.getAnnotationMirrors()) {
            if (types.isSameType(a.getAnnotationType(), annotationType)) {
                return a;
            }
        }

        return null;
    }

    public AnnotationMirror getAnnotation(Class<?> c, Element type) {
        AnnotationMirror annotation = findAnnotation(c, type);

        if (annotation != null) {
            return annotation;
        }

        throw new CompilerException(type, "Could not find annotation " + c.getSimpleName());
    }
}