// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Threading;
using Xunit;
 
namespace Roslyn.Test.Utilities
{
    public static class TaskJoinExtensions
    {
        /// <summary>
        /// Joins a <see cref="Task"/> to the current thread with a <see cref="Dispatcher"/> message pump in place
        /// during the join operation.
        /// </summary>
        public static void JoinUsingDispatcher(this Task task, CancellationToken cancellationToken)
        {
            JoinUsingDispatcherNoResult(task, cancellationToken);
 
            // Handle task completion by throwing the appropriate exception on failure
            task.GetAwaiter().GetResult();
        }
 
        /// <summary>
        /// Joins a <see cref="Task{TResult}"/> to the current thread with a <see cref="Dispatcher"/> message pump in
        /// place during the join operation.
        /// </summary>
        public static TResult JoinUsingDispatcher<TResult>(this Task<TResult> task, CancellationToken cancellationToken)
        {
            JoinUsingDispatcherNoResult(task, cancellationToken);
 
            // Handle task completion by throwing the appropriate exception on failure
            return task.GetAwaiter().GetResult();
        }
 
        private static void JoinUsingDispatcherNoResult(Task task, CancellationToken cancellationToken)
        {
            var frame = new DispatcherFrame();
 
            // When the task completes or cancellation is requested, mark the frame so we leave the message pump
            task.ContinueWith(
                t => frame.Continue = false,
                CancellationToken.None,
                TaskContinuationOptions.ExecuteSynchronously,
                TaskScheduler.Default);
 
            using (var registration = cancellationToken.Register(() => frame.Continue = false))
            {
                Dispatcher.PushFrame(frame);
            }
 
            // Handle cancellation by throwing an exception
            if (!task.IsCompleted)
            {
                Assert.True(cancellationToken.IsCancellationRequested);
                cancellationToken.ThrowIfCancellationRequested();
            }
        }
    }
}