// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

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

namespace Roslyn.Utilities
{
    internal sealed class ImmutableSetWithInsertionOrder<T> : IEnumerable<T>
        where T : notnull
    {
        public static readonly ImmutableSetWithInsertionOrder<T> Empty = new ImmutableSetWithInsertionOrder<T>(ImmutableDictionary.Create<T, uint>(), 0u);

        private readonly ImmutableDictionary<T, uint> _map;
        private readonly uint _nextElementValue;

        private ImmutableSetWithInsertionOrder(ImmutableDictionary<T, uint> map, uint nextElementValue)
        {
            _map = map;
            _nextElementValue = nextElementValue;
        }

        public int Count
        {
            get { return _map.Count; }
        }

        public bool Contains(T value)
        {
            return _map.ContainsKey(value);
        }

        public ImmutableSetWithInsertionOrder<T> Add(T value)
        {
            // no reason to cause allocations if value is already in the set
            if (_map.ContainsKey(value))
            {
                return this;
            }

            return new ImmutableSetWithInsertionOrder<T>(_map.Add(value, _nextElementValue), _nextElementValue + 1u);
        }

        public ImmutableSetWithInsertionOrder<T> AddRange(List<T> values)
        {
            ImmutableDictionary<T, uint>.Builder? builder = null;
            var nextElementValue = _nextElementValue;

            foreach (var value in values)
            {
                // no reason to cause allocations if value is already in the set
                if (builder == null)
                {
                    if (_map.ContainsKey(value))
                    {
                        continue;
                    }

                    builder = _map.ToBuilder();
                }
                else if (builder.ContainsKey(value))
                {
                    continue;
                }

                builder.Add(value, nextElementValue);

                nextElementValue++;
            }

            if (builder == null)
            {
                return this;
            }

            return new ImmutableSetWithInsertionOrder<T>(builder.ToImmutable(), nextElementValue);
        }

        public ImmutableSetWithInsertionOrder<T> Remove(T value)
        {
            var modifiedMap = _map.Remove(value);
            if (modifiedMap == _map)
            {
                // no reason to cause allocations if value is missing
                return this;
            }

            return this.Count == 1 ? Empty : new ImmutableSetWithInsertionOrder<T>(modifiedMap, _nextElementValue);
        }

        public ImmutableSetWithInsertionOrder<T> RemoveRange(List<T> values)
        {
            ImmutableDictionary<T, uint>.Builder? builder = null;

            foreach (var value in values)
            {
                // no reason to cause allocations if value is missing
                if (builder == null)
                {
                    if (!_map.ContainsKey(value))
                    {
                        continue;
                    }

                    builder = _map.ToBuilder();
                }

                builder.Remove(value);
            }

            if (builder == null)
            {
                return this;
            }

            return new ImmutableSetWithInsertionOrder<T>(builder.ToImmutable(), _nextElementValue);
        }

        public IEnumerable<T> InInsertionOrder
        {
            get { return _map.OrderBy(kv => kv.Value).Select(kv => kv.Key); }
        }

        public override string ToString()
        {
            return "{" + string.Join(", ", this) + "}";
        }

        public IEnumerator<T> GetEnumerator()
        {
            return _map.Keys.GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return _map.Keys.GetEnumerator();
        }
    }
}