using System;
using System.Threading;
using ZB.MOM.WW.MxGateway.Worker.Sta;

namespace ZB.MOM.WW.MxGateway.Worker.Tests.TestSupport;

/// <summary>
/// Test <see cref="IStaMessagePump"/> that counts pump calls and can be armed
/// to throw from the next pump, so tests can observe the STA loop's pump
/// interleaving and drive a post-start loop death without a real Win32
/// message queue. The idle wait blocks on the command wake event exactly
/// like the production pump's <c>MsgWaitForMultipleObjectsEx</c> does.
/// </summary>
internal sealed class ControllableMessagePump : IStaMessagePump
{
    private int pumpCount;
    private Exception? throwOnNextPump;

    /// <summary>Gets the number of <see cref="PumpPendingMessages"/> calls so far.</summary>
    public int PumpCount => Volatile.Read(ref pumpCount);

    /// <summary>Arms the pump to throw <paramref name="exception"/> from its next pump call.</summary>
    /// <param name="exception">The exception to throw once.</param>
    public void ThrowOnNextPump(Exception exception)
    {
        Volatile.Write(ref throwOnNextPump, exception);
    }

    /// <inheritdoc />
    public void WaitForWorkOrMessages(WaitHandle commandWakeEvent, TimeSpan timeout)
    {
        commandWakeEvent.WaitOne(timeout);
    }

    /// <inheritdoc />
    public int PumpPendingMessages()
    {
        Interlocked.Increment(ref pumpCount);
        Exception? exception = Interlocked.Exchange(ref throwOnNextPump, null);
        if (exception is not null)
        {
            throw exception;
        }

        return 0;
    }
}