package anastore.util;

import java.lang.ref.*;
import java.util.*;

/**
 * A weak value hash map.  An entries in this map will be
 * automatically deleted when the value is no longer reachable.  This
 * is particularly useful for caches.
 *<p>
 * As for the standard weak key map, the behavior of this can be
 * somewhat unexpected.  For example, subsequent calls to size may
 * return different values.
 *<p>
 * Note that this implementation is <b>not</b> synchronized.
 *<p>
 * Also, be careful, because this hasn't been tested at all.
 */
public class WeakValueMap<K, V> extends AbstractMap<K, V>
{
    private final Map<K, Ref<K, V>> _map = new HashMap<K, Ref<K, V>>();
    private final ReferenceQueue<V> _q = new ReferenceQueue<V>();

    private final EntrySet _entrySet = new EntrySet();

    private static class Ref<K, V> extends WeakReference<V>
    {
        private final K _key;
        private boolean removed;

        public Ref(K key, V value, ReferenceQueue<V> q)
        {
            super(value, q);
            _key = key;
            removed = false;
        }
    }

    public int size()
    {
        preen();
        return _map.size();
    }

    public boolean containsKey(Object key)
    {
        return _map.containsKey(key);
    }

    public V get(Object key)
    {
        preen();
        Ref<K, V> ref = _map.get(key);
        if (ref == null)
            return null;
        return ref.get();
    }

    public V put(K key, V value)
    {
        if (value == null)
            throw new NullPointerException();
        preen();
        Ref<K, V> cur = _map.put(key, new Ref<K, V>(key, value, _q));
        if (cur == null)
            return null;
        cur.removed = true;
        return cur.get();
    }

    public V remove(Object key)
    {
        // Don't preen
        Ref<K, V> cur = _map.remove(key);
        if (cur == null)
            return null;
        cur.removed = true;
        return cur.get();
    }

    public void clear()
    {
        // No need to preen
        for (Ref<K,V> ref : _map.values())
            ref.removed = true;
        _map.clear();
    }

    public Set<K> keySet()
    {
        return _map.keySet();
    }

    private class EntrySet extends AbstractSet<Map.Entry<K,V>>
    {
        public Iterator<Map.Entry<K,V>> iterator()
        {
            return new EntryIterator();
        }

        public int size()
        {
            return WeakValueMap.this.size();
        }

        public boolean contains(Object o)
        {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry o2 = (Map.Entry)o;
            V val = WeakValueMap.this.get(o2.getKey());
            if (val == null)
                return false;
            return val.equals(o2.getValue());
        }

        public boolean add(Map.Entry<K, V> o)
        {
            V old = WeakValueMap.this.put(o.getKey(), o.getValue());
            if (old == null)
                return false;
            return !old.equals(o.getValue());
        }

        public boolean remove(Object o)
        {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry o2 = (Map.Entry)o;
            return (WeakValueMap.this.remove(o2.getValue()) != null);
        }

        public void clear()
        {
            WeakValueMap.this.clear();
        }
    }

    private class EntryIterator implements Iterator<Map.Entry<K,V>>
    {
        private final Iterator<Map.Entry<K, Ref<K, V>>> _it =
            _map.entrySet().iterator();
        private boolean _nextValid = false;
        private Ref<K, V> _nextRef = null;
        private V _nextVal = null;
        private Ref<K, V> _cur = null;

        public boolean hasNext()
        {
            if (!_nextValid) {
                nextImpl();
            }
            return _nextVal != null;
        }

        public Map.Entry<K, V> next()
            throws NoSuchElementException
        {
            V val;
            if (!_nextValid) {
                nextImpl();
            }
            _cur = _nextRef;
            val = _nextVal;
            _nextValid = false;
            _nextRef = null;
            _nextVal = null;
            if (val == null)
                throw new NoSuchElementException();
            return new Entry(_cur._key, val);
        }

        private void nextImpl()
        {
            do {
                Map.Entry<K, Ref<K, V>> er;
                try {
                    er = _it.next();
                } catch (NoSuchElementException e) {
                    _nextRef = null;
                    _nextVal = null;
                    break;
                }
                _nextRef = er.getValue();
                _nextVal = _nextRef.get();
            } while (_nextVal == null);
            _nextValid = true;
        }

        public void remove()
        {
            if (_cur == null)
                throw new IllegalStateException("No element to remove");
            if (!_cur.removed) {
                _map.remove(_cur._key);
                _cur.removed = true;
            }
            _cur = null;
        }
    }

    private class Entry implements Map.Entry<K, V>
    {
        private final K _key;
        private V _value;

        public Entry(K key, V value)
        {
            _key = key;
            _value = value;
        }

        public K getKey()
        {
            return _key;
        }

        public V getValue()
        {
            return _value;
        }

        public V setValue(V value)
        {
            _value = value;
            return WeakValueMap.this.put(_key, value);
        }

        public boolean equals(Object o)
        {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry e1 = this;
            Map.Entry e2 = (Map.Entry)o;

            return ((e1.getKey()==null ?
                     e2.getKey()==null : e1.getKey().equals(e2.getKey())) &&
                    (e1.getValue()==null ?
                     e2.getValue()==null : e1.getValue().equals(e2.getValue())));
        }

        public int hashCode()
        {
            return ((getKey()==null   ? 0 : getKey().hashCode()) ^
                    (getValue()==null ? 0 : getValue().hashCode()));
        }
    }

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

    @SuppressWarnings("unchecked")
    private void preen()
    {
        Ref<K, V> ref;
        while ((ref = (Ref<K,V>)_q.poll()) != null) {
            if (!ref.removed) {
                _map.remove(ref._key);
                ref.removed = true;
            }
        }
    }
}