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

namespace NATS.Server.Benchmark.Tests.JetStream;

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

    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task JSSyncPublish_16B_MemoryStore()
    {
        const int payloadSize = 16;

        var dotnetResult = await RunSyncPublish("JS Sync Publish (16B Memory)", "DotNet", payloadSize, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunSyncPublish("JS Sync Publish (16B Memory)", "Go", payloadSize, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private async Task<BenchmarkResult> RunSyncPublish(string name, string serverType, int payloadSize, Func<NatsConnection> createClient)
    {
        var payload = new byte[payloadSize];
        var streamName = $"BENCH_SYNC_{serverType.ToUpperInvariant()}_{Guid.NewGuid():N}"[..30];
        var subject = $"bench.js.sync.{serverType.ToLowerInvariant()}";

        await using var nats = createClient();
        await nats.ConnectAsync();
        var js = new NatsJSContext(nats);

        await js.CreateStreamAsync(new StreamConfig(streamName, [subject])
        {
            Storage = StreamConfigStorage.Memory,
            Retention = StreamConfigRetention.Limits,
            MaxMsgs = 1_000_000,
        });

        try
        {
            var result = await _runner.MeasureThroughputAsync(name, serverType, payloadSize,
                async _ => await js.PublishAsync(subject, payload));
            return result;
        }
        finally
        {
            await js.DeleteStreamAsync(streamName);
        }
    }
}