using System;
using System.Collections;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;

using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Fallback;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Services.Agent;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class ProviderRetryTests
    {
        private string _workspace;


        [SetUp]
        public void Setup()
        {
            _workspace =
                Path.Combine(
                    Application.temporaryCachePath,
                    "AIRuntimeProviderRetryTests");


            if (Directory.Exists(_workspace))
            {
                Directory.Delete(
                    _workspace,
                    true);
            }


            Directory.CreateDirectory(
                _workspace);
        }


        [TearDown]
        public void Cleanup()
        {
            if (Directory.Exists(_workspace))
            {
                Directory.Delete(
                    _workspace,
                    true);
            }
        }


        //--------------------------------------------------
        // Retry Policy Defaults
        //--------------------------------------------------

        [Test]
        public void RetryPolicy_HasExpectedDefaults()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.AreEqual(
                1,
                policy.MaxRetries);


            Assert.AreEqual(
                TimeSpan.FromMilliseconds(500),
                policy.InitialDelay);


            Assert.AreEqual(
                2.0,
                policy.BackoffMultiplier);


            Assert.AreEqual(
                TimeSpan.FromSeconds(5),
                policy.MaximumDelay);


            Assert.IsTrue(
                policy.ShouldRetry(
                    ProviderFailureKind.Timeout));


            Assert.IsTrue(
                policy.ShouldRetry(
                    ProviderFailureKind.Unavailable));


            Assert.IsTrue(
                policy.ShouldRetry(
                    ProviderFailureKind.RateLimited));


            Assert.IsTrue(
                policy.ShouldRetry(
                    ProviderFailureKind.EmptyResponse));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.InvalidResponse));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.Authentication));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.ModelUnavailable));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.InvalidRequest));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.Unknown));


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.Cancelled));
        }


        //--------------------------------------------------
        // Policy Validation
        //--------------------------------------------------

        [Test]
        public void RetryPolicy_RejectsNegativeRetryCount()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.Throws<ArgumentOutOfRangeException>(
                () =>
                {
                    policy.MaxRetries =
                        -1;
                });
        }


        [Test]
        public void RetryPolicy_RejectsNegativeInitialDelay()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.Throws<ArgumentOutOfRangeException>(
                () =>
                {
                    policy.InitialDelay =
                        TimeSpan.FromMilliseconds(
                            -1);
                });
        }


        [Test]
        public void RetryPolicy_RejectsBackoffBelowOne()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.Throws<ArgumentOutOfRangeException>(
                () =>
                {
                    policy.BackoffMultiplier =
                        0.5;
                });
        }


        [Test]
        public void RetryPolicy_RejectsNegativeMaximumDelay()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.Throws<ArgumentOutOfRangeException>(
                () =>
                {
                    policy.MaximumDelay =
                        TimeSpan.FromMilliseconds(
                            -1);
                });
        }


        //--------------------------------------------------
        // Backoff Calculation
        //--------------------------------------------------

        [Test]
        public void RetryPolicy_CalculatesExponentialBackoff()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy
                {
                    InitialDelay =
                        TimeSpan.FromSeconds(1),

                    BackoffMultiplier =
                        2.0,

                    MaximumDelay =
                        TimeSpan.FromSeconds(10)
                };


            Assert.AreEqual(
                TimeSpan.FromSeconds(1),
                policy.GetDelay(1));


            Assert.AreEqual(
                TimeSpan.FromSeconds(2),
                policy.GetDelay(2));


            Assert.AreEqual(
                TimeSpan.FromSeconds(4),
                policy.GetDelay(3));
        }


        [Test]
        public void RetryPolicy_CapsDelayAtMaximum()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy
                {
                    InitialDelay =
                        TimeSpan.FromSeconds(2),

                    BackoffMultiplier =
                        3.0,

                    MaximumDelay =
                        TimeSpan.FromSeconds(5)
                };


            Assert.AreEqual(
                TimeSpan.FromSeconds(2),
                policy.GetDelay(1));


            Assert.AreEqual(
                TimeSpan.FromSeconds(5),
                policy.GetDelay(2));


            Assert.AreEqual(
                TimeSpan.FromSeconds(5),
                policy.GetDelay(3));
        }


        [Test]
        public void RetryPolicy_AllowsCustomRetryableFailures()
        {
            ProviderRetryPolicy policy =
                new ProviderRetryPolicy();


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.Authentication));


            policy.SetRetryable(
                ProviderFailureKind.Authentication,
                true);


            Assert.IsTrue(
                policy.ShouldRetry(
                    ProviderFailureKind.Authentication));


            policy.SetRetryable(
                ProviderFailureKind.Authentication,
                false);


            Assert.IsFalse(
                policy.ShouldRetry(
                    ProviderFailureKind.Authentication));
        }


        //--------------------------------------------------
        // Runtime Configuration
        //--------------------------------------------------

        [Test]
        public void Runtime_ExposesRetryConfiguration()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntimeWithoutProvider();


            runtime.ProviderRetryCount =
                3;


            runtime.ProviderRetryDelay =
                TimeSpan.FromSeconds(1);


            runtime.ProviderRetryBackoffMultiplier =
                3.0;


            runtime.ProviderMaximumRetryDelay =
                TimeSpan.FromSeconds(8);


            Assert.AreEqual(
                3,
                runtime.ProviderRetryCount);


            Assert.AreEqual(
                TimeSpan.FromSeconds(1),
                runtime.ProviderRetryDelay);


            Assert.AreEqual(
                3.0,
                runtime.ProviderRetryBackoffMultiplier);


            Assert.AreEqual(
                TimeSpan.FromSeconds(8),
                runtime.ProviderMaximumRetryDelay);
        }


        [Test]
        public void Runtime_CanConfigureRetryableFailureKinds()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntimeWithoutProvider();


            Assert.IsFalse(
                runtime.IsProviderFailureRetryable(
                    ProviderFailureKind.Authentication));


            runtime.SetProviderFailureRetryable(
                ProviderFailureKind.Authentication,
                true);


            Assert.IsTrue(
                runtime.IsProviderFailureRetryable(
                    ProviderFailureKind.Authentication));
        }


        [Test]
        public void Runtime_ResetRetryPolicy_RestoresDefaults()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntimeWithoutProvider();


            runtime.ProviderRetryCount =
                5;


            runtime.ProviderRetryDelay =
                TimeSpan.Zero;


            runtime.SetProviderFailureRetryable(
                ProviderFailureKind.Authentication,
                true);


            runtime.ResetProviderRetryPolicy();


            Assert.AreEqual(
                1,
                runtime.ProviderRetryCount);


            Assert.AreEqual(
                TimeSpan.FromMilliseconds(500),
                runtime.ProviderRetryDelay);


            Assert.IsFalse(
                runtime.IsProviderFailureRetryable(
                    ProviderFailureKind.Authentication));


            Assert.IsTrue(
                runtime.IsProviderFailureRetryable(
                    ProviderFailureKind.Timeout));
        }


        //--------------------------------------------------
        // Retry Success
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator TransientFailure_RetriesSameProvider()
        {
            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")),
                    AttemptResult.Success(
                        "Recovered on retry."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                1;


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Retry the primary provider.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Transient retry test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            Assert.AreEqual(
                "Recovered on retry.",
                task.Result.Response);


            Assert.AreEqual(
                2,
                primary.CallCount);
        }


        [UnityTest]
        public IEnumerator ProviderUsesAllRetriesBeforeFallback()
        {
            List<string> attempts =
                new List<string>();


            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    attempts,
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")),
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")),
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")));


            SequenceProvider fallback =
                new SequenceProvider(
                    "Fallback",
                    attempts,
                    AttemptResult.Success(
                        "Fallback response."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                2;


            runtime.AddFallbackProvider(
                fallback);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Exhaust primary retries.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Retry-before-fallback test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            Assert.AreEqual(
                "Fallback response.",
                task.Result.Response);


            Assert.AreEqual(
                3,
                primary.CallCount);


            Assert.AreEqual(
                1,
                fallback.CallCount);


            CollectionAssert.AreEqual(
                new[]
                {
                    "Primary",
                    "Primary",
                    "Primary",
                    "Fallback"
                },
                attempts);
        }


        [UnityTest]
        public IEnumerator EveryFallbackProviderReceivesItsOwnRetries()
        {
            List<string> attempts =
                new List<string>();


            SequenceProvider primary =
                CreateUnavailableProvider(
                    "Primary",
                    attempts,
                    2);


            SequenceProvider fallback1 =
                CreateUnavailableProvider(
                    "Fallback 1",
                    attempts,
                    2);


            SequenceProvider fallback2 =
                new SequenceProvider(
                    "Fallback 2",
                    attempts,
                    AttemptResult.Success(
                        "Final provider succeeded."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                1;


            runtime.AddFallbackProvider(
                fallback1);


            runtime.AddFallbackProvider(
                fallback2);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Retry each provider.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Per-provider retry test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            CollectionAssert.AreEqual(
                new[]
                {
                    "Primary",
                    "Primary",
                    "Fallback 1",
                    "Fallback 1",
                    "Fallback 2"
                },
                attempts);
        }


        //--------------------------------------------------
        // Retry Disabled
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ZeroRetries_MovesDirectlyToFallback()
        {
            List<string> attempts =
                new List<string>();


            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    attempts,
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")));


            SequenceProvider fallback =
                new SequenceProvider(
                    "Fallback",
                    attempts,
                    AttemptResult.Success(
                        "Fallback succeeded."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            runtime.ProviderRetryCount =
                0;


            runtime.AddFallbackProvider(
                fallback);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Do not retry.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Zero-retry test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            Assert.AreEqual(
                1,
                primary.CallCount);


            Assert.AreEqual(
                1,
                fallback.CallCount);


            CollectionAssert.AreEqual(
                new[]
                {
                    "Primary",
                    "Fallback"
                },
                attempts);
        }


        //--------------------------------------------------
        // Non-Retryable Failure
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ModelUnavailable_DoesNotRetrySameProvider()
        {
            List<string> attempts =
                new List<string>();


            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    attempts,
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Model not found.")));


            SequenceProvider fallback =
                new SequenceProvider(
                    "Fallback",
                    attempts,
                    AttemptResult.Success(
                        "Fallback model succeeded."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                5;


            runtime.AddFallbackProvider(
                fallback);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Use fallback model.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Model-unavailable retry test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            Assert.AreEqual(
                1,
                primary.CallCount);


            Assert.AreEqual(
                1,
                fallback.CallCount);
        }


        [UnityTest]
        public IEnumerator InvalidRequest_StopsRetriesAndFallback()
        {
            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "400 Bad Request: invalid request.")));


            SequenceProvider fallback =
                new SequenceProvider(
                    "Fallback",
                    AttemptResult.Success(
                        "Must not execute."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                5;


            runtime.AddFallbackProvider(
                fallback);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Invalid provider request.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Invalid-request retry test");


            Assert.AreEqual(
                AgentRunStatus.Failed,
                task.Result.Status);


            Assert.AreEqual(
                1,
                primary.CallCount);


            Assert.AreEqual(
                0,
                fallback.CallCount);
        }


        //--------------------------------------------------
        // Configurable Retry Category
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator AuthenticationCanBeConfiguredAsRetryable()
        {
            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "401 Unauthorized: invalid API key.")),
                    AttemptResult.Success(
                        "Authentication recovered."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            DisableRetryDelay(
                runtime);


            runtime.ProviderRetryCount =
                1;


            runtime.SetProviderFailureRetryable(
                ProviderFailureKind.Authentication,
                true);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Retry authentication.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Custom retry-category test");


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                task.Result.Status);


            Assert.AreEqual(
                2,
                primary.CallCount);
        }


        //--------------------------------------------------
        // Cancellation During Delay
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator CancellationDuringRetryDelay_StopsImmediately()
        {
            SequenceProvider primary =
                new SequenceProvider(
                    "Primary",
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable.")),
                    AttemptResult.Success(
                        "Must not complete."));


            SequenceProvider fallback =
                new SequenceProvider(
                    "Fallback",
                    AttemptResult.Success(
                        "Must not execute."));


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    primary);


            runtime.ProviderRetryCount =
                1;


            runtime.ProviderRetryDelay =
                TimeSpan.FromSeconds(
                    10);


            runtime.ProviderMaximumRetryDelay =
                TimeSpan.FromSeconds(
                    10);


            runtime.AddFallbackProvider(
                fallback);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Cancel during retry delay.",
                    cancellation.Token);


            yield return WaitForCallCount(
                primary,
                1);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Retry-delay cancellation test");


            Assert.AreEqual(
                AgentRunStatus.Cancelled,
                task.Result.Status);


            Assert.AreEqual(
                1,
                primary.CallCount);


            Assert.AreEqual(
                0,
                fallback.CallCount);
        }


        //--------------------------------------------------
        // Configuration Lock
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator RetryConfiguration_DuringRun_IsRejected()
        {
            BlockingProvider provider =
                new BlockingProvider(
                    "Blocking Provider");


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    provider);


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    10);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Keep running.",
                    cancellation.Token);


            yield return WaitForProviderStart(
                provider);


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.ProviderRetryCount =
                            3;
                    });


            Assert.AreEqual(
                "Provider configuration cannot be changed while an agent run is in progress.",
                exception.Message);


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Retry configuration lock test");


            Assert.AreEqual(
                AgentRunStatus.Cancelled,
                task.Result.Status);
        }


        //--------------------------------------------------
        // Runtime Creation
        //--------------------------------------------------

        private AIRuntime.Runtime.AIRuntime CreateRuntimeWithoutProvider()
        {
            return new AIRuntime.Runtime.AIRuntime(
                _workspace);
        }


        private AIRuntime.Runtime.AIRuntime CreateRuntime(
            IProvider provider)
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntimeWithoutProvider();


            runtime.ProviderTimeout =
                TimeSpan.FromSeconds(
                    2);


            runtime.RunTimeout =
                TimeSpan.FromSeconds(
                    10);


            runtime.UseProvider(
                provider);


            return runtime;
        }


        private static void DisableRetryDelay(
            AIRuntime.Runtime.AIRuntime runtime)
        {
            runtime.ProviderRetryDelay =
                TimeSpan.Zero;


            runtime.ProviderMaximumRetryDelay =
                TimeSpan.Zero;
        }


        private static SequenceProvider CreateUnavailableProvider(
            string name,
            List<string> attempts,
            int failureCount)
        {
            AttemptResult[] results =
                new AttemptResult[
                    failureCount];


            for (
                int index = 0;
                index < failureCount;
                index++)
            {
                results[index] =
                    AttemptResult.Failure(
                        new InvalidOperationException(
                            "Service unavailable."));
            }


            return new SequenceProvider(
                name,
                attempts,
                results);
        }


        //--------------------------------------------------
        // Fake Providers
        //--------------------------------------------------

        private sealed class AttemptResult
        {
            private AttemptResult(
                string response,
                Exception exception)
            {
                Response =
                    response;


                Exception =
                    exception;
            }


            public string Response
            {
                get;
            }


            public Exception Exception
            {
                get;
            }


            public static AttemptResult Success(
                string response)
            {
                return new AttemptResult(
                    response,
                    exception: null);
            }


            public static AttemptResult Failure(
                Exception exception)
            {
                return new AttemptResult(
                    response: null,
                    exception);
            }
        }


        private sealed class SequenceProvider :
            IProvider
        {
            private readonly Queue<AttemptResult> _results;

            private readonly List<string> _attempts;


            public SequenceProvider(
                string name,
                params AttemptResult[] results)
                :
                this(
                    name,
                    attempts: null,
                    results)
            {
            }


            public SequenceProvider(
                string name,
                List<string> attempts,
                params AttemptResult[] results)
            {
                if (
                    results == null ||
                    results.Length == 0)
                {
                    throw new ArgumentException(
                        "At least one attempt result is required.",
                        nameof(results));
                }


                Name =
                    name;


                _attempts =
                    attempts;


                _results =
                    new Queue<AttemptResult>(
                        results);
            }


            public string Name
            {
                get;
            }


            public int CallCount
            {
                get;
                private set;
            }


            public Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken = default)
            {
                cancellationToken
                    .ThrowIfCancellationRequested();


                CallCount++;


                _attempts?.Add(
                    Name);


                if (_results.Count == 0)
                {
                    return Task.FromException<ChatCompletion>(
                        new InvalidOperationException(
                            $"Provider '{Name}' received more calls than configured."));
                }


                AttemptResult result =
                    _results.Dequeue();


                if (result.Exception != null)
                {
                    return Task.FromException<ChatCompletion>(
                        result.Exception);
                }


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            result.Response
                    });
            }
        }


        private sealed class BlockingProvider :
            IProvider
        {
            public BlockingProvider(
                string name)
            {
                Name =
                    name;
            }


            public string Name
            {
                get;
            }


            public bool HasStarted
            {
                get;
                private set;
            }


            public async Task<ChatCompletion> ChatAsync(
                IReadOnlyList<ConversationMessage> messages,
                IReadOnlyList<ProviderTool> tools,
                CancellationToken cancellationToken = default)
            {
                HasStarted =
                    true;


                await Task.Delay(
                    Timeout.Infinite,
                    cancellationToken);


                throw new InvalidOperationException(
                    "Blocking provider completed unexpectedly.");
            }
        }


        //--------------------------------------------------
        // Coroutine Helpers
        //--------------------------------------------------

        private static IEnumerator WaitForTask(
            Task task)
        {
            while (!task.IsCompleted)
            {
                yield return null;
            }
        }


        private static IEnumerator WaitForCallCount(
            SequenceProvider provider,
            int expectedCallCount)
        {
            while (
                provider.CallCount <
                expectedCallCount)
            {
                yield return null;
            }
        }


        private static IEnumerator WaitForProviderStart(
            BlockingProvider provider)
        {
            while (!provider.HasStarted)
            {
                yield return null;
            }
        }


        //--------------------------------------------------
        // Assertions
        //--------------------------------------------------

        private static void AssertTaskCompletedSuccessfully(
            Task task,
            string operation)
        {
            Assert.IsNotNull(
                task,
                operation +
                " returned a null task.");


            if (task.IsCanceled)
            {
                Assert.Fail(
                    operation +
                    " was unexpectedly cancelled.");
            }


            if (!task.IsFaulted)
                return;


            Exception exception =
                UnwrapException(
                    task.Exception);


            Assert.Fail(
                operation +
                " failed:\n" +
                (exception?.ToString() ??
                 "Unknown task failure."));
        }


        private static Exception UnwrapException(
            Exception exception)
        {
            if (exception == null)
                return null;


            if (!(exception is AggregateException aggregate))
                return exception;


            AggregateException flattened =
                aggregate.Flatten();


            if (flattened.InnerExceptions.Count == 1)
            {
                return flattened.InnerExceptions[0];
            }


            return flattened;
        }
    }
}
