package io.quarkus.security.jpa.deployment;

import java.lang.reflect.Modifier;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.function.BiConsumer;

import javax.inject.Singleton;
import javax.persistence.EntityManager;
import javax.persistence.Query;

import org.hibernate.Session;
import org.hibernate.SimpleNaturalIdLoadAccess;
import org.hibernate.annotations.NaturalId;
import org.jboss.jandex.AnnotationInstance;
import org.jboss.jandex.AnnotationTarget;
import org.jboss.jandex.AnnotationValue;
import org.jboss.jandex.ClassInfo;
import org.jboss.jandex.DotName;
import org.jboss.jandex.Index;
import org.jboss.jandex.Type;

import io.quarkus.arc.deployment.GeneratedBeanBuildItem;
import io.quarkus.arc.deployment.GeneratedBeanGizmoAdaptor;
import io.quarkus.arc.deployment.UnremovableBeanBuildItem;
import io.quarkus.deployment.Capabilities;
import io.quarkus.deployment.annotations.BuildProducer;
import io.quarkus.deployment.annotations.BuildStep;
import io.quarkus.deployment.builditem.ApplicationIndexBuildItem;
import io.quarkus.deployment.builditem.CapabilityBuildItem;
import io.quarkus.deployment.builditem.FeatureBuildItem;
import io.quarkus.gizmo.AssignableResultHandle;
import io.quarkus.gizmo.BranchResult;
import io.quarkus.gizmo.BytecodeCreator;
import io.quarkus.gizmo.ClassCreator;
import io.quarkus.gizmo.MethodCreator;
import io.quarkus.gizmo.MethodDescriptor;
import io.quarkus.gizmo.ResultHandle;
import io.quarkus.panache.common.deployment.PanacheEntityClassesBuildItem;
import io.quarkus.security.identity.SecurityIdentity;
import io.quarkus.security.identity.request.UsernamePasswordAuthenticationRequest;
import io.quarkus.security.jpa.Password;
import io.quarkus.security.jpa.PasswordType;
import io.quarkus.security.jpa.Roles;
import io.quarkus.security.jpa.RolesValue;
import io.quarkus.security.jpa.UserDefinition;
import io.quarkus.security.jpa.Username;
import io.quarkus.security.jpa.deployment.JpaSecurityDefinition.FieldOrMethod;
import io.quarkus.security.jpa.runtime.JpaIdentityProvider;
import io.quarkus.security.runtime.QuarkusSecurityIdentity;

class QuarkusSecurityJpaProcessor {

    static final DotName DOTNAME_OBJECT = DotName.createSimple(Object.class.getName());

    private static final DotName DOTNAME_STRING = DotName.createSimple(String.class.getName());
    private static final DotName DOTNAME_LIST = DotName.createSimple(List.class.getName());
    private static final DotName DOTNAME_SET = DotName.createSimple(Set.class.getName());
    private static final DotName DOTNAME_COLLECTION = DotName.createSimple(Collection.class.getName());

    private static final DotName DOTNAME_NATURAL_ID = DotName.createSimple(NaturalId.class.getName());

    private static final DotName DOTNAME_USER_DEFINITION = DotName.createSimple(UserDefinition.class.getName());
    private static final DotName DOTNAME_USERNAME = DotName.createSimple(Username.class.getName());
    private static final DotName DOTNAME_PASSWORD = DotName.createSimple(Password.class.getName());
    private static final DotName DOTNAME_ROLES = DotName.createSimple(Roles.class.getName());
    private static final DotName DOTNAME_ROLES_VALUE = DotName.createSimple(RolesValue.class.getName());

    @BuildStep
    FeatureBuildItem feature() {
        return new FeatureBuildItem(FeatureBuildItem.SECURITY_JPA);
    }

    @BuildStep
    CapabilityBuildItem capability() {
        return new CapabilityBuildItem(Capabilities.SECURITY_JPA);
    }

    @BuildStep
    void configureJpaAuthConfig(ApplicationIndexBuildItem index,
            BuildProducer<UnremovableBeanBuildItem> unremovable,
            BuildProducer<GeneratedBeanBuildItem> beanProducer,
            List<PanacheEntityClassesBuildItem> panacheEntityClasses) throws Exception {

        // Generate an IdentityProvider if we have a @UserDefinition
        List<AnnotationInstance> userDefinitions = index.getIndex().getAnnotations(DOTNAME_USER_DEFINITION);
        if (userDefinitions.size() > 1) {
            throw new RuntimeException("You can only annotate one class with @UserDefinition");
        } else if (!userDefinitions.isEmpty()) {
            ClassInfo userDefinitionClass = userDefinitions.get(0).target().asClass();
            AnnotationTarget annotatedUsername = getSingleAnnotatedElement(index.getIndex(), DOTNAME_USERNAME);
            // FIXME: check that it's in the same class hierarchy
            // FIXME: look up fields with default names? name, username…
            AnnotationTarget annotatedPassword = getSingleAnnotatedElement(index.getIndex(), DOTNAME_PASSWORD);
            AnnotationTarget annotatedRoles = getSingleAnnotatedElement(index.getIndex(), DOTNAME_ROLES);
            Set<String> panacheEntities = collectPanacheEntities(panacheEntityClasses);
            // collect associated getters if required
            JpaSecurityDefinition jpaSecurityDefinition = new JpaSecurityDefinition(index.getIndex(),
                    userDefinitionClass,
                    isPanache(userDefinitionClass, panacheEntities),
                    annotatedUsername,
                    annotatedPassword,
                    annotatedRoles);
            AnnotationInstance passAnnotation = jpaSecurityDefinition.password.annotation(DOTNAME_PASSWORD);
            AnnotationValue passwordType = passAnnotation.value();
            generateIdentityProvider(index.getIndex(), jpaSecurityDefinition,
                    passwordType != null ? passwordType.asEnum() : PasswordType.MCF.name(),
                    beanProducer, panacheEntities);
        }
    }

    private boolean isPanache(ClassInfo annotatedClass, Set<String> panacheEntities) {
        return panacheEntities.contains(annotatedClass.name().toString());
    }

    private Set<String> collectPanacheEntities(List<PanacheEntityClassesBuildItem> panacheEntityClassesBuildItems) {
        Set<String> modelClasses = new HashSet<>();
        for (PanacheEntityClassesBuildItem panacheEntityClasses : panacheEntityClassesBuildItems) {
            modelClasses.addAll(panacheEntityClasses.getEntityClasses());
        }
        return modelClasses;
    }

    private AnnotationTarget getSingleAnnotatedElement(Index index, DotName annotation) {
        List<AnnotationInstance> annotations = index.getAnnotations(annotation);
        if (annotations.isEmpty()) {
            return null;
        } else if (annotations.size() > 1) {
            throw new RuntimeException("You can only annotate one field or method with @" + annotation);
        }
        return annotations.get(0).target();
    }

    private void generateIdentityProvider(Index index, JpaSecurityDefinition jpaSecurityDefinition, String passwordType,
            BuildProducer<GeneratedBeanBuildItem> beanProducer, Set<String> panacheClasses) {
        GeneratedBeanGizmoAdaptor gizmoAdaptor = new GeneratedBeanGizmoAdaptor(beanProducer);

        String name = jpaSecurityDefinition.annotatedClass.name() + "__JpaIdentityProviderImpl";
        try (ClassCreator classCreator = ClassCreator.builder()
                .className(name)
                .superClass(JpaIdentityProvider.class)
                .classOutput(gizmoAdaptor)
                .build()) {
            classCreator.addAnnotation(Singleton.class);
            try (MethodCreator methodCreator = classCreator.getMethodCreator("authenticate", SecurityIdentity.class,
                    EntityManager.class, UsernamePasswordAuthenticationRequest.class)) {
                methodCreator.setModifiers(Modifier.PUBLIC);

                ResultHandle username = methodCreator.invokeVirtualMethod(
                        MethodDescriptor.ofMethod(UsernamePasswordAuthenticationRequest.class, "getUsername", String.class),
                        methodCreator.getMethodParam(1));

                // two strategies, depending on whether the username is natural id
                AnnotationInstance naturalIdAnnotation = jpaSecurityDefinition.username.annotation(DOTNAME_NATURAL_ID);
                ResultHandle user;
                if (naturalIdAnnotation != null) {
                    // Session session = em.unwrap(Session.class);
                    ResultHandle session = methodCreator.invokeInterfaceMethod(
                            MethodDescriptor.ofMethod(EntityManager.class, "unwrap", Object.class, Class.class),
                            methodCreator.getMethodParam(0),
                            methodCreator.loadClass(Session.class));
                    // SimpleNaturalIdLoadAccess<PlainUserEntity> naturalIdLoadAccess = session.bySimpleNaturalId(PlainUserEntity.class);
                    ResultHandle naturalIdLoadAccess = methodCreator.invokeInterfaceMethod(
                            MethodDescriptor.ofMethod(Session.class, "bySimpleNaturalId",
                                    SimpleNaturalIdLoadAccess.class, Class.class),
                            methodCreator.checkCast(session, Session.class),
                            methodCreator.loadClass(jpaSecurityDefinition.annotatedClass.name().toString()));
                    // PlainUserEntity user = naturalIdLoadAccess.load(request.getUsername());
                    user = methodCreator.invokeInterfaceMethod(
                            MethodDescriptor.ofMethod(SimpleNaturalIdLoadAccess.class, "load",
                                    Object.class, Object.class),
                            naturalIdLoadAccess, username);
                } else {
                    // Query query = entityManager.createQuery("FROM Entity WHERE field = :name")
                    String hql = "FROM " + jpaSecurityDefinition.annotatedClass.simpleName() + " WHERE "
                            + jpaSecurityDefinition.username.name() + " = :name";
                    ResultHandle query = methodCreator.invokeInterfaceMethod(
                            MethodDescriptor.ofMethod(EntityManager.class, "createQuery", Query.class, String.class),
                            methodCreator.getMethodParam(0), methodCreator.load(hql));
                    // query.setParameter("name", request.getUsername())
                    ResultHandle query2 = methodCreator.invokeInterfaceMethod(
                            MethodDescriptor.ofMethod(Query.class, "setParameter", Query.class, String.class, Object.class),
                            query, methodCreator.load("name"), username);

                    // UserEntity user = getSingleUser(query2);
                    user = methodCreator.invokeVirtualMethod(
                            MethodDescriptor.ofMethod(name, "getSingleUser", Object.class, Query.class),
                            methodCreator.getThis(), query2);
                }

                String declaringClassName = jpaSecurityDefinition.annotatedClass.name().toString();
                String declaringClassTypeDescriptor = "L" + declaringClassName.replace('.', '/') + ";";
                AssignableResultHandle userVar = methodCreator.createVariable(declaringClassTypeDescriptor);
                methodCreator.assign(userVar, methodCreator.checkCast(user, declaringClassName));

                // if(user == null) return null;
                try (BytecodeCreator trueBranch = methodCreator.ifNull(userVar).trueBranch()) {
                    trueBranch.returnValue(trueBranch.loadNull());
                }

                // :pass = user.pass | user.getPass()
                ResultHandle pass = jpaSecurityDefinition.password.readValue(methodCreator, userVar);
                String getPasswordMethod;
                if (passwordType == null) {
                    passwordType = PasswordType.MCF.name();
                }
                switch (PasswordType.valueOf(passwordType)) {
                    case CLEAR:
                        getPasswordMethod = "getClearPassword";
                        break;
                    case MCF:
                        getPasswordMethod = "getMcfPassword";
                        break;
                    default:
                        throw new RuntimeException("Unknown password type: " + passwordType);
                }
                // :getPasswordMethod(:pass);
                ResultHandle storedPassword = methodCreator.invokeVirtualMethod(
                        MethodDescriptor.ofMethod(name, getPasswordMethod, org.wildfly.security.password.Password.class,
                                String.class),
                        methodCreator.getThis(), pass);

                // Builder builder = checkPassword(storedPassword, request);
                ResultHandle builder = methodCreator.invokeVirtualMethod(MethodDescriptor.ofMethod(name, "checkPassword",
                        QuarkusSecurityIdentity.Builder.class,
                        org.wildfly.security.password.Password.class,
                        UsernamePasswordAuthenticationRequest.class),
                        methodCreator.getThis(),
                        storedPassword, methodCreator.getMethodParam(1));
                AssignableResultHandle builderVar = methodCreator.createVariable(QuarkusSecurityIdentity.Builder.class);
                methodCreator.assign(builderVar, builder);

                ResultHandle role = jpaSecurityDefinition.roles.readValue(methodCreator, userVar);
                // role: user.getRole()
                boolean handledRole = false;
                Type rolesType = jpaSecurityDefinition.roles.type();
                switch (rolesType.kind()) {
                    case ARRAY:
                        // FIXME: support non-JPA-backed array roles?
                        break;
                    case CLASS:
                        if (rolesType.name().equals(DOTNAME_STRING)) {
                            // addRoles(builder, :role)
                            methodCreator.invokeVirtualMethod(
                                    MethodDescriptor.ofMethod(name, "addRoles", void.class,
                                            QuarkusSecurityIdentity.Builder.class, String.class),
                                    methodCreator.getThis(),
                                    builderVar,
                                    role);
                            handledRole = true;
                        }
                        break;
                    case PARAMETERIZED_TYPE:
                        DotName roleType = rolesType.name();
                        if (roleType.equals(DOTNAME_LIST)
                                || roleType.equals(DOTNAME_COLLECTION)
                                || roleType.equals(DOTNAME_SET)) {
                            Type elementType = rolesType.asParameterizedType().arguments().get(0);
                            String elementClassName = elementType.name().toString();
                            String elementClassTypeDescriptor = "L" + elementClassName.replace('.', '/') + ";";
                            FieldOrMethod rolesFieldOrMethod;
                            if (!elementType.name().equals(DOTNAME_STRING)) {
                                ClassInfo roleClass = index.getClassByName(elementType.name());
                                if (roleClass == null) {
                                    throw new RuntimeException(
                                            "The role element type must be indexed by Jandex: " + elementType);
                                }
                                AnnotationTarget annotatedRolesValue = getSingleAnnotatedElement(index, DOTNAME_ROLES_VALUE);
                                rolesFieldOrMethod = JpaSecurityDefinition.getFieldOrMethod(index, roleClass,
                                        annotatedRolesValue, isPanache(roleClass, panacheClasses));
                                if (rolesFieldOrMethod == null) {
                                    throw new RuntimeException(
                                            "Missing @RoleValue annotation on (non-String) role element type: " + elementType);
                                }
                            } else {
                                rolesFieldOrMethod = null;
                            }
                            // for(:elementType roleElement : :role){
                            //    ret.addRoles(:role.roleField);
                            //    // or for String collections:
                            //    ret.addRoles(:role);
                            // }
                            foreach(methodCreator, role, elementClassTypeDescriptor, (creator, var) -> {
                                ResultHandle roleElement;
                                if (rolesFieldOrMethod != null) {
                                    roleElement = rolesFieldOrMethod.readValue(creator, var);
                                } else {
                                    roleElement = var;
                                }
                                creator.invokeVirtualMethod(
                                        MethodDescriptor.ofMethod(name, "addRoles", void.class,
                                                QuarkusSecurityIdentity.Builder.class, String.class),
                                        methodCreator.getThis(),
                                        builderVar,
                                        roleElement);
                            });
                            handledRole = true;
                        }
                        break;
                }
                if (!handledRole) {
                    throw new RuntimeException("Unsupported @Roles field/getter type: " + rolesType);
                }

                // return builder.build()
                methodCreator.returnValue(methodCreator.invokeVirtualMethod(
                        MethodDescriptor.ofMethod(QuarkusSecurityIdentity.Builder.class,
                                "build",
                                QuarkusSecurityIdentity.class),
                        builderVar));
            }
        }
    }

    private void foreach(MethodCreator method, ResultHandle iterable, String type,
            BiConsumer<BytecodeCreator, AssignableResultHandle> body) {
        ResultHandle iterator = method.invokeInterfaceMethod(
                MethodDescriptor.ofMethod(Iterable.class, "iterator", Iterator.class),
                iterable);
        try (BytecodeCreator loop = method.createScope()) {
            ResultHandle hasNextValue = loop.invokeInterfaceMethod(
                    MethodDescriptor.ofMethod(Iterator.class, "hasNext", boolean.class),
                    iterator);

            BranchResult hasNextBranch = loop.ifNonZero(hasNextValue);
            try (BytecodeCreator hasNext = hasNextBranch.trueBranch()) {
                ResultHandle next = hasNext.invokeInterfaceMethod(
                        MethodDescriptor.ofMethod(Iterator.class, "next", Object.class),
                        iterator);

                AssignableResultHandle var = hasNext.createVariable(type);
                hasNext.assign(var, next);

                body.accept(hasNext, var);
                hasNext.continueScope(loop);
            }
            try (BytecodeCreator doesNotHaveNext = hasNextBranch.falseBranch()) {
                doesNotHaveNext.breakScope(loop);
            }
        }
    }
}
