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<>();
  	private final ConcurrentMap<K, CountDownLatch> latches = new ConcurrentHashMap<>();
  	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) {
  		latchFor(key);
  		return executor.submit(() -> 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(() -> 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;
 	}
 
 }