// 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.
 
#nullable disable
 
using System;
using System.Reflection;
using System.Threading;
using System.Windows.Threading;
using Microsoft.CodeAnalysis.Test.Utilities;
 
namespace Roslyn.Test.Utilities;
 
public sealed class StaTaskScheduler : IDisposable
{
    /// <summary>Gets a <see cref="StaTaskScheduler"/> for the current <see cref="AppDomain"/>.</summary>
    /// <remarks>We use a count of 1, because the editor ends up re-using <see cref="DispatcherObject"/>
    /// instances between tests, so we need to always use the same thread for our Sta tests.</remarks>
    public static StaTaskScheduler DefaultSta { get; } = new StaTaskScheduler();
 
    /// <summary>The STA threads used by the scheduler.</summary>
    public Thread StaThread { get; }
 
    public bool IsRunningInScheduler => StaThread.ManagedThreadId == Environment.CurrentManagedThreadId;
 
    static StaTaskScheduler()
    {
        // Overwrite xunit's app domain handling to not call AppDomain.Unload
        var getDefaultDomain = typeof(AppDomain).GetMethod("GetDefaultDomain", BindingFlags.NonPublic | BindingFlags.Static);
        var defaultDomain = (AppDomain)getDefaultDomain.Invoke(null, null);
        var hook = (XunitDisposeHook)defaultDomain.CreateInstanceFromAndUnwrap(typeof(XunitDisposeHook).Assembly.CodeBase, typeof(XunitDisposeHook).FullName, ignoreCase: false, BindingFlags.CreateInstance, binder: null, args: null, culture: null, activationAttributes: null);
        hook.Execute();
 
        // We've created an STA thread, which has some extra requirements for COM Runtime
        // Callable Wrappers (RCWs). If any COM object is created on the STA thread, calls to that
        // object must be made from that thread; when the RCW is no longer being used by any
        // managed code, the RCW is put into the finalizer queue, but to actually finalize it
        // it has to marshal to the STA thread to do the work. This means that in order to safely
        // clean up any RCWs, we need to ensure that the thread is pumping past the point of
        // all RCWs being finalized
        //
        // This constraint is particularly problematic if our tests are running in an AppDomain:
        // when the AppDomain is unloaded, any threads (including our STA thread) are going to be
        // aborted. Once the thread and AppDomain is being torn down, the CLR is going to try cleaning up
        // any RCWs associated them, because if the thread is gone for good there's no way
        // it could ever clean anything further up. The code there waits for the finalizer queue
        // -- but the finalizer queue might be already trying to clean up an RCW, which is marshaling
        // to the STA thread. This could then deadlock.
        //
        // The suggested workaround from the CLR team is to do an explicit GC.Collect and
        // WaitForPendingFinalizers before we let the AppDomain shut down. The belief is subscribing
        // to DomainUnload is a reasonable place to do it.
        AppDomain.CurrentDomain.DomainUnload += (sender, e) =>
        {
            GC.Collect();
            GC.WaitForPendingFinalizers();
        };
    }
 
    /// <summary>Initializes a new instance of the <see cref="StaTaskScheduler"/> class.</summary>
    public StaTaskScheduler()
    {
        using (var threadStartedEvent = new ManualResetEventSlim(initialState: false))
        {
            DispatcherSynchronizationContext synchronizationContext = null;
            StaThread = new Thread(() =>
            {
                var oldContext = SynchronizationContext.Current;
                try
                {
                    // All WPF Tests need a DispatcherSynchronizationContext and we dont want to block pending keyboard
                    // or mouse input from the user. So use background priority which is a single level below user input.
                    synchronizationContext = new DispatcherSynchronizationContext();
 
                    // xUnit creates its own synchronization context and wraps any existing context so that messages are
                    // still pumped as necessary. So we are safe setting it here, where we are not safe setting it in test.
                    SynchronizationContext.SetSynchronizationContext(synchronizationContext);
 
                    threadStartedEvent.Set();
 
                    Dispatcher.Run();
                }
                finally
                {
                    SynchronizationContext.SetSynchronizationContext(oldContext);
                }
            });
            StaThread.Name = $"{nameof(StaTaskScheduler)} thread";
            StaThread.IsBackground = true;
            StaThread.SetApartmentState(ApartmentState.STA);
            StaThread.Start();
 
            threadStartedEvent.Wait();
            DispatcherSynchronizationContext = synchronizationContext;
        }
 
        // Work around the WeakEventTable Shutdown race conditions
        AppContext.SetSwitch("Switch.MS.Internal.DoNotInvokeInWeakEventTableShutdownListener", isEnabled: true);
    }
 
    public DispatcherSynchronizationContext DispatcherSynchronizationContext
    {
        get;
    }
 
    /// <summary>
    /// Cleans up the scheduler by indicating that no more tasks will be queued.
    /// This method blocks until all threads successfully shutdown.
    /// </summary>
    public void Dispose()
    {
        if (StaThread.IsAlive)
        {
            DispatcherSynchronizationContext.Post(_ => Dispatcher.ExitAllFrames(), null);
            StaThread.Join();
        }
    }
}