using Beam.Models;

namespace Beam {
    public static class IAsyncEnumeratorExtensions {
        public static async IAsyncEnumerator<T> UnwrapFragmented<T>(this IAsyncEnumerator<Fragment<T>> fragmented) {
            if (fragmented is null)
                throw new ArgumentNullException();
            try {
                while(await fragmented.MoveNextAsync().ConfigureAwait(false)) {
                    if (fragmented.Current is null)
                        yield break;
                    if (!fragmented.Current.IsComplete)
                        yield break;
                    while (fragmented.Current.TryTake(out var item))
                        if (item is null)
                            yield break;
                        else
                            yield return item;
                }
            } finally {
                await fragmented.DisposeAsync().ConfigureAwait(false);
            }
        }

        public static async IAsyncEnumerator<Ordered<T>> WrapOrdered<T>(this IAsyncEnumerator<T> enumerator) {
            if (enumerator is null)
                throw new ArgumentNullException();
            try {
                int index = 0;
                while (await enumerator.MoveNextAsync().ConfigureAwait(false))
                    if (enumerator.Current is null)
                        yield break;
                    else
                        yield return new Ordered<T>(enumerator.Current, index);
            } finally {
                await enumerator.DisposeAsync().ConfigureAwait(false);
            }
        }
    }
}
