package com.cosium.spring.data.jpa.entity.graph.repository.support;

import javax.persistence.EntityManager;
import java.util.List;

import com.cosium.spring.data.jpa.entity.graph.domain.EntityGraph;
import com.cosium.spring.data.jpa.entity.graph.domain.EntityGraphUtils;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.springframework.aop.framework.ProxyFactory;
import org.springframework.core.NamedThreadLocal;
import org.springframework.data.jpa.repository.query.JpaEntityGraph;
import org.springframework.data.repository.core.RepositoryInformation;
import org.springframework.data.repository.core.support.RepositoryProxyPostProcessor;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

/**
 * Captures {@link EntityGraph} on repositories method calls.
 * Created on 22/11/16.
 *
 * @author Reda.Housni-Alaoui
 */
class RepositoryMethodEntityGraphExtractor implements RepositoryProxyPostProcessor {

	private static final ThreadLocal<EntityGraphBean> CURRENT_ENTITY_GRAPH =
			new NamedThreadLocal<EntityGraphBean>("Thread local holding the current spring data jpa repository entity graph");

	private final EntityManager entityManager;

	RepositoryMethodEntityGraphExtractor(EntityManager entityManager) {
		this.entityManager = entityManager;
	}

	@Override
	public void postProcess(ProxyFactory factory, RepositoryInformation repositoryInformation) {
		factory.addAdvice(new JpaEntityGraphMethodInterceptor(entityManager, repositoryInformation.getDomainType()));
	}

	static EntityGraphBean getCurrentJpaEntityGraph() {
		return CURRENT_ENTITY_GRAPH.get();
	}

	private static class JpaEntityGraphMethodInterceptor implements MethodInterceptor {

		private static final String DEFAULT_ENTITYGRAPH_NAME_SUFFIX = ".default";
		private final Class domainClass;
		private final EntityGraphBean defaultEntityGraph;

		JpaEntityGraphMethodInterceptor(EntityManager entityManager, Class domainClass) {
			this.domainClass = domainClass;
			this.defaultEntityGraph = findDefaultEntityGraph(entityManager);
		}

		/**
		 * @param entityManager
		 * @return The default entity graph if it exists. Null otherwise.
		 */
		private EntityGraphBean findDefaultEntityGraph(EntityManager entityManager){
			List<javax.persistence.EntityGraph<?>> entityGraphs = entityManager.getEntityGraphs(domainClass);
			for (javax.persistence.EntityGraph entityGraph : entityGraphs) {
				if (entityGraph.getName().endsWith(DEFAULT_ENTITYGRAPH_NAME_SUFFIX)) {
					return buildEntityGraphBean(EntityGraphUtils.fromName(entityGraph.getName()));
				}
			}
			return null;
		}

		@Override
		public Object invoke(MethodInvocation invocation) throws Throwable {
			Object[] arguments = invocation.getArguments();
			EntityGraph entityGraph = null;
			for (Object argument : arguments) {
				if (!(argument instanceof EntityGraph)) {
					continue;
				}
				entityGraph = (EntityGraph) argument;
				break;
			}
			EntityGraphBean entityGraphBean = buildEntityGraphBean(entityGraph);
			if (entityGraphBean != null) {
				CURRENT_ENTITY_GRAPH.set(entityGraphBean);
			}
			try {
				return invocation.proceed();
			} finally {
				if (entityGraphBean != null) {
					CURRENT_ENTITY_GRAPH.remove();
				}
			}
		}

		private EntityGraphBean buildEntityGraphBean(EntityGraph entityGraph) {
			if (EntityGraphUtils.isEmpty(entityGraph)) {
				return defaultEntityGraph;
			}

			Assert.notNull(entityGraph.getEntityGraphType());

			org.springframework.data.jpa.repository.EntityGraph.EntityGraphType type;
			switch (entityGraph.getEntityGraphType()) {
				case FETCH:
					type = org.springframework.data.jpa.repository.EntityGraph.EntityGraphType.FETCH;
					break;
				case LOAD:
					type = org.springframework.data.jpa.repository.EntityGraph.EntityGraphType.LOAD;
					break;
				default:
					throw new RuntimeException("Unknown entity graph type");
			}

			List<String> attributePaths = entityGraph.getEntityGraphAttributePaths();
			JpaEntityGraph jpaEntityGraph = new JpaEntityGraph(
					StringUtils.hasText(entityGraph.getEntityGraphName()) ? entityGraph.getEntityGraphName() : domainClass.getName() + "-_-_-_-_-_-",
					type,
					attributePaths != null ? attributePaths.toArray(new String[attributePaths.size()]) : null
			);

			return new EntityGraphBean(jpaEntityGraph, domainClass);
		}

	}
}

@@ -2289,16 +2289,59 @@
 nager)%7B%0A
+%09%09%09EntityGraphBean entityGraphBean = null;%0A
 %09%09%09List%3C
@@ -2581,22 +2581,236 @@
  %7B%0A%09%09%09%09%09
-return
+if(entityGraphBean != null)%7B%0A%09%09%09%09%09%09throw new RuntimeException(%22Multiple default entity graphs detected : %22 + entityGraph.getName() + %22 and %22 + entityGraphBean.getJpaEntityGraph().getName());%0A%09%09%09%09%09%7D%0A%09%09%09%09%09entityGraphBean =
  buildEn
@@ -2895,20 +2895,31 @@
 %09return 
-null
+entityGraphBean
 ;%0A%09%09%7D%0A%0A%09

Add unique constraint on default method