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

namespace NATS.Server.Benchmark.Tests.CorePubSub;

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

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task PubNoSub_16B()
    {
        const int payloadSize = 16;
        var payload = new byte[payloadSize];
        const string subject = "bench.pub.nosub.16";

        var dotnetResult = await RunThroughput("Single Publisher (16B)", "DotNet", subject, payload, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunThroughput("Single Publisher (16B)", "Go", subject, payload, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task PubNoSub_128B()
    {
        const int payloadSize = 128;
        var payload = new byte[payloadSize];
        const string subject = "bench.pub.nosub.128";

        var dotnetResult = await RunThroughput("Single Publisher (128B)", "DotNet", subject, payload, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunThroughput("Single Publisher (128B)", "Go", subject, payload, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private async Task<BenchmarkResult> RunThroughput(string name, string serverType, string subject, byte[] payload, Func<NatsConnection> createClient)
    {
        await using var client = createClient();
        await client.ConnectAsync();

        return await _runner.MeasureThroughputAsync(name, serverType, payload.Length,
            async _ => await client.PublishAsync(subject, payload));
    }
}