using ZB.MOM.WW.MxGateway.Contracts.Proto;

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

/// <summary>Coordinates a fake worker's incoming frames and event emission gate.</summary>
internal sealed class GatedWorkerEnvelopeReader(
    Func<CancellationToken, Task<WorkerEnvelope>> read,
    SemaphoreSlim gate,
    CancellationToken cancellationToken) : IAsyncDisposable
{
    private readonly CancellationTokenSource _lifetime = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
    private Task<WorkerEnvelope>? _pendingRead;
    private Task? _pendingGate;

    /// <summary>Reads an envelope, or returns null when emission should be checked.</summary>
    /// <param name="watchGate">Whether the worker is currently subscribed and emitting.</param>
    /// <returns>The next complete envelope, or null for an emission signal.</returns>
    public async Task<WorkerEnvelope?> ReadOrGateAsync(bool watchGate = true)
    {
        // A gate wakeup must never cancel or replace a partially consumed frame.
        _pendingRead ??= read(_lifetime.Token);
        if (watchGate)
        {
            _pendingGate ??= gate.WaitAsync(_lifetime.Token);
            await Task.WhenAny(_pendingRead, _pendingGate).ConfigureAwait(false);
            if (!_pendingRead.IsCompleted)
            {
                await _pendingGate.ConfigureAwait(false);
                _pendingGate = null;
                return null;
            }
        }

        WorkerEnvelope envelope = await _pendingRead.ConfigureAwait(false);
        _pendingRead = null;
        return envelope;
    }

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        await _lifetime.CancelAsync().ConfigureAwait(false);
        try
        {
            if (_pendingRead is not null)
            {
                await _pendingRead.ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (_lifetime.IsCancellationRequested)
        {
        }
        finally
        {
            try
            {
                if (_pendingGate is not null)
                {
                    await _pendingGate.ConfigureAwait(false);
                }
            }
            catch (OperationCanceledException) when (_lifetime.IsCancellationRequested)
            {
            }

            _lifetime.Dispose();
        }
    }
}