using System.Diagnostics;
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 DurableConsumerFetchTests(JetStreamServerPairFixture fixture, ITestOutputHelper output)
{
    [Fact]
    [Trait("Category", "Benchmark")]
    public async Task JSDurableFetch_Throughput()
    {
        const int payloadSize = 128;
        const int messageCount = 5_000;
        const int fetchBatchSize = 500;

        var dotnetResult = await RunDurableFetch("JS Durable Fetch (128B)", "DotNet", payloadSize, messageCount, fetchBatchSize, fixture.CreateDotNetClient);

        if (fixture.GoAvailable)
        {
            var goResult = await RunDurableFetch("JS Durable Fetch (128B)", "Go", payloadSize, messageCount, fetchBatchSize, fixture.CreateGoClient);
            BenchmarkResultWriter.WriteComparison(output, goResult, dotnetResult);
        }
        else
        {
            BenchmarkResultWriter.WriteSingle(output, dotnetResult);
        }
    }

    private static async Task<BenchmarkResult> RunDurableFetch(
        string name, string serverType, int payloadSize, int messageCount, int fetchBatchSize,
        Func<NatsConnection> createClient)
    {
        var payload = new byte[payloadSize];
        var streamName = $"BENCH_DUR_{serverType.ToUpperInvariant()}_{Guid.NewGuid():N}"[..30];
        var subject = $"bench.js.durable.{serverType.ToLowerInvariant()}";
        var consumerName = $"bench-dur-{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 = 10_000_000,
        });

        try
        {
            // Pre-populate stream
            var pubTasks = new List<ValueTask<PubAckResponse>>(1000);
            for (var i = 0; i < messageCount; i++)
            {
                pubTasks.Add(js.PublishAsync(subject, payload));
                if (pubTasks.Count >= 1000)
                {
                    foreach (var t in pubTasks)
                        await t;
                    pubTasks.Clear();
                }
            }

            foreach (var t in pubTasks)
                await t;

            // Create durable consumer
            var consumer = await js.CreateOrUpdateConsumerAsync(streamName, new ConsumerConfig(consumerName)
            {
                AckPolicy = ConsumerConfigAckPolicy.None,
            });

            // Fetch in batches
            var received = 0;
            var sw = Stopwatch.StartNew();

            while (received < messageCount)
            {
                await foreach (var msg in consumer.FetchAsync<byte[]>(new NatsJSFetchOpts { MaxMsgs = fetchBatchSize }))
                {
                    received++;
                    if (received >= messageCount)
                        break;
                }
            }

            sw.Stop();

            return new BenchmarkResult
            {
                Name = name,
                ServerType = serverType,
                TotalMessages = received,
                TotalBytes = (long)received * payloadSize,
                Duration = sw.Elapsed,
            };
        }
        finally
        {
            await js.DeleteStreamAsync(streamName);
        }
    }
}