package fschmidt.util.java;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

// based on Memoizer from Java Concurrency in Practice

public final class Memoizer<A,V> implements Computable<A,V> {
	private final ConcurrentMap<A, Future<V>> cache
			= new ConcurrentHashMap<A, Future<V>>();
	private final Computable<A,V> comp;

	public Memoizer(Computable<A,V> comp) {
		this.comp = comp;
	}

	public V get(final A arg) throws ComputationException {
		while (true) {
			Future<V> f = cache.get(arg);
			if (f == null) {
				Callable<V> eval = new Callable<V>() {
					public V call() throws Exception {
						return comp.get(arg);
					}
				};
				FastFuture<V> ft = new FastFuture<V>(eval);
				f = cache.putIfAbsent(arg, ft);
				if (f == null) {
					f = ft;
					ft.run();
				}
			}
			try {
				return f.get();
			} catch (InterruptedException e) {
				throw new RuntimeException(e);
			} catch (CancellationException e) {
				cache.remove(arg, f);
			} catch (ExecutionException e) {
				throw ComputationException.newInstance(e);
			}
		}
	}

	public Map<A,V> get(final Collection<A> args) throws ComputationException {
		outer:
		while (true) {
			Map<A,Future<V>> fmap = new HashMap<A,Future<V>>();
			for( final A arg : args ) {
				Future<V> f = cache.get(arg);
				if (f == null) {
					Callable<V> eval = new Callable<V>() {
						public V call() throws Exception {
							return comp.get(arg);
						}
					};
					FastFuture<V> ft = new FastFuture<V>(eval);
					f = cache.putIfAbsent(arg, ft);
					if (f == null) {
						f = ft;
						ft.run();
					}
				}
				fmap.put(arg,f);
			}
			Map<A,V> map = new HashMap<A,V>();
			for( Map.Entry<A,Future<V>> entry : fmap.entrySet() ) {
				A arg = entry.getKey();
				Future<V> f = entry.getValue();
				try {
					map.put( arg, f.get() );
				} catch (InterruptedException e) {
					throw new RuntimeException(e);
				} catch (CancellationException e) {
					cache.remove(arg, f);
					continue outer;
				} catch (ExecutionException e) {
					throw ComputationException.newInstance(e);
				}
			}
			return map;
		}
	}

	public void remove(final A arg) {
		cache.remove(arg);
	}
/*
	public void clear() {
		cache.clear();
	}
*/
}