using NATS.Client.Core;
using NATS.Server.Benchmark.Tests.Harness;
using NATS.Server.Benchmark.Tests.Infrastructure;
using Xunit.Abstractions;

namespace NATS.Server.Benchmark.Tests.RequestReply;

[Collection("Benchmark-Core")]
public class SingleClientLatencyTests(CoreServerPairFixture fixture, ITestOutputHelper output)
{
    private readonly BenchmarkRunner _runner = new() { WarmupCount = 500, MeasurementCount = 10_000 };

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task RequestReply_SingleClient_128B()
    {
        const int payloadSize = 128;
        const string subject = "bench.reqrep.single";

        var dotnetResult = await RunLatency("Request-Reply Single (128B)", "DotNet", subject, payloadSize, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunLatency("Request-Reply Single (128B)", "Go", subject, payloadSize, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private async Task<BenchmarkResult> RunLatency(string name, string serverType, string subject, int payloadSize, Func<NatsConnection> createClient)
    {
        var payload = new byte[payloadSize];

        await using var serviceClient = createClient();
        await using var requestClient = createClient();
        await serviceClient.ConnectAsync();
        await requestClient.ConnectAsync();

        // Start service responder
        var sub = await serviceClient.SubscribeCoreAsync<byte[]>(subject);
        var responderTask = Task.Run(async () =>
        {
            await foreach (var msg in sub.Msgs.ReadAllAsync())
            {
                if (msg.ReplyTo is not null)
                    await serviceClient.PublishAsync(msg.ReplyTo, payload);
            }
        });

        await Task.Delay(50);

        var result = await _runner.MeasureLatencyAsync(name, serverType, payloadSize,
            async _ =>
            {
                await requestClient.RequestAsync<byte[], byte[]>(subject, payload);
            });

        await sub.UnsubscribeAsync();
        return result;
    }
}