// 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.Generic;
using System.Threading;
using System.Threading.Tasks;
namespace System.Linq
{
    public static partial class AsyncEnumerable
    {
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// A function to extract the key for each element.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  is null.
        public static IAsyncEnumerable> GroupBy(this IAsyncEnumerable source, Func keySelector) =>
            new GroupedAsyncEnumerable(source, keySelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function and comparer.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// A function to extract the key for each element.
        /// An equality comparer to compare keys with.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  is null.
        public static IAsyncEnumerable> GroupBy(this IAsyncEnumerable source, Func keySelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerable(source, keySelector, comparer);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector) =>
            new GroupedAsyncEnumerableWithTask(source, keySelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function and comparer.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An equality comparer to compare keys with.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerableWithTask(source, keySelector, comparer);
#if !NO_DEEP_CANCELLATION
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector) =>
            new GroupedAsyncEnumerableWithTaskAndCancellation(source, keySelector, comparer: null);
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerableWithTaskAndCancellation(source, keySelector, comparer);
#endif
        /// 
        /// Groups the elements of an async-enumerable sequence and selects the resulting elements by using a specified function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of the elements within the groups computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// A function to extract the key for each element.
        /// A function to map each source element to an element in an async-enumerable group.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  is null.
        public static IAsyncEnumerable> GroupBy(this IAsyncEnumerable source, Func keySelector, Func elementSelector) =>
            new GroupedAsyncEnumerable(source, keySelector, elementSelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function and comparer and selects the resulting elements by using a specified function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of the elements within the groups computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// A function to extract the key for each element.
        /// A function to map each source element to an element in an async-enumerable group.
        /// An equality comparer to compare keys with.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  or  is null.
        public static IAsyncEnumerable> GroupBy(this IAsyncEnumerable source, Func keySelector, Func elementSelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerable(source, keySelector, elementSelector, comparer);
        /// 
        /// Groups the elements of an async-enumerable sequence and selects the resulting elements by using a specified function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of the elements within the groups computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to map each source element to an element in an async-enumerable group.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector) =>
            new GroupedAsyncEnumerableWithTask(source, keySelector, elementSelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence and selects the resulting elements by using a specified function.
        /// 
        /// The type of the elements in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of the elements within the groups computed for each element in the source sequence.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to map each source element to an element in an async-enumerable group.
        /// An equality comparer to use to compare keys.
        /// A sequence of async-enumerable groups, each of which corresponds to a unique key value, containing all elements that share that same key value.
        ///  or  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerableWithTask(source, keySelector, elementSelector, comparer);
#if !NO_DEEP_CANCELLATION
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector) =>
            new GroupedAsyncEnumerableWithTaskAndCancellation(source, keySelector, elementSelector, comparer: null);
        [GenerateAsyncOverload]
        private static IAsyncEnumerable> GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, IEqualityComparer? comparer) =>
            new GroupedAsyncEnumerableWithTaskAndCancellation(source, keySelector, elementSelector, comparer);
#endif
        public static IAsyncEnumerable GroupBy(this IAsyncEnumerable source, Func keySelector, Func, TResult> resultSelector) =>
            new GroupedResultAsyncEnumerable(source, keySelector, resultSelector, comparer: null);
        public static IAsyncEnumerable GroupBy(this IAsyncEnumerable source, Func keySelector, Func, TResult> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerable(source, keySelector, resultSelector, comparer);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function, and then applies a result selector function to each group.
        /// 
        /// Type of element in the source sequence.
        /// Type of the grouping key computed for each element in the source sequence.
        /// The result type returned by the result selector function.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to transform each group into the result type.
        /// An async-enumerable sequence of results obtained by invoking and awaiting the result-selector function on each group.
        ///  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func, ValueTask> resultSelector) =>
            new GroupedResultAsyncEnumerableWithTask(source, keySelector, resultSelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key selector function, and then applies a result selector function to each group.
        /// 
        /// Type of element in the source sequence.
        /// Type of the grouping key computed for each element in the source sequence.
        /// The result type returned by the result selector function.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to transform each group into the result type.
        /// An equality comparer to use to compare keys.
        /// An async-enumerable sequence of results obtained by invoking and awaiting the result-selector function on each group.
        ///  or  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func, ValueTask> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerableWithTask(source, keySelector, resultSelector, comparer);
#if !NO_DEEP_CANCELLATION
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func, CancellationToken, ValueTask> resultSelector) =>
            new GroupedResultAsyncEnumerableWithTaskAndCancellation(source, keySelector, resultSelector, comparer: null);
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func, CancellationToken, ValueTask> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerableWithTaskAndCancellation(source, keySelector, resultSelector, comparer);
#endif
        public static IAsyncEnumerable GroupBy(this IAsyncEnumerable source, Func keySelector, Func elementSelector, Func, TResult> resultSelector) =>
            new GroupedResultAsyncEnumerable(source, keySelector, elementSelector, resultSelector, comparer: null);
        public static IAsyncEnumerable GroupBy(this IAsyncEnumerable source, Func keySelector, Func elementSelector, Func, TResult> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerable(source, keySelector, elementSelector, resultSelector, comparer);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key-selector function, applies an element selector to each element of each group, then applies a result selector to each transformed group.
        /// 
        /// The type of element in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of element computed by the element selector.
        /// The type of the final result, computed by applying the result selector to each transformed group of elements.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to apply to each element of each group. 
        /// An asynchronous function to transform each group into the result type.
        /// An async-enumerable sequence of results obtained by invoking the result selector function on each group and awaiting the result.
        ///  or  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, Func, ValueTask> resultSelector) =>
            new GroupedResultAsyncEnumerableWithTask(source, keySelector, elementSelector, resultSelector, comparer: null);
        /// 
        /// Groups the elements of an async-enumerable sequence according to a specified key-selector function, applies an element selector to each element of each group, then applies a result selector to each transformed group.
        /// 
        /// The type of element in the source sequence.
        /// The type of the grouping key computed for each element in the source sequence.
        /// The type of element computed by the element selector.
        /// The type of the final result, computed by applying the result selector to each transformed group of elements.
        /// An async-enumerable sequence whose elements to group.
        /// An asynchronous function to extract the key for each element.
        /// An asynchronous function to apply to each element of each group. 
        /// An asynchronous function to transform each group into the result type.
        /// An equality comparer to use to compare keys.
        /// An async-enumerable sequence of results obtained by invoking the result selector function on each group and awaiting the result.
        ///  or  or  or  or  is .
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, Func, ValueTask> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerableWithTask(source, keySelector, elementSelector, resultSelector, comparer);
#if !NO_DEEP_CANCELLATION
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, Func, CancellationToken, ValueTask> resultSelector) =>
            new GroupedResultAsyncEnumerableWithTaskAndCancellation(source, keySelector, elementSelector, resultSelector, comparer: null);
        [GenerateAsyncOverload]
        private static IAsyncEnumerable GroupByAwaitWithCancellationCore(this IAsyncEnumerable source, Func> keySelector, Func> elementSelector, Func, CancellationToken, ValueTask> resultSelector, IEqualityComparer? comparer) =>
            new GroupedResultAsyncEnumerableWithTaskAndCancellation(source, keySelector, elementSelector, resultSelector, comparer);
#endif
        private sealed class GroupedResultAsyncEnumerable : AsyncIterator, IAsyncIListProvider
        {
            private readonly IAsyncEnumerable _source;
            private readonly Func _keySelector;
            private readonly Func, TResult> _resultSelector;
            private readonly IEqualityComparer? _comparer;
            private Internal.Lookup? _lookup;
            private IEnumerator? _enumerator;
            public GroupedResultAsyncEnumerable(IAsyncEnumerable source, Func keySelector, Func, TResult> resultSelector, IEqualityComparer? comparer)
            {
                _source = source ?? throw Error.ArgumentNull(nameof(source));
                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
                _comparer = comparer;
            }
            public override AsyncIteratorBase Clone()
            {
                return new GroupedResultAsyncEnumerable(_source, _keySelector, _resultSelector, _comparer);
            }
            public override async ValueTask DisposeAsync()
            {
                if (_enumerator != null)
                {
                    _enumerator.Dispose();
                    _enumerator = null;
                    _lookup = null;
                }
                await base.DisposeAsync().ConfigureAwait(false);
            }
            protected override async ValueTask MoveNextCore()
            {
                switch (_state)
                {
                    case AsyncIteratorState.Allocated:
                        _lookup = await Internal.Lookup.CreateAsync(_source, _keySelector, _comparer, _cancellationToken).ConfigureAwait(false);
                        _enumerator = _lookup.ApplyResultSelector(_resultSelector).GetEnumerator();
                        _state = AsyncIteratorState.Iterating;
                        goto case AsyncIteratorState.Iterating;
                    case AsyncIteratorState.Iterating:
                        if (_enumerator!.MoveNext())
                        {
                            _current = _enumerator.Current;
                            return true;
                        }
                        await DisposeAsync().ConfigureAwait(false);
                        break;
                }
                return false;
            }
            public async ValueTask ToArrayAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.Lookup.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return l.ToArray(_resultSelector);
            }
            public async ValueTask> ToListAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.Lookup.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return l.ToList(_resultSelector);
            }
            public ValueTask GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
            {
                if (onlyIfCheap)
                {
                    return new ValueTask(-1);
                }
                return Core();
                async ValueTask Core()
                {
                    var l = await Internal.Lookup.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                    return l.Count;
                }
            }
        }
        private sealed class GroupedResultAsyncEnumerableWithTask : AsyncIterator, IAsyncIListProvider
        {
            private readonly IAsyncEnumerable _source;
            private readonly Func> _keySelector;
            private readonly Func, ValueTask> _resultSelector;
            private readonly IEqualityComparer? _comparer;
            private Internal.LookupWithTask? _lookup;
            private IAsyncEnumerator? _enumerator;
            public GroupedResultAsyncEnumerableWithTask(IAsyncEnumerable source, Func> keySelector, Func, ValueTask> resultSelector, IEqualityComparer? comparer)
            {
                _source = source ?? throw Error.ArgumentNull(nameof(source));
                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
                _comparer = comparer;
            }
            public override AsyncIteratorBase Clone()
            {
                return new GroupedResultAsyncEnumerableWithTask(_source, _keySelector, _resultSelector, _comparer);
            }
            public override async ValueTask DisposeAsync()
            {
                if (_enumerator != null)
                {
                    await _enumerator.DisposeAsync().ConfigureAwait(false);
                    _enumerator = null;
                    _lookup = null;
                }
                await base.DisposeAsync().ConfigureAwait(false);
            }
            protected override async ValueTask MoveNextCore()
            {
                switch (_state)
                {
                    case AsyncIteratorState.Allocated:
                        _lookup = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, _cancellationToken).ConfigureAwait(false);
                        _enumerator = _lookup.SelectAwaitCore(async g => await _resultSelector(g.Key, g).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
                        _state = AsyncIteratorState.Iterating;
                        goto case AsyncIteratorState.Iterating;
                    case AsyncIteratorState.Iterating:
                        if (await _enumerator!.MoveNextAsync().ConfigureAwait(false))
                        {
                            _current = _enumerator.Current;
                            return true;
                        }
                        await DisposeAsync().ConfigureAwait(false);
                        break;
                }
                return false;
            }
            public async ValueTask ToArrayAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return await l.ToArray(_resultSelector).ConfigureAwait(false);
            }
            public async ValueTask> ToListAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return await l.ToList(_resultSelector).ConfigureAwait(false);
            }
            public ValueTask GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
            {
                if (onlyIfCheap)
                {
                    return new ValueTask(-1);
                }
                return Core();
                async ValueTask Core()
                {
                    var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                    return l.Count;
                }
            }
        }
#if !NO_DEEP_CANCELLATION
        private sealed class GroupedResultAsyncEnumerableWithTaskAndCancellation : AsyncIterator, IAsyncIListProvider
        {
            private readonly IAsyncEnumerable _source;
            private readonly Func> _keySelector;
            private readonly Func, CancellationToken, ValueTask> _resultSelector;
            private readonly IEqualityComparer? _comparer;
            private Internal.LookupWithTask? _lookup;
            private IAsyncEnumerator? _enumerator;
            public GroupedResultAsyncEnumerableWithTaskAndCancellation(IAsyncEnumerable source, Func> keySelector, Func, CancellationToken, ValueTask> resultSelector, IEqualityComparer? comparer)
            {
                _source = source ?? throw Error.ArgumentNull(nameof(source));
                _keySelector = keySelector ?? throw Error.ArgumentNull(nameof(keySelector));
                _resultSelector = resultSelector ?? throw Error.ArgumentNull(nameof(resultSelector));
                _comparer = comparer;
            }
            public override AsyncIteratorBase Clone()
            {
                return new GroupedResultAsyncEnumerableWithTaskAndCancellation(_source, _keySelector, _resultSelector, _comparer);
            }
            public override async ValueTask DisposeAsync()
            {
                if (_enumerator != null)
                {
                    await _enumerator.DisposeAsync().ConfigureAwait(false);
                    _enumerator = null;
                    _lookup = null;
                }
                await base.DisposeAsync().ConfigureAwait(false);
            }
            protected override async ValueTask MoveNextCore()
            {
                switch (_state)
                {
                    case AsyncIteratorState.Allocated:
                        _lookup = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, _cancellationToken).ConfigureAwait(false);
                        _enumerator = _lookup.SelectAwaitCore(async g => await _resultSelector(g.Key, g, _cancellationToken).ConfigureAwait(false)).GetAsyncEnumerator(_cancellationToken); // REVIEW: Introduce another ApplyResultSelector?
                        _state = AsyncIteratorState.Iterating;
                        goto case AsyncIteratorState.Iterating;
                    case AsyncIteratorState.Iterating:
                        if (await _enumerator!.MoveNextAsync().ConfigureAwait(false))
                        {
                            _current = _enumerator.Current;
                            return true;
                        }
                        await DisposeAsync().ConfigureAwait(false);
                        break;
                }
                return false;
            }
            public async ValueTask ToArrayAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return await l.ToArray(_resultSelector, cancellationToken).ConfigureAwait(false);
            }
            public async ValueTask> ToListAsync(CancellationToken cancellationToken)
            {
                var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                return await l.ToList(_resultSelector, cancellationToken).ConfigureAwait(false);
            }
            public ValueTask GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
            {
                if (onlyIfCheap)
                {
                    return new ValueTask(-1);
                }
                return Core();
                async ValueTask Core()
                {
                    var l = await Internal.LookupWithTask.CreateAsync(_source, _keySelector, _comparer, cancellationToken).ConfigureAwait(false);
                    return l.Count;
                }
            }
        }
#endif
        private sealed class GroupedResultAsyncEnumerable : AsyncIterator, IAsyncIListProvider
        {
            private readonly IAsyncEnumerable _source;
            private readonly Func _keySelector;
            private readonly Func _elementSelector;
            private readonly Func, TResult> _resultSelector;
            private readonly IEqualityComparer? _comparer;
            private Internal.Lookup? _lookup;
            private IEnumerator? _enumerator;
            public GroupedResultAsyncEnumerable(IAsyncEnumerable source, Func