using System;
using System.Threading;
using System.Windows.Forms;
using NUnit.Framework;
using NUnit.Framework.Interfaces;
using NUnit.Framework.Internal;
using NUnit.Framework.Internal.Commands;

/// <summary>
/// Requires NUnit 3.11 or later.
/// </summary>
[AttributeUsage(AttributeTargets.Method)]
public sealed class WindowsFormsMessagePumpAttribute : PropertyAttribute, IWrapSetUpTearDown
{
    public WindowsFormsMessagePumpAttribute()
    {
        Properties.Add(PropertyNames.ApartmentState, ApartmentState.STA);
    }

    public TestCommand Wrap(TestCommand command) => new Command(command);

    private sealed class Command : DelegatingTestCommand
    {
        public Command(TestCommand innerCommand) : base(innerCommand)
        {
        }

        /* NUnit 3.12 version (change class <summary> element if this version is used)
        public override TestResult Execute(TestExecutionContext context)
        {
            SynchronizationContext.SetSynchronizationContext(new WindowsFormsSynchronizationContext());
            return innerCommand.Execute(context);
        }
        */

        // NUnit 3.11 version
        public override TestResult Execute(TestExecutionContext context)
        {
            var previousSyncContext = SynchronizationContext.Current;
            if (!(previousSyncContext is WindowsFormsSynchronizationContext winFormsSyncContext))
            {
                winFormsSyncContext = new WindowsFormsSynchronizationContext();
                SynchronizationContext.SetSynchronizationContext(winFormsSyncContext);
            }

            try
            {
                return innerCommand.Execute(context);
            }
            finally
            {
                if (previousSyncContext != winFormsSyncContext)
                    SynchronizationContext.SetSynchronizationContext(previousSyncContext);
            }
        }
    }
}