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

using AIRuntime.Runtime.Capabilities;
using AIRuntime.Runtime.Capabilities.Permissions;
using AIRuntime.Runtime.Models;
using AIRuntime.Runtime.Providers;
using AIRuntime.Runtime.Providers.Models;
using AIRuntime.Runtime.Services.Agent;

using NUnit.Framework;

using UnityEngine;
using UnityEngine.TestTools;


namespace AIRuntime.Tests
{
    public sealed class CapabilityPermissionTests
    {
        private const string PolicyLockedMessage =
            "Capability permissions cannot be changed while an agent run is in progress.";


        private string _workspace;


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


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


            Directory.CreateDirectory(
                _workspace);
        }


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


        //--------------------------------------------------
        // Default Policy
        //--------------------------------------------------

        [Test]
        public void DefaultPolicy_AllowsReadOnlyCapabilities()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.ReadOnly));
        }


        [Test]
        public void DefaultPolicy_RequiresApprovalForWriteCapabilities()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));
        }


        [Test]
        public void DefaultPolicy_RequiresApprovalForDangerousPermissions()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Destructive));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Process));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Network));
        }


        //--------------------------------------------------
        // Read-Only Execution
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator ReadOnlyCapability_IsAllowedByDefault()
        {
            CountingReadOnlyCapability capability =
                new CountingReadOnlyCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    "read_only_test",
                    "{}");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Read-only capability execution");


            Assert.AreEqual(
                "Read-only result.",
                task.GetAwaiter().GetResult());


            Assert.AreEqual(
                1,
                capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Write Approval
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator WriteCapability_RequiresApprovalByDefault()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    "write_test",
                    "{}");


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsFaulted);


            Exception exception =
                UnwrapException(
                    task.Exception);


            Assert.IsInstanceOf<CapabilityPermissionException>(
                exception);


            CapabilityPermissionException permissionException =
                (CapabilityPermissionException)exception;


            Assert.AreEqual(
                "write_test",
                permissionException.CapabilityName);


            Assert.AreEqual(
                CapabilityPermission.Write,
                permissionException.Permission);


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                permissionException.Decision);


            StringAssert.Contains(
                "requires approval",
                permissionException.Message);


            Assert.AreEqual(
                0,
                capability.ExecutionCount,
                "A capability requiring approval must not execute.");
        }


        //--------------------------------------------------
        // Explicit Allow
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator AllowedWriteCapability_ExecutesNormally()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.AllowPermission(
                CapabilityPermission.Write);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    "write_test",
                    "{}");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Allowed write capability");


            Assert.AreEqual(
                "Write result.",
                task.GetAwaiter().GetResult());


            Assert.AreEqual(
                1,
                capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Explicit Denial
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator DeniedCapability_DoesNotExecute()
        {
            CountingReadOnlyCapability capability =
                new CountingReadOnlyCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.RegisterCapability(
                capability);


            runtime.DenyPermission(
                CapabilityPermission.ReadOnly);


            Task<string> task =
                runtime.ExecuteCapabilityAsync(
                    "read_only_test",
                    "{}");


            yield return WaitForTask(
                task);


            Assert.IsTrue(
                task.IsFaulted);


            Exception exception =
                UnwrapException(
                    task.Exception);


            Assert.IsInstanceOf<CapabilityPermissionException>(
                exception);


            CapabilityPermissionException permissionException =
                (CapabilityPermissionException)exception;


            Assert.AreEqual(
                CapabilityPermissionDecision.Deny,
                permissionException.Decision);


            Assert.AreEqual(
                CapabilityPermission.ReadOnly,
                permissionException.Permission);


            StringAssert.Contains(
                "was denied",
                permissionException.Message);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);
        }


        //--------------------------------------------------
        // Require Approval Configuration
        //--------------------------------------------------

        [Test]
        public void RequirePermissionApproval_ChangesDecision()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.AllowPermission(
                CapabilityPermission.Write);


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));


            runtime.RequirePermissionApproval(
                CapabilityPermission.Write);


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));
        }


        //--------------------------------------------------
        // Policy Reset
        //--------------------------------------------------

        [Test]
        public void ResetPermissionPolicy_RestoresDefaultDecisions()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            runtime.DenyPermission(
                CapabilityPermission.ReadOnly);


            runtime.AllowPermission(
                CapabilityPermission.Destructive);


            runtime.AllowPermission(
                CapabilityPermission.Network);


            runtime.ResetPermissionPolicy();


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.ReadOnly));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Destructive));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Process));


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Network));
        }


        //--------------------------------------------------
        // Active Run Protection
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator PermissionPolicy_CannotAllowDuringRun()
        {
            BlockingProvider provider =
                new BlockingProvider();


            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 this run active.",
                    cancellation.Token);


            yield return WaitForProviderStart(
                provider);


            Assert.IsTrue(
                runtime.IsRunning);


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.AllowPermission(
                            CapabilityPermission.Destructive);
                    });


            Assert.AreEqual(
                PolicyLockedMessage,
                exception.Message);


            Assert.AreEqual(
                CapabilityPermissionDecision.RequireApproval,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Destructive));


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Permission allow lock test");


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


        [UnityTest]
        public IEnumerator PermissionPolicy_CannotDenyDuringRun()
        {
            BlockingProvider provider =
                new BlockingProvider();


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


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


            yield return WaitForProviderStart(
                provider);


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.DenyPermission(
                            CapabilityPermission.ReadOnly);
                    });


            Assert.AreEqual(
                PolicyLockedMessage,
                exception.Message);


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.ReadOnly));


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Permission deny lock test");


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


        [UnityTest]
        public IEnumerator PermissionPolicy_CannotRequireApprovalDuringRun()
        {
            BlockingProvider provider =
                new BlockingProvider();


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


            runtime.AllowPermission(
                CapabilityPermission.Write);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


            yield return WaitForProviderStart(
                provider);


            InvalidOperationException exception =
                Assert.Throws<InvalidOperationException>(
                    () =>
                    {
                        runtime.RequirePermissionApproval(
                            CapabilityPermission.Write);
                    });


            Assert.AreEqual(
                PolicyLockedMessage,
                exception.Message);


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Permission approval lock test");


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


        [UnityTest]
        public IEnumerator PermissionPolicy_CannotResetDuringRun()
        {
            BlockingProvider provider =
                new BlockingProvider();


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


            runtime.DenyPermission(
                CapabilityPermission.ReadOnly);


            using CancellationTokenSource cancellation =
                new CancellationTokenSource();


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


            yield return WaitForProviderStart(
                provider);


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


            Assert.AreEqual(
                PolicyLockedMessage,
                exception.Message);


            Assert.AreEqual(
                CapabilityPermissionDecision.Deny,
                runtime.GetPermissionDecision(
                    CapabilityPermission.ReadOnly));


            cancellation.Cancel();


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Permission reset lock test");


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


        //--------------------------------------------------
        // Changes After Run
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator PermissionPolicy_CanChangeAfterRunCompletes()
        {
            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime();


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Complete normally.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Completed run");


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


            Assert.IsFalse(
                runtime.IsRunning);


            runtime.AllowPermission(
                CapabilityPermission.Write);


            Assert.AreEqual(
                CapabilityPermissionDecision.Allow,
                runtime.GetPermissionDecision(
                    CapabilityPermission.Write));
        }


        //--------------------------------------------------
        // Agent Permission Failure
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator PermissionFailure_ReturnsFailedAgentStatus()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            AIRuntime.Runtime.AIRuntime runtime =
                CreateRuntime(
                    new SingleToolCallProvider(
                        "write_test"));


            runtime.RegisterCapability(
                capability);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Attempt the write capability.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Permission-failed agent run");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.IsTrue(
                result.Failed);


            Assert.IsFalse(
                result.Succeeded);


            Assert.IsFalse(
                result.Cancelled);


            Assert.IsFalse(
                result.TimedOut);


            StringAssert.Contains(
                "requires approval",
                result.ErrorMessage);


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);


            Assert.AreEqual(
                0,
                runtime.MessageCount,
                "The denied run should roll back its conversation.");
        }


        //--------------------------------------------------
        // Agent Allowed Execution
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator AllowedPermission_ExecutesThroughAgentLoop()
        {
            CountingWriteCapability capability =
                new CountingWriteCapability();


            ToolThenResponseProvider provider =
                new ToolThenResponseProvider(
                    "write_test");


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


            runtime.RegisterCapability(
                capability);


            runtime.AllowPermission(
                CapabilityPermission.Write);


            Task<AgentRunResult> task =
                runtime.RunAsync(
                    "Execute the permitted write capability.");


            yield return WaitForTask(
                task);


            AssertTaskCompletedSuccessfully(
                task,
                "Allowed agent capability run");


            AgentRunResult result =
                task.GetAwaiter()
                    .GetResult();


            Assert.AreEqual(
                AgentRunStatus.Succeeded,
                result.Status);


            Assert.IsTrue(
                result.Succeeded);


            Assert.AreEqual(
                1,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                1,
                capability.ExecutionCount);


            Assert.AreEqual(
                "Capability completed.",
                result.Response);


            Assert.AreEqual(
                4,
                runtime.MessageCount,
                "The conversation should contain user, tool request, tool result and final response.");
        }


        //--------------------------------------------------
        // Isolated Permission Failure
        //--------------------------------------------------

        [UnityTest]
        public IEnumerator IsolatedPermissionFailure_PreservesPersistentHistory()
        {
            SwitchableProvider provider =
                new SwitchableProvider();


            CountingWriteCapability capability =
                new CountingWriteCapability();


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


            runtime.RegisterCapability(
                capability);


            Task<AgentRunResult> persistentTask =
                runtime.RunAsync(
                    "Store persistent history.");


            yield return WaitForTask(
                persistentTask);


            AssertTaskCompletedSuccessfully(
                persistentTask,
                "Persistent setup run");


            Assert.IsTrue(
                persistentTask.Result.Succeeded);


            Assert.AreEqual(
                2,
                runtime.MessageCount);


            provider.RequestWriteCapability =
                true;


            Task<AgentRunResult> isolatedTask =
                runtime.RunIsolatedAsync(
                    "Attempt an isolated write.");


            yield return WaitForTask(
                isolatedTask);


            AssertTaskCompletedSuccessfully(
                isolatedTask,
                "Isolated permission failure");


            AgentRunResult result =
                isolatedTask.GetAwaiter()
                    .GetResult();


            Assert.AreEqual(
                AgentRunStatus.Failed,
                result.Status);


            Assert.AreEqual(
                0,
                result.ToolCallsExecuted);


            Assert.AreEqual(
                0,
                capability.ExecutionCount);


            Assert.AreEqual(
                2,
                runtime.MessageCount,
                "The isolated permission failure must preserve persistent history.");
        }


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

        private AIRuntime.Runtime.AIRuntime CreateRuntime()
        {
            return CreateRuntime(
                new SuccessfulProvider());
        }


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


            runtime.UseProvider(
                provider);


            return runtime;
        }


        //--------------------------------------------------
        // Test Capabilities
        //--------------------------------------------------

        private sealed class CountingReadOnlyCapability :
            CapabilityBase<EmptyArguments>
        {
            public int ExecutionCount
            {
                get;
                private set;
            }


            public override string Name =>
                "read_only_test";


            public override string Description =>
                "A read-only test capability.";


            public override CapabilityPermission Permission =>
                CapabilityPermission.ReadOnly;


            protected override Task<object> ExecuteAsync(
                EmptyArguments arguments,
                CapabilityContext context)
            {
                context.CancellationToken
                    .ThrowIfCancellationRequested();


                ExecutionCount++;


                return Task.FromResult<object>(
                    "Read-only result.");
            }
        }


        private sealed class CountingWriteCapability :
            CapabilityBase<EmptyArguments>
        {
            public int ExecutionCount
            {
                get;
                private set;
            }


            public override string Name =>
                "write_test";


            public override string Description =>
                "A write test capability.";


            public override CapabilityPermission Permission =>
                CapabilityPermission.Write;


            protected override Task<object> ExecuteAsync(
                EmptyArguments arguments,
                CapabilityContext context)
            {
                context.CancellationToken
                    .ThrowIfCancellationRequested();


                ExecutionCount++;


                return Task.FromResult<object>(
                    "Write result.");
            }
        }


        private sealed class EmptyArguments
        {
        }


        //--------------------------------------------------
        // Providers
        //--------------------------------------------------

        private sealed class SuccessfulProvider :
            IProvider
        {
            public string Name =>
                "SuccessfulProvider";


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


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Successful response."
                    });
            }
        }


        private sealed class BlockingProvider :
            IProvider
        {
            public string Name =>
                "BlockingProvider";


            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(
                    "The blocking provider should not complete normally.");
            }
        }


        private sealed class SingleToolCallProvider :
            IProvider
        {
            private readonly string _capabilityName;


            public SingleToolCallProvider(
                string capabilityName)
            {
                _capabilityName =
                    capabilityName;
            }


            public string Name =>
                "SingleToolCallProvider";


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


                return Task.FromResult(
                    CreateToolCompletion(
                        _capabilityName,
                        "single_permission_call"));
            }
        }


        private sealed class ToolThenResponseProvider :
            IProvider
        {
            private readonly string _capabilityName;

            private bool _toolRequested;


            public ToolThenResponseProvider(
                string capabilityName)
            {
                _capabilityName =
                    capabilityName;
            }


            public string Name =>
                "ToolThenResponseProvider";


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


                if (!_toolRequested)
                {
                    _toolRequested =
                        true;


                    return Task.FromResult(
                        CreateToolCompletion(
                            _capabilityName,
                            "allowed_permission_call"));
                }


                return Task.FromResult(
                    new ChatCompletion
                    {
                        Content =
                            "Capability completed."
                    });
            }
        }


        private sealed class SwitchableProvider :
            IProvider
        {
            public bool RequestWriteCapability
            {
                get;
                set;
            }


            public string Name =>
                "SwitchableProvider";


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


                if (!RequestWriteCapability)
                {
                    return Task.FromResult(
                        new ChatCompletion
                        {
                            Content =
                                "Persistent response."
                        });
                }


                return Task.FromResult(
                    CreateToolCompletion(
                        "write_test",
                        "isolated_permission_call"));
            }
        }


        private static ChatCompletion CreateToolCompletion(
            string capabilityName,
            string callId)
        {
            ChatCompletion completion =
                new ChatCompletion();


            completion.ToolCalls.Add(
                new ToolCall
                {
                    Id =
                        callId,

                    Name =
                        capabilityName,

                    Arguments =
                        "{}"
                });


            return completion;
        }


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

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


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


        //--------------------------------------------------
        // Assertion Helpers
        //--------------------------------------------------

        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 is AggregateException aggregate))
                return exception;


            AggregateException flattened =
                aggregate.Flatten();


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


            return flattened;
        }
    }
}