using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace ln.collections
{
	public class WeakKeyDictionary<K, V> : IDictionary<K, V> where K : class
	{
		BTreeValueList<int, WeakKeyValuePair<K, V>> store = new BTreeValueList<int, WeakKeyValuePair<K, V>>();

		public WeakKeyDictionary()
		{
		}

		public V this[K key]
		{
			get
			{  
				if (TryGetValue(key, out V value))
					return value;
				throw new KeyNotFoundException();
			}
			set
			{
				SetValue(key, value, true);
			}
		}

        public virtual bool KeyEquals(K key1,K key2)
        {
            return object.Equals(key1, key2);
        }
        public virtual int GetKeyHashcode(K key) => key.GetHashCode();

        private WeakKeyValuePair<K, V> FindKeyValuePair(K key)
		{
            int keyHashCode = GetKeyHashcode(key);

			if (store.ContainsKey(keyHashCode))
			foreach (WeakKeyValuePair<K, V> weakKeyValuePair in store[keyHashCode].ToArray())
			{
				K weakKey = weakKeyValuePair.Key;
				if (weakKey == null)
				{
					store.Remove(keyHashCode, weakKeyValuePair);
				}
				else if (KeyEquals(key, weakKey))
				{
					return weakKeyValuePair;
				}
			}
			return null;
		}

		public bool TryGetValue(K key, out V value)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(key);
			if (weakKeyValuePair != null)
			{
				value = weakKeyValuePair.Value;
				return true;
			}
			value = default(V);
			return false;
		}

        public K GetKeyInstance(K key)
        {
            WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(key);
            return weakKeyValuePair.Key;
        }

        public void SetValue(K key, V value) => SetValue(key, value, false);
		public bool SetValue(K key, V value, bool replace)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(key);
			if (weakKeyValuePair == null)
			{
                int keyHashCode = GetKeyHashcode(key);
                weakKeyValuePair = new WeakKeyValuePair<K, V>(keyHashCode, key, value);
				store.Add(keyHashCode, weakKeyValuePair);
				return false;
			} else if (replace)
			{
				weakKeyValuePair.Value = value;
				return true;
			} else
			{
				throw new ArgumentException("Key already present");
			}
		}

		public ICollection<K> Keys
		{
			get
			{
				List<K> keys = new List<K>();

				foreach (WeakKeyValuePair<K,V> weakKeyValuePair in store.Values.ToArray())
				{
					if (weakKeyValuePair.IsStrong)
						keys.Add(weakKeyValuePair.Key);
					else
						store.TryRemove(weakKeyValuePair.keyHashCode, weakKeyValuePair);
				}

				return keys;
			}
		}

		public ICollection<V> Values
		{
			get
			{
				List<V> values = new List<V>();

				foreach (WeakKeyValuePair<K, V> weakKeyValuePair in store.Values)
				{
					if (weakKeyValuePair.IsStrong)
						values.Add(weakKeyValuePair.Value);
					else
						store.Remove(weakKeyValuePair.keyHashCode, weakKeyValuePair);
				}

				return values;
			}
		}

		public int Count => store.Keys.Count();
		public bool IsReadOnly => false;

		public void Add(K key, V value)
		{
			SetValue(key, value);
		}
		public void Add(KeyValuePair<K, V> item) => Add(item.Key, item.Value);

		public void Clear() => store.Clear();
		public bool Contains(KeyValuePair<K, V> item)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(item.Key);
			if ((weakKeyValuePair != null) && weakKeyValuePair.IsStrong)
			{
				return Object.Equals(item.Value, weakKeyValuePair.Value);
			}
			return false;
		}
		public bool ContainsKey(K key)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(key);
			return (weakKeyValuePair != null) && weakKeyValuePair.IsStrong;
		}

		public void CopyTo(KeyValuePair<K, V>[] array, int arrayIndex)
		{
			foreach (WeakKeyValuePair<K, V> weakKeyValuePair in store.Values)
			{
				if (weakKeyValuePair.IsStrong)
					array[arrayIndex++] = new KeyValuePair<K, V>(weakKeyValuePair.Key,weakKeyValuePair.Value);
			}
		}

		public IEnumerator<KeyValuePair<K, V>> GetEnumerator()
		{
			foreach (WeakKeyValuePair<K,V> weakKeyValuePair in store.Values)
			{
				if (weakKeyValuePair.IsStrong)
					yield return new KeyValuePair<K, V>(weakKeyValuePair.Key,weakKeyValuePair.Value);
			}
		}

		public bool Remove(K key)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(key);
			if (weakKeyValuePair != null)
			{
				store.Remove(weakKeyValuePair.keyHashCode, weakKeyValuePair);
				return true;
			}
			return false;
		}

		public bool Remove(KeyValuePair<K, V> item)
		{
			WeakKeyValuePair<K, V> weakKeyValuePair = FindKeyValuePair(item.Key);
			if ((weakKeyValuePair != null) && (object.Equals(item.Value, weakKeyValuePair.Value)))
			{
				store.Remove(weakKeyValuePair.keyHashCode, weakKeyValuePair);
				return true;
			}
			return false;
		}

		IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

		class WeakKeyValuePair<TK,TV> where TK: class
		{
			public readonly int keyHashCode;

			WeakReference<TK> reference;

			public TK Key => reference.TryGetTarget(out TK target) ? target : null;
			public bool IsStrong => reference.TryGetTarget(out TK target);

			public V Value { get; set; }

			public WeakKeyValuePair(int keyHashCode,TK key,V value)
			{
				reference = new WeakReference<TK>(key);
                this.keyHashCode = keyHashCode;
				Value = value;
			}

			public override int GetHashCode() => keyHashCode;
			public override bool Equals(object obj)
			{
                if (ReferenceEquals(this, obj))
                    return true;

				TK key = Key;
				if ((key != null) && (obj is WeakKeyValuePair<TK,TV>))
				{
					WeakKeyValuePair<TK, TV> other = obj as WeakKeyValuePair<TK, TV>;
					return Equals(key, other.Key);
				}
				return false;
			}
		}


	}
}