package ca.uhn.hl7v2.concurrent;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

/**
 * Default Implementation of a {@link BlockingMap}.
 * <p>
 * Note: While it is not actively prevented that more then one thread waits for
 * an entry, it is not guaranteed that all waiting threads will receive the
 * entry once it became available. Other implementations may choose to count the
 * waiting threads and/or to remove an available value after a grace period.
 * 
 * @param <K>
 * @param <V>
 */
public class BlockingHashMap<K, V> implements BlockingMap<K, V> {

        private final ConcurrentMap<K, V> map = new ConcurrentHashMap<K, V>();
        private final ConcurrentMap<K, CountDownLatch> latches = new ConcurrentHashMap<K, CountDownLatch>();
        private final ExecutorService executor;
        
        public BlockingHashMap() {
                this(Executors.newCachedThreadPool());
        }
        
        public BlockingHashMap(ExecutorService executor) {
                super();
                this.executor = executor;
        }

        /**
         * Returns the keys of available entries
         * 
         * @see java.util.Map#keySet()
         */
        public Set<K> keySet() {
                return map.keySet();
        }

        /**
         * Returns an available entry without removing it from the map
         * 
         * @see java.util.Map#get(java.lang.Object)
         */
        public V get(Object key) {
                return map.get(key);
        }

        /**
         * Returns <code>true</code> if an entry with the given key is available
         * 
         * @see java.util.Map#containsKey(java.lang.Object)
         */
        public boolean containsKey(Object key) {
                return map.containsKey(key);
        }

        /**
         * @see java.util.Map#put(java.lang.Object, java.lang.Object)
         */
        synchronized public V put(K key, V value) {
                V result = map.put(key, value);
                latchFor(key).countDown();
                return result;
        }

        /**
         * @see ca.uhn.hl7v2.concurrent.BlockingMap#give(java.lang.Object,
         *      java.lang.Object)
         */
        synchronized public boolean give(K key, V value) {
                if (!latches.containsKey(key)) {
                        return false;
                }
                put(key, value);
                return true;
        }

        public V take(K key) throws InterruptedException {
                latchFor(key).await();
                latches.remove(key);
                return map.remove(key); // likely to fail there are n > 1 consumers
        }
        

        public Future<V> asyncTake(final K key) throws InterruptedException {
                latchFor(key);
                return executor.submit(new Callable<V>() {

                        public V call() throws Exception {
                                return take(key);
                        }
                });
        }

        public V poll(K key, long timeout, TimeUnit unit)
                        throws InterruptedException {
                if (latchFor(key).await(timeout, unit)) {
                        latches.remove(key);
                        return map.remove(key);
                }
                return null;
        }
        
        public Future<V> asyncPoll(final K key, final long timeout, final TimeUnit unit) {
                latchFor(key);
                return executor.submit(new Callable<V>() {

                        public V call() throws Exception {
                                return poll(key, timeout, unit);
                        }
                });             
        }
        

        /**
         * Returns true if no entry is available for consumers
         * 
         * @see java.util.Map#isEmpty()
         */
        public boolean isEmpty() {
                return map.isEmpty();
        }

        /**
         * Returns the number of available values
         * 
         * @see java.util.Map#size()
         */
        public int size() {
                return map.size();
        }

        /**
         * Removes an entry, regardless whether a value has been set or not. Waiting
         * consumers will receive a null value.
         * 
         * @see java.util.Map#remove(java.lang.Object)
         */
        synchronized public V remove(Object key) {
                V result = map.remove(key);
                CountDownLatch latch = latches.remove(key);
                if (latch != null)
                        latch.countDown();
                return result;
        }

        /**
         * Clears all existing entries. Waiting consumers will receive a null value
         * for each removed entry.
         * 
         * @see java.util.Map#clear()
         */
        public void clear() {
                for (K key : latches.keySet()) {
                        remove(key);
                }
        }

        public Collection<V> values() {
                return map.values();
        }

        public Set<java.util.Map.Entry<K, V>> entrySet() {
                return map.entrySet();
        }

        public void putAll(Map<? extends K, ? extends V> t) {
                for (Entry<? extends K, ? extends V> entry : t.entrySet()) {
                        put(entry.getKey(), entry.getValue());
                }
        }

        public boolean containsValue(Object value) {
                return map.containsValue(value);
        }

        private synchronized CountDownLatch latchFor(K key) {
                CountDownLatch latch = latches.get(key);
                if (latch == null) {
                        latch = new CountDownLatch(1);
                        latches.put(key, latch);
                }
                return latch;
        }

}