﻿/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma warning disable 8601
#pragma warning disable 8625

namespace DotPulsar.StressTests
{
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Linq;
    using System.Threading.Tasks;

    public static class EnumerableValueTaskExtensions
    {
        [DebuggerStepThrough]
        public static async ValueTask<TResult[]> WhenAll<TResult>(this IEnumerable<ValueTask<TResult>> source) where TResult : notnull
        {
            // the volatile property IsCompleted must be accessed only once
            var tasks = source.Select(GetInfo).ToArray();

            // run incomplete tasks
            foreach (var task in tasks)
            {
                if (task.Task is not null && !task.Task.IsCompleted)
                    await task.Task.ConfigureAwait(false);
            }

            // return ordered mixed tasks \m/
            return tasks
                .Select(x => x.Task is not null ? x.Task.Result : x.Result)
                .ToArray();
        }

        [DebuggerStepThrough]
        public static async Task<TResult[]> WhenAllAsTask<TResult>(this IEnumerable<ValueTask<TResult>> source) where TResult : notnull
            => await source.WhenAll().ConfigureAwait(false);

        [DebuggerStepThrough]
        public static async IAsyncEnumerable<TResult> Enumerate<TResult>(this IEnumerable<ValueTask<TResult>> source) where TResult : notnull
        {
            foreach (var operation in source.Select(GetInfo))
            {
                yield return operation.Task is not null
                    ? await operation.Task.ConfigureAwait(false)
                    : operation.Result;
            }
        }

        private static ValueTaskInfo<TResult> GetInfo<TResult>(this ValueTask<TResult> source) where TResult : notnull
            => source.IsCompleted
                ? new ValueTaskInfo<TResult>(source.Result)
                : new ValueTaskInfo<TResult>(source.AsTask());

        private readonly struct ValueTaskInfo<TResult> where TResult : notnull
        {
            public ValueTaskInfo(Task<TResult> task)
            {
                Result = default;
                Task = task;
            }

            public ValueTaskInfo(TResult result)
            {
                Result = result;
                Task = default;
            }

            public TResult Result { get; }
            public Task<TResult>? Task { get; }
        }
    }

    public static class EnumerableTaskExtensions
    {
        [DebuggerStepThrough]
        public static Task WhenAll(this IEnumerable<Task> source)
            => Task.WhenAll(source);

        [DebuggerStepThrough]
        public static Task<T[]> WhenAll<T>(this IEnumerable<Task<T>> source)
            => Task.WhenAll(source);
    }
}