// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.InteropServices;
 
namespace Microsoft.AspNetCore.Razor.PooledObjects;
 
/// <summary>
///  Wraps a pooled <see cref="HashSet{T}"/> but doesn't allocate it until
///  it's needed. Note: Dispose this to ensure that the pooled set is returned
///  to the pool.
/// </summary>
internal ref struct PooledHashSet<T>
{
    private readonly IEqualityComparer<T> _comparer;
 
    // Used in NET only code below. Called API doesn't exist on framework.
#pragma warning disable IDE0052
    private readonly int? _capacity;
#pragma warning restore IDE0052
 
    private HashSetPool<T>? _pool;
    private (bool hasValue, T value) _item;
    private HashSet<T>? _set;
 
    public PooledHashSet()
    {
        _comparer = EqualityComparer<T>.Default;
    }
 
    public PooledHashSet(int capacity)
    {
        _comparer = EqualityComparer<T>.Default;
        _capacity = capacity;
    }
 
    public PooledHashSet(IEqualityComparer<T> comparer)
    {
        ValidateComparer(comparer);
        _comparer = comparer;
    }
 
    public PooledHashSet(IEqualityComparer<T> comparer, int capacity)
    {
        ValidateComparer(comparer);
        _comparer = comparer;
        _capacity = capacity;
    }
 
    public PooledHashSet(HashSetPool<T> pool, int capacity)
    {
        _comparer = pool.Comparer;
        _pool = pool;
        _capacity = capacity;
    }
 
    public PooledHashSet(HashSetPool<T> pool)
    {
        _comparer = pool.Comparer;
        _pool = pool;
    }
 
    [Conditional("DEBUG")]
    private static void ValidateComparer(IEqualityComparer<T> comparer)
    {
        if (!ReferenceEquals(comparer, EqualityComparer<T>.Default) &&
            !ReferenceEquals(comparer, StringComparer.Ordinal) &&
            !ReferenceEquals(comparer, StringComparer.OrdinalIgnoreCase))
        {
            ThrowHelper.ThrowArgumentException(nameof(comparer),
                "Only EqualityComparer<T>.Default, StringComparer.Ordinal, and StringComparer.OrdinalIgnoreCase are supported. Please provide a pool.");
        }
    }
 
    private readonly IEqualityComparer<T> Comparer => _comparer ?? EqualityComparer<T>.Default;
 
    private readonly bool HasSingleItem => _item.hasValue && _set is null;
 
    [MemberNotNullWhen(false, nameof(_set))]
    private readonly bool SetIsNullOrEmpty => _set is null || _set.Count == 0;
 
    public void Dispose()
    {
        ClearAndFree();
    }
 
    public readonly int Count
    {
        get
        {
            if (HasSingleItem)
            {
                return 1;
            }
 
            if (SetIsNullOrEmpty)
            {
                return 0;
            }
 
            return _set.Count;
        }
    }
 
    public bool Add(T item)
    {
        if (_set is null)
        {
            // Optimized for the single item case.
            if (!HasSingleItem)
            {
                _item.value = item;
                _item.hasValue = true;
                return true;
            }
 
            if (Comparer.Equals(_item.value, item))
            {
                // Duplicate of single item.
                return false;
            }
 
            _set = AcquireHashSet();
        }
 
        return _set.Add(item);
    }
 
    public void ClearAndFree()
    {
        if (_pool is not null && _set is { } set)
        {
            _pool.Return(set);
        }
 
        _set = null;
        _pool = null;
        _item = default;
    }
 
    public readonly bool Contains(T item)
    {
        if (_set is null)
        {
            return _item.hasValue && Comparer.Equals(_item.value, item);
        }
 
        return _set.Contains(item);
    }
 
    public readonly T[] ToArray()
    {
        if (HasSingleItem)
        {
            return [_item.value];
        }
 
        if (SetIsNullOrEmpty)
        {
            return [];
        }
 
        var result = new T[_set.Count];
        _set.CopyTo(result);
 
        return result;
    }
 
    public readonly ImmutableArray<T> ToImmutableArray()
    {
        return ImmutableCollectionsMarshal.AsImmutableArray(ToArray());
    }
 
    public readonly ImmutableArray<T> OrderByAsArray<TKey>(Func<T, TKey> keySelector)
    {
        if (HasSingleItem)
        {
            return [_item.value];
        }
 
        if (SetIsNullOrEmpty)
        {
            return [];
        }
 
        return _set.OrderByAsArray(keySelector);
    }
 
    public void UnionWith(ImmutableArray<T> other)
    {
        if (other.IsDefaultOrEmpty)
        {
            return;
        }
 
        if (other.Length == 1)
        {
            Add(other[0]);
            return;
        }
 
        _set ??= AcquireHashSet();
 
        // Avoid boxing the ImmutableArray<T>
        var array = ImmutableCollectionsMarshal.AsArray(other)!;
 
        _set.UnionWith(array);
    }
 
    public void UnionWith(IReadOnlyList<T>? other)
    {
        if (other is null || other.Count == 0)
        {
            return;
        }
 
        if (other.Count == 1)
        {
            Add(other[0]);
            return;
        }
 
        _set ??= AcquireHashSet();
 
        _set.UnionWith(other);
    }
 
    private HashSet<T> AcquireHashSet()
    {
        Debug.Assert(_set is null);
 
        _pool ??= TrySelectPool(Comparer);
 
        var result = _pool is not null
            ? _pool.Get()
            : new HashSet<T>(Comparer);
 
#if NET
        if (_capacity is int capacity)
        {
            result.EnsureCapacity(capacity);
        }
#endif
 
        if (HasSingleItem)
        {
            result.Add(_item.value);
            _item = default;
        }
 
        return result;
 
        static HashSetPool<T>? TrySelectPool(IEqualityComparer<T> comparer)
        {
            if (ReferenceEquals(comparer, EqualityComparer<T>.Default))
            {
                return HashSetPool<T>.Default;
            }
 
            if (typeof(T) == typeof(string))
            {
                if (ReferenceEquals(comparer, StringComparer.Ordinal))
                {
                    return (HashSetPool<T>)(object)SpecializedPools.StringHashSet.Ordinal;
                }
 
                if (ReferenceEquals(comparer, StringComparer.OrdinalIgnoreCase))
                {
                    return (HashSetPool<T>)(object)SpecializedPools.StringHashSet.OrdinalIgnoreCase;
                }
            }
 
            return null;
        }
    }
}