using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using MongoDB.Driver;

namespace xDataService.Mongo {
    public class XAsyncEnumerableAdapter<T> : IAsyncEnumerable<T> {
        //
        private readonly IAsyncCursorSource<T> asyncCursorSource;

        //
        public XAsyncEnumerableAdapter (IAsyncCursorSource<T> asyncCursorSource) {
            this.asyncCursorSource = asyncCursorSource;
        }

        //
        public IAsyncEnumerator<T> GetAsyncEnumerator (CancellationToken cancellationToken) =>
        new XAsyncEnumeratorAdapter<T> (asyncCursorSource, cancellationToken);
    }

    public class XAsyncEnumeratorAdapter<T> : IAsyncEnumerator<T> {
        private readonly IAsyncCursorSource<T> asyncCursorSource;
        private readonly CancellationToken cancellationToken;
        private IAsyncCursor<T> asyncCursor;
        private IEnumerator<T> batchEnumerator;

        public T Current => batchEnumerator.Current;

        public XAsyncEnumeratorAdapter (
            IAsyncCursorSource<T> asyncCursorSource,
            CancellationToken cancellationToken
        ) {
            //
            this.asyncCursorSource = asyncCursorSource;
            this.cancellationToken = cancellationToken;
        }

        public async ValueTask<bool> MoveNextAsync () {
            //
            if (asyncCursor == null) {
                asyncCursor = await asyncCursorSource
                    .ToCursorAsync (cancellationToken);
            }

            //
            if (batchEnumerator != null &&
                batchEnumerator.MoveNext ()) {
                return true;
            }

            //
            if (asyncCursor != null &&
                await asyncCursor.MoveNextAsync (cancellationToken)) {
                //
                batchEnumerator?.Dispose ();
                batchEnumerator = asyncCursor.Current.GetEnumerator ();

                //
                return batchEnumerator.MoveNext ();
            }

            //
            return false;
        }

        public async ValueTask DisposeAsync () {
            //
            await Task.CompletedTask;
            asyncCursor?.Dispose ();
            asyncCursor = null;
        }
    }
}