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

namespace InABox.Core
{
    public interface IMutableLookup<TKey, TElement> : ILookup<TKey, TElement>
    {
        void Add(TKey key, TElement element);
        void AddRange(TKey key, IEnumerable<TElement> elements);
        void Remove(TKey key);
    }

    public class MutableLookup<TKey, TElement> : IMutableLookup<TKey, TElement>
    {
        private readonly Dictionary<TKey, List<TElement>> _data = new Dictionary<TKey, List<TElement>>();

        public MutableLookup()
        {
        }

        public MutableLookup(ILookup<TKey, TElement> source)
        {
            foreach (var grouping in source)
            foreach (var element in grouping)
                Add(grouping.Key, element);
        }

        public IEnumerable<TElement> All => (from key in _data.Keys
                select _data[key])
            .SelectMany(list => list);

        public int Count => All.Count();

        public IEnumerable<TElement> this[TKey key]
        {
            get
            {
                List<TElement> result;
                if (_data.TryGetValue(key, out result))
                    return result;
                return Array.Empty<TElement>();
            }
        }

        public bool Contains(TKey key)
        {
            return _data.ContainsKey(key);
        }

        public void Add(TKey key, TElement element)
        {
            if (!_data.TryGetValue(key, out var list))
            {
                list = new List<TElement>();
                _data.Add(key, list);
            }

            if (!list.Contains(element))
                list.Add(element);
        }

        public void AddRange(TKey key, IEnumerable<TElement> elements)
        {
            if (!_data.TryGetValue(key, out var list))
            {
                list = new List<TElement>();
                _data.Add(key, list);
            }

            list.AddRange(elements.Where(x => list.Contains(x)));
        }

        public void Remove(TKey key)
        {
            _data.Remove(key);
        }

        public IEnumerator<IGrouping<TKey, TElement>> GetEnumerator()
        {
            return GetGroupings().GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return (GetGroupings() as IEnumerable).GetEnumerator();
        }

        private IEnumerable<IGrouping<TKey, TElement>> GetGroupings()
        {
            return from key in _data.Keys
                select new LookupDictionaryGrouping<TKey, TElement>
                {
                    Key = key,
                    Elements = _data[key]
                } as IGrouping<TKey, TElement>;
        }
    }

    public class LookupDictionaryGrouping<TKey, TElement> : IGrouping<TKey, TElement>
    {
        public IEnumerable<TElement> Elements { get; set; }
        public TKey Key { get; set; }

        public IEnumerator<TElement> GetEnumerator()
        {
            return Elements.GetEnumerator();
        }


        IEnumerator IEnumerable.GetEnumerator()
        {
            return (Elements as IEnumerable).GetEnumerator();
        }
    }
}