namespace NesterovskyBros.Utils
{
  using System;
  using System.Collections.Concurrent;
  using System.Collections.Generic;
  using System.Collections.ObjectModel;
  using System.Runtime.CompilerServices;
  using System.Runtime.Serialization;
  using System.Threading;

  /// <summary>
  /// A shared state container.
  /// </summary>
  /// <typeparam name="K">A key type.</typeparam>
  /// <typeparam name="V">A value type.</typeparam>
  public class State<K, V>
  {
    /// <summary>
    /// Creates an instance of a shared state.
    /// </summary>
    /// <param name="protectDictionary">
    /// true to copy and wrap input dictionaries, and 
    /// false to store dictionary passed.
    /// </param>
    /// <param name="keyComparer">Optional key comparer.</param>
    /// <param name="valueComparer">Optional value comparer.</param>
    public State(
      bool protectDictionary = true,
      IEqualityComparer<K> keyComparer = null,
      IEqualityComparer<V> valueComparer = null)
    {
      if (keyComparer == null)
      {
        keyComparer = EqualityComparer<K>.Default;
      }

      if (valueComparer == null)
      {
        valueComparer = EqualityComparer<V>.Default;
      }

      this.protectDictionary = protectDictionary;
      this.keyComparer = keyComparer;
      this.valueComparer = valueComparer;
      this.states = new WeakTable<object, IDictionary<K, V>>(
        new EqualityComparer { self = this });
    }

    /// <summary>
    /// Gets a hash code for a dictionary.
    /// </summary>
    /// <param name="dictionary">A dictionary to get hash code.</param>
    /// <returns>A hash code value.</returns>
    public int GetHashCode(IDictionary<K, V> dictionary)
    {
      if (dictionary == null)
      {
        return 0;
      }

      Tag tag;

      if (tags.TryGetValue(dictionary, out tag))
      {
        return tag.hashCode;
      }

      var hashCode = 0;

      foreach(var entry in dictionary)
      {
        hashCode += GetHashCode(entry.Key, entry.Value);
      }

      return hashCode;
    }

    /// <summary>
    /// Gets a typed value for a dictionary and a key.
    /// </summary>
    /// <typeparam name="T">A value type.</typeparam>
    /// <param name="dictionary">A dictionary instance or null.</param>
    /// <param name="key">A key to get a value for.</param>
    /// <returns>A value or a default if no value is available.</returns>
    public T Get<T>(IDictionary<K, V> dictionary, K key)
      where T: V
    {
      V value;

      if ((dictionary == null) || !dictionary.TryGetValue(key, out value))
      {
        return default(T);
      }

      return (T)value;
    }

    /// <summary>
    /// Sets or remove a typed value for a dictionary and a key.
    /// </summary>
    /// <typeparam name="T">A value type.</typeparam>
    /// <param name="dictionary">A dictionary instance or null.</param>
    /// <param name="key">A key to get a value for.</param>
    /// <param name="value">A value to set.</param>
    /// <returns>A dictionary instance with a value set.</returns>
    public IDictionary<K, V> Set<T>(
      IDictionary<K, V> dictionary, 
      K key, 
      T value)
      where T : V
    {
      return value == null ? Remove(dictionary, key) : 
        Set(dictionary, key, (V)value);
    }

    /// <summary>
    /// Get a shared dictionary that contains the same value,
    /// as an input dictionary.
    /// </summary>
    /// <param name="dictionary">An input dictionary.</param>
    /// <returns>Return a shared dictionary, or null on null input.</returns>
    public IDictionary<K, V> Get(IDictionary<K, V> dictionary)
    {
      if (dictionary == null)
      {
        return null;
      }

      var hashCode = 0;
      Tag tag;

      if (tags.TryGetValue(dictionary, out tag))
      {
        if (tag.key)
        {
          return dictionary;
        }

        hashCode = tag.hashCode;
      }
      else
      {
        foreach(var entry in dictionary)
        {
          hashCode += GetHashCode(entry.Key, entry.Value);
        }
      }

      var stateKey = new GetValue
      {
        self = this,
        dictionary = dictionary,
        hashCode = hashCode
      };

      IDictionary<K, V> stateValue;

      states.TryGetValue(stateKey, out stateValue);

      if (stateValue != null)
      {
        return stateValue;
      }

      dictionary = new Dictionary<K, V>(dictionary, keyComparer);

      if (protectDictionary)
      {
        dictionary = new ReadOnlyDictionary<K, V>(dictionary);
      }

      tag = tags.GetValue(
        dictionary,
        d => new Tag { hashCode = hashCode });

      var result = states.GetValue(dictionary, d => d as IDictionary<K, V>);

      if (result == dictionary)
      {
        Volatile.Write(ref tag.key, true);
      }

      return result;
    }

    /// <summary>
    /// Get a shared dictionary that contains the same value,
    /// as an input dictionary, with a key set to a value.
    /// </summary>
    /// <param name="dictionary">An input dictionary.</param>
    /// <param name="key">A key to set.</param>
    /// <param name="value">A value to set.</param>
    /// <returns>
    /// Return a shared dictionary, or null if result dictionary is empty.
    /// </returns>
    public IDictionary<K, V> Set(
      IDictionary<K, V> dictionary,
      K key, 
      V value)
    {
      var hashCode = GetHashCode(key, value);
      var hasValue = false;
      Tag tag;

      if (dictionary != null)
      {
        V oldValue;

        hasValue = dictionary.TryGetValue(key, out oldValue);
        hashCode = -GetHashCode(key, oldValue);

        if (tags.TryGetValue(dictionary, out tag))
        {
          hashCode += tag.hashCode;
        }
        else
        {
          foreach (var entry in dictionary)
          {
            hashCode += GetHashCode(entry.Key, entry.Value);
          }
        }
      }

      var stateKey = hasValue ?
        new SetValue
        {
          self = this,
          dictionary = dictionary,
          hashCode = hashCode,
          key = key,
          value = value
        } :
        dictionary != null ?
        new AddValue
        {
          self = this,
          dictionary = dictionary,
          hashCode = hashCode,
          key = key,
          value = value
        } :
        new SingleValue 
        {
          self = this,
          hashCode = hashCode,
          key = key,
          value = value
        } as Action;

      IDictionary<K, V> stateValue;

      states.TryGetValue(stateKey, out stateValue);

      if (stateValue != null)
      {
        return stateValue;
      }

      dictionary = dictionary != null ?
        new Dictionary<K, V>(dictionary, keyComparer) :
        new Dictionary<K, V>(keyComparer);

      dictionary[key] = value;

      if (protectDictionary)
      {
        dictionary = new ReadOnlyDictionary<K, V>(dictionary);
      }

      tag = tags.GetValue(
        dictionary,
        d => new Tag { hashCode = hashCode });

      var result = states.GetValue(dictionary, d => d as IDictionary<K, V>);

      if (result == dictionary)
      {
        Volatile.Write(ref tag.key, true);
      }

      return result;
    }

    /// <summary>
    /// Get a shared dictionary that contains the same value,
    /// as an input dictionary, with a key removed.
    /// </summary>
    /// <param name="dictionary">An input dictionary.</param>
    /// <param name="key">A key to remove.</param>
    /// <returns>
    /// Return a shared dictionary, or null if result dictionary is empty.
    /// </returns>
    public IDictionary<K, V> Remove(
      IDictionary<K, V> dictionary,
      K key)
    {
      if (dictionary == null)
      {
        return null;
      }

      V oldValue;

      if (!dictionary.TryGetValue(key, out oldValue))
      {
        return Get(dictionary);
      }

      if (dictionary.Count == 1)
      {
        return null;
      }

      var hashCode = GetHashCode(key, oldValue);
      Tag tag;

      if (tags.TryGetValue(dictionary, out tag))
      {
        hashCode = +tag.hashCode;
      }
      else
      {
        foreach(var entry in dictionary)
        {
          hashCode += GetHashCode(entry.Key, entry.Value);
        }
      }

      var stateKey = new RemoveValue
      {
        self = this,
        dictionary = dictionary,
        hashCode = hashCode,
        key = key
      };

      IDictionary<K, V> stateValue;

      states.TryGetValue(stateKey, out stateValue);

      if (stateValue != null)
      {
        return stateValue;
      }

      dictionary = new Dictionary<K, V>(dictionary, keyComparer);
      dictionary.Remove(key);

      if (protectDictionary)
      {
        dictionary = new ReadOnlyDictionary<K, V>(dictionary);
      }

      tag = tags.GetValue(
        dictionary,
        d => new Tag { hashCode = hashCode });

      var result = states.GetValue(dictionary, d => d as IDictionary<K, V>);

      if (result == dictionary)
      {
        Volatile.Write(ref tag.key, true);
      }

      return result;
    }

    /// <summary>
    /// A key comparer.
    /// </summary>
    public IEqualityComparer<K> KeyComparer
    {
      get { return keyComparer; }
    }

    /// <summary>
    /// A value comparer.
    /// </summary>
    public IEqualityComparer<V> ValueComparer
    {
      get { return valueComparer; }
    }

    /// <summary>
    /// Internal equality comparer.
    /// </summary>
    private class EqualityComparer: IEqualityComparer<object>
    {
      /// <summary>
      /// A reference to the state.
      /// </summary>
      public State<K, V> self;

      /// <summary>
      /// Compares two instances.
      /// </summary>
      /// <param name="x">A first instance to compare.</param>
      /// <param name="y">A second instance to compare.</param>
      /// <returns>
      /// true if instances are considered equal, and false otherwise.
      /// </returns>
      public new bool Equals(object x, object y)
      {
        if (x == y)
        {
          return true;
        }

        if ((x == null) || (y == null))
        {
          return false;
        }

        return ((y is Action) && y.Equals(x)) || ((x is Action) && x.Equals(y));
      }

      /// <summary>
      /// Gets a hash code for the object.
      /// </summary>
      /// <param name="obj">An object to get hash code for.</param>
      /// <returns>a hash code value.</returns>
      public int GetHashCode(object obj)
      {
        var action = obj as Action;

        if (action != null)
        {
          return action.hashCode;
        }

        var dictionary = obj as IDictionary<K, V>;

        if (dictionary == null)
        {
          return 0;
        }

        Tag tag;

        if (self.tags.TryGetValue(dictionary, out tag))
        {
          return tag.hashCode;
        }

        return 0;
      }
    }

    /// <summary>
    /// Internal action class to minimize copying.
    /// </summary>
    private abstract class Action
    {
      /// <summary>
      /// A reference to the state.
      /// </summary>
      public State<K, V> self;

      /// <summary>
      /// A reference dictionary.
      /// </summary>
      public IDictionary<K, V> dictionary;

      /// <summary>
      /// A hash code.
      /// </summary>
      public int hashCode;

      /// <summary>
      /// Gets a hash code.
      /// </summary>
      /// <returns>A hash code value.</returns>
      public override int GetHashCode()
      {
        return hashCode;
      }

      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      public override bool Equals(object obj)
      {
        if (this.dictionary == obj)
        {
          return true;
        }

        var dictionary = obj as IDictionary<K, V>;

        if (dictionary == null)
        {
          return false;
        }

        return Equals(dictionary);
      }

      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected abstract bool Equals(IDictionary<K, V> dictionary);
    }

    /// <summary>
    /// A get value action.
    /// </summary>
    private class GetValue: Action
    {
      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected override bool Equals(IDictionary<K, V> dictionary)
      {
        if (this.dictionary.Count != dictionary.Count)
        {
          return false;
        }

        V value;

        foreach(var entry in this.dictionary)
        {
          if (!dictionary.TryGetValue(entry.Key, out value) ||
            !self.valueComparer.Equals(entry.Value, value))
          {
            return false;
          }
        }

        return true;
      }
    }

    /// <summary>
    /// Init dictionary with a single value.
    /// </summary>
    private class SingleValue : Action
    {
      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected override bool Equals(IDictionary<K, V> dictionary)
      {
        if (dictionary.Count != 1)
        {
          return false;
        }

        V value;

        if (!dictionary.TryGetValue(this.key, out value) ||
          !self.valueComparer.Equals(this.value, value))
        {
          return false;
        }

        return true;
      }

      /// <summary>
      /// Added key.
      /// </summary>
      public K key;

      /// <summary>
      /// Added value.
      /// </summary>
      public V value;
    }

    /// <summary>
    /// Add value to a dictionary.
    /// </summary>
    private class AddValue : Action
    {
      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected override bool Equals(IDictionary<K, V> dictionary)
      {
        if (this.dictionary.Count + 1 != dictionary.Count)
        {
          return false;
        }

        V value;

        foreach(var entry in this.dictionary)
        {
          if (!dictionary.TryGetValue(entry.Key, out value) ||
            !self.valueComparer.Equals(entry.Value, value))
          {
            return false;
          }
        }

        if (!dictionary.TryGetValue(this.key, out value) ||
          !self.valueComparer.Equals(this.value, value))
        {
          return false;
        }

        return true;
      }

      /// <summary>
      /// Added key.
      /// </summary>
      public K key;

      /// <summary>
      /// Added value.
      /// </summary>
      public V value;
    }

    /// <summary>
    /// Set dictionary value.
    /// </summary>
    private class SetValue : Action
    {
      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected override bool Equals(IDictionary<K, V> dictionary)
      {
        if (this.dictionary.Count != dictionary.Count)
        {
          return false;
        }

        var found = false;
        V value;

        foreach(var entry in this.dictionary)
        {
          if (!found && self.keyComparer.Equals(entry.Key, this.key))
          {
            if (!dictionary.TryGetValue(this.key, out value) ||
              !self.valueComparer.Equals(this.value, value))
            {
              return false;
            }

            found = true;
          }
          else
          {
            if (!dictionary.TryGetValue(entry.Key, out value) ||
              !self.valueComparer.Equals(entry.Value, value))
            {
              return false;
            }
          }
        }

        return true;
      }

      /// <summary>
      /// Added key.
      /// </summary>
      public K key;

      /// <summary>
      /// Added value.
      /// </summary>
      public V value;
    }

    /// <summary>
    /// Remove dictionary value.
    /// </summary>
    private class RemoveValue: Action
    {
      /// <summary>
      /// Compares action with dictionary.
      /// </summary>
      /// <param name="obj">A dictionary instance to compare with.</param>
      /// <returns>
      /// true if action and dictionary are considered equal, and false otherwise.
      /// </returns>
      protected override bool Equals(IDictionary<K, V> dictionary)
      {
        if (this.dictionary.Count - 1 != dictionary.Count)
        {
          return false;
        }

        var found = false;
        V value;

        foreach (var entry in this.dictionary)
        {
          if (!found && self.keyComparer.Equals(entry.Key, this.key))
          {
            found = true;

            continue;
          }

          if (!dictionary.TryGetValue(entry.Key, out value) ||
            !self.valueComparer.Equals(entry.Value, value))
          {
            return false;
          }
        }

        return true;
      }

      /// <summary>
      /// Removed key.
      /// </summary>
      public K key;
    }

    /// <summary>
    /// Tag to cache hash code for the dictionary.
    /// </summary>
    private class Tag
    {
      /// <summary>
      /// Dictionary hash code.
      /// </summary>
      public int hashCode;

      /// <summary>
      /// A key indicator.
      /// </summary>
      public bool key;
    }

    /// <summary>
    /// Gets a hash code for a key and value.
    /// </summary>
    /// <param name="key">A key.</param>
    /// <param name="value">A value.</param>
    /// <returns>A hash code of the (key, value) entry.</returns>
    private int GetHashCode(K key, V value)
    {
      if ((key == null) || (value == null))
      {
        return 0;
      }

      var k = keyComparer.GetHashCode(key);
      var v = valueComparer.GetHashCode(value);

      k = ((k << 5) + k);
      v = (v << (k >> 3)) + v;

      return k ^ v;
    }

    /// <summary>
    /// Indicates whether to copy dictionaries kept in the storage.
    /// </summary>
    private bool protectDictionary;

    /// <summary>
    /// A key comparer.
    /// </summary>
    private IEqualityComparer<K> keyComparer;

    /// <summary>
    /// A value comparer.
    /// </summary>
    private IEqualityComparer<V> valueComparer;

    /// <summary>
    /// States stored in the dictionary.
    /// </summary>
    private WeakTable<object, IDictionary<K, V>> states;

    /// <summary>
    /// Tags attached to dictionaries.
    /// </summary>
    private ConditionalWeakTable<IDictionary<K, V>, Tag> tags =
      new ConditionalWeakTable<IDictionary<K, V>, Tag>();
  }
}