/*
 * Copyright 2006-2007 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.batch.support.transaction;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CopyOnWriteArraySet;

import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.springframework.aop.framework.ProxyFactory;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationAdapter;
import org.springframework.transaction.support.TransactionSynchronizationManager;

/**
 * <p>
 * Factory for transaction aware objects (like lists, sets, maps). If a
 * transaction is active when a method is called on an instance created by the
 * factory, it makes a copy of the target object and carries out all operations
 * on the copy. Only when the transaction commits is the target re-initialised
 * with the copy.
 * </p>
 * 
 * <p>
 * Works well with collections and maps for testing transactional behaviour
 * without needing a database. The base implementation handles lists, sets and
 * maps. Subclasses can implement {@link #begin(Object)} and
 * {@link #commit(Object, Object)} to provide support for other resources.
 * </p>
 * 
 * <p>
 * Generally not intended for multi-threaded use, but the
 * {@link #createAppendOnlyTransactionalMap() append only version} of
 * collections gives isolation between threads operating on different keys in a
 * map, provided they only append to the map. (Threads are limited to removing
 * entries that were created in the same transaction.)
 * </p>
 * 
 * @author Dave Syer
 * 
 */
public class TransactionAwareProxyFactory<T> {

	private final T target;

	private final boolean appendOnly;

	private TransactionAwareProxyFactory(T target) {
		this(target, false);

	}

	private TransactionAwareProxyFactory(T target, boolean appendOnly) {
		super();
		this.target = target;
		this.appendOnly = appendOnly;
	}

	/**
	 * Make a copy of the target that can be used inside a transaction to
	 * isolate changes from the original. Also called from the factory
	 * constructor to isolate the target from the original value passed in.
	 * 
	 * @param target the target object (List, Set or Map)
	 * @return an independent copy
	 */
	@SuppressWarnings({ "unchecked", "rawtypes" })
	protected final T begin(T target) {
		// Unfortunately in Java 5 this method has to synchronized
		// (works OK without in Java 6).
		synchronized (target) {
			if (target instanceof List) {
				if (appendOnly) {
					return (T) new ArrayList();
				}
				return (T) new ArrayList((List) target);
			}
			else if (target instanceof Set) {
				if (appendOnly) {
					return (T) new HashSet();
				}
				return (T) new HashSet((Set) target);
			}
			else if (target instanceof Map) {
				if (appendOnly) {
					return (T) new HashMap();
				}
				return (T) new HashMap((Map) target);
			}
			else {
				throw new UnsupportedOperationException("Cannot copy target for this type: " + target.getClass());
			}
		}
	}

	/**
	 * Take the working copy state and commit it back to the original target.
	 * The target then reflects all the changes applied to the copy during a
	 * transaction.
	 * 
	 * @param copy the working copy.
	 * @param target the original target of the factory.
	 */
	@SuppressWarnings({ "unchecked", "rawtypes" })
	protected void commit(T copy, T target) {
		// Unfortunately in Java 5 this method has to be synchronized
		// (works OK without in Java 6).
		synchronized (target) {
			if (target instanceof Collection) {
				if (!appendOnly) {
					((Collection) target).clear();
				}
				((Collection) target).addAll((Collection) copy);
			}
			else {
				if (!appendOnly) {
					((Map) target).clear();
				}
				((Map) target).putAll((Map) copy);
			}
		}
	}

	private T createInstance() {

		synchronized (target) {

			ProxyFactory factory = new ProxyFactory(target);
			factory.addAdvice(new TransactionAwareInterceptor());
			@SuppressWarnings("unchecked")
			T instance = (T) factory.getProxy();
			return instance;

		}

	}

	public static <K, V> Map<K, V> createTransactionalMap() {
		return (Map<K, V>) new TransactionAwareProxyFactory<ConcurrentHashMap<K, V>>(new ConcurrentHashMap<K, V>()).createInstance();
	}

	public static <K, V> Map<K, V> createTransactionalMap(Map<K, V> map) {
		return (Map<K, V>) new TransactionAwareProxyFactory<ConcurrentHashMap<K, V>>(new ConcurrentHashMap<K, V>(map)).createInstance();
	}

	public static <K, V> ConcurrentMap<K, V> createAppendOnlyTransactionalMap() {
		return new TransactionAwareProxyFactory<ConcurrentHashMap<K, V>>(new ConcurrentHashMap<K, V>(), true).createInstance();
	}

	public static <T> Set<T> createAppendOnlyTransactionalSet() {
		return (Set<T>) new TransactionAwareProxyFactory<CopyOnWriteArraySet<T>>(new CopyOnWriteArraySet<T>(), true).createInstance();
	}

	public static <T> Set<T> createTransactionalSet() {
		return (Set<T>) new TransactionAwareProxyFactory<CopyOnWriteArraySet<T>>(new CopyOnWriteArraySet<T>()).createInstance();
	}

	public static <T> Set<T> createTransactionalSet(Set<T> set) {
		return (Set<T>) new TransactionAwareProxyFactory<CopyOnWriteArraySet<T>>(new CopyOnWriteArraySet<T>(set)).createInstance();
	}

	public static <T> List<T> createAppendOnlyTransactionalList() {
		return (List<T>) new TransactionAwareProxyFactory<CopyOnWriteArrayList<T>>(new CopyOnWriteArrayList<T>(), true).createInstance();
	}

	public static <T> List<T> createTransactionalList() {
		return (List<T>) new TransactionAwareProxyFactory<CopyOnWriteArrayList<T>>(new CopyOnWriteArrayList<T>()).createInstance();
	}

	public static <T> List<T> createTransactionalList(List<T> list) {
		return (List<T>) new TransactionAwareProxyFactory<CopyOnWriteArrayList<T>>(new CopyOnWriteArrayList<T>(list)).createInstance();
	}

	private class TargetSynchronization extends TransactionSynchronizationAdapter {

		private final T cache;

		private final Object key;

		public TargetSynchronization(Object key, T cache) {
			super();
			this.cache = cache;
			this.key = key;
		}

        @Override
		public void afterCompletion(int status) {
			super.afterCompletion(status);
			if (status == TransactionSynchronization.STATUS_COMMITTED) {
				synchronized (target) {
					commit(cache, target);
				}
			}
			TransactionSynchronizationManager.unbindResource(key);
		}
	}

	private class TransactionAwareInterceptor implements MethodInterceptor {

        @Override
		public Object invoke(MethodInvocation invocation) throws Throwable {

			if (!TransactionSynchronizationManager.isActualTransactionActive()) {
				return invocation.proceed();
			}

			T cache;

			if (!TransactionSynchronizationManager.hasResource(this)) {
				cache = begin(target);
				TransactionSynchronizationManager.bindResource(this, cache);
				TransactionSynchronizationManager.registerSynchronization(new TargetSynchronization(this, cache));
			}
			else {
				@SuppressWarnings("unchecked")
				T retrievedCache = (T) TransactionSynchronizationManager.getResource(this);
				cache = retrievedCache;
			}

			Object result = invocation.getMethod().invoke(cache, invocation.getArguments());

			if (appendOnly) {
				String methodName = invocation.getMethod().getName();
				if ((result == null && methodName.equals("get"))
						|| (Boolean.FALSE.equals(result) && (methodName.startsWith("contains")) || (Boolean.TRUE
								.equals(result) && methodName.startsWith("isEmpty")))) {
					// In appendOnly mode the result of a get might not be
					// in the cache...
					return invocation.proceed();
				}
				if (result instanceof Collection<?>) {
					HashSet<Object> set = new HashSet<Object>((Collection<?>) result);
					set.addAll((Collection<?>) invocation.proceed());
					result = set;
				}
			}

			return result;

		}
	}

}
