using Grpc.Core;

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

/// <summary>
/// Thread-safe <see cref="IServerStreamWriter{T}"/> that records every written message
/// and lets a test await the first message with a timeout.
/// </summary>
/// <typeparam name="T">The streamed message type.</typeparam>
public sealed class RecordingServerStreamWriter<T> : IServerStreamWriter<T>
{
    private readonly object _syncRoot = new();
    private readonly TaskCompletionSource<T> _firstMessage = new(TaskCreationOptions.RunContinuationsAsynchronously);
    private readonly List<T> _messages = [];

    /// <summary>Gets the messages written to this stream, in order.</summary>
    public IReadOnlyList<T> Messages
    {
        get
        {
            lock (_syncRoot)
            {
                return _messages.ToArray();
            }
        }
    }

    /// <summary>Gets or sets options for writing messages to the stream.</summary>
    public WriteOptions? WriteOptions { get; set; }

    /// <summary>Records the supplied message.</summary>
    /// <param name="message">The message to record.</param>
    /// <returns>A completed task.</returns>
    public Task WriteAsync(T message)
    {
        lock (_syncRoot)
        {
            _messages.Add(message);
        }

        _firstMessage.TrySetResult(message);
        return Task.CompletedTask;
    }

    /// <summary>Waits for the first message to be written within the specified timeout.</summary>
    /// <param name="timeout">Maximum time to wait for the first message.</param>
    /// <returns>The first message written to this stream.</returns>
    public async Task<T> WaitForFirstMessageAsync(TimeSpan timeout) =>
        await _firstMessage.Task.WaitAsync(timeout).ConfigureAwait(false);
}
