// 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;
using System.Composition;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Host.Mef;
using Microsoft.VisualStudio.Threading;
 
namespace Microsoft.CodeAnalysis.Editor.Shared.Utilities
{
    /// <summary>
    /// Implements <see cref="IThreadingContext"/>, which provides an implementation of
    /// <see cref="VisualStudio.Threading.JoinableTaskFactory"/> to Roslyn code.
    /// </summary>
    /// <remarks>
    /// <para>The <see cref="VisualStudio.Threading.JoinableTaskFactory"/> is constructed from the
    /// <see cref="VisualStudio.Threading.JoinableTaskContext"/> provided by the MEF container, if available. If no
    /// <see cref="VisualStudio.Threading.JoinableTaskContext"/> is available, a new instance is constructed using the
    /// synchronization context of the current thread as the main thread.</para>
    /// </remarks>
    [Export(typeof(IThreadingContext))]
    [Shared]
    internal sealed class ThreadingContext : IThreadingContext, IDisposable
    {
        private readonly CancellationTokenSource _disposalTokenSource = new();
 
        [ImportingConstructor]
        [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
        public ThreadingContext(JoinableTaskContext joinableTaskContext)
        {
            HasMainThread = joinableTaskContext.MainThread.IsAlive;
            JoinableTaskContext = joinableTaskContext;
            JoinableTaskFactory = joinableTaskContext.Factory;
            ShutdownBlockingTasks = new JoinableTaskCollection(JoinableTaskContext);
            ShutdownBlockingTaskFactory = JoinableTaskContext.CreateFactory(ShutdownBlockingTasks);
        }
 
        /// <inheritdoc/>
        public bool HasMainThread
        {
            get;
        }
 
        /// <inheritdoc/>
        public JoinableTaskContext JoinableTaskContext
        {
            get;
        }
 
        /// <inheritdoc/>
        public JoinableTaskFactory JoinableTaskFactory
        {
            get;
        }
 
        public JoinableTaskCollection ShutdownBlockingTasks { get; }
 
        private JoinableTaskFactory ShutdownBlockingTaskFactory { get; }
 
        public CancellationToken DisposalToken => _disposalTokenSource.Token;
 
        public JoinableTask RunWithShutdownBlockAsync(Func<CancellationToken, Task> func)
        {
            return ShutdownBlockingTaskFactory.RunAsync(() =>
            {
                DisposalToken.ThrowIfCancellationRequested();
                return func(DisposalToken);
            });
        }
 
        public void Dispose()
        {
            // https://github.com/Microsoft/vs-threading/blob/main/doc/cookbook_vs.md#how-to-write-a-fire-and-forget-method-responsibly
            _disposalTokenSource.Cancel();
 
            try
            {
                // Block Dispose until all async work has completed.
                JoinableTaskContext.Factory.Run(ShutdownBlockingTasks.JoinTillEmptyAsync);
            }
            catch (OperationCanceledException)
            {
                // this exception is expected because we signaled the cancellation token
            }
            catch (AggregateException ex)
            {
                // ignore AggregateException containing only OperationCanceledException
                ex.Handle(inner => inner is OperationCanceledException);
            }
        }
    }
}