// Go: consumer.go:500-600 — Priority group tests for sticky consumer assignment.
// Validates that the lowest-priority-numbered consumer is "active" and that
// failover occurs correctly when consumers register/unregister.
using System.Collections.Concurrent;
using System.Text;
using NATS.Server.JetStream;
using NATS.Server.JetStream.Consumers;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.Tests.JetStream.Consumers;

public class PriorityGroupTests
{
    // -------------------------------------------------------------------------
    // Test 1 — Single consumer registered is active
    //
    // Go reference: consumer.go:500 — when only one consumer is in a priority
    // group, it is unconditionally the active consumer.
    // -------------------------------------------------------------------------
    [Fact]
    public void Register_SingleConsumer_IsActive()
    {
        var mgr = new PriorityGroupManager();
        mgr.Register("group1", "consumer-a", priority: 1);

        mgr.IsActive("group1", "consumer-a").ShouldBeTrue();
        mgr.GetActiveConsumer("group1").ShouldBe("consumer-a");
    }

    // -------------------------------------------------------------------------
    // Test 2 — Multiple consumers: lowest priority number wins
    //
    // Go reference: consumer.go:510 — the consumer with the lowest priority
    // number is the active consumer. Priority 1 < Priority 5, so 1 wins.
    // -------------------------------------------------------------------------
    [Fact]
    public void Register_MultipleConsumers_LowestPriorityIsActive()
    {
        var mgr = new PriorityGroupManager();
        mgr.Register("group1", "consumer-high", priority: 5);
        mgr.Register("group1", "consumer-low", priority: 1);
        mgr.Register("group1", "consumer-mid", priority: 3);

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-low");
        mgr.IsActive("group1", "consumer-low").ShouldBeTrue();
        mgr.IsActive("group1", "consumer-high").ShouldBeFalse();
        mgr.IsActive("group1", "consumer-mid").ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Test 3 — Unregister active consumer: next takes over
    //
    // Go reference: consumer.go:530 — when the active consumer disconnects,
    // the next-lowest-priority consumer becomes active (failover).
    // -------------------------------------------------------------------------
    [Fact]
    public void Unregister_ActiveConsumer_NextTakesOver()
    {
        var mgr = new PriorityGroupManager();
        mgr.Register("group1", "consumer-a", priority: 1);
        mgr.Register("group1", "consumer-b", priority: 2);
        mgr.Register("group1", "consumer-c", priority: 3);

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-a");

        mgr.Unregister("group1", "consumer-a");

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-b");
        mgr.IsActive("group1", "consumer-b").ShouldBeTrue();
        mgr.IsActive("group1", "consumer-a").ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Test 4 — Unregister non-active consumer: active unchanged
    //
    // Go reference: consumer.go:540 — removing a non-active consumer does not
    // change the active assignment.
    // -------------------------------------------------------------------------
    [Fact]
    public void Unregister_NonActiveConsumer_ActiveUnchanged()
    {
        var mgr = new PriorityGroupManager();
        mgr.Register("group1", "consumer-a", priority: 1);
        mgr.Register("group1", "consumer-b", priority: 2);

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-a");

        mgr.Unregister("group1", "consumer-b");

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-a");
        mgr.IsActive("group1", "consumer-a").ShouldBeTrue();
    }

    // -------------------------------------------------------------------------
    // Test 5 — Same priority: first registered wins
    //
    // Go reference: consumer.go:520 — when two consumers share the same
    // priority, the first to register is treated as the active consumer.
    // -------------------------------------------------------------------------
    [Fact]
    public void Register_SamePriority_FirstRegisteredWins()
    {
        var mgr = new PriorityGroupManager();
        mgr.Register("group1", "consumer-first", priority: 1);
        mgr.Register("group1", "consumer-second", priority: 1);

        mgr.GetActiveConsumer("group1").ShouldBe("consumer-first");
        mgr.IsActive("group1", "consumer-first").ShouldBeTrue();
        mgr.IsActive("group1", "consumer-second").ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Test 6 — Empty group returns null
    //
    // Go reference: consumer.go:550 — calling GetActiveConsumer on an empty
    // or nonexistent group returns nil (null).
    // -------------------------------------------------------------------------
    [Fact]
    public void GetActiveConsumer_EmptyGroup_ReturnsNull()
    {
        var mgr = new PriorityGroupManager();

        mgr.GetActiveConsumer("nonexistent").ShouldBeNull();
        mgr.IsActive("nonexistent", "any-consumer").ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Test 7 — Idle heartbeat sent after timeout
    //
    // Go reference: consumer.go:5222 — sendIdleHeartbeat is invoked by a
    // background timer when no data frames are delivered within HeartbeatMs.
    // -------------------------------------------------------------------------
    [Fact]
    public async Task IdleHeartbeat_SentAfterTimeout()
    {
        var engine = new PushConsumerEngine();
        var consumer = new ConsumerHandle("TEST-STREAM", new ConsumerConfig
        {
            DurableName = "HB-CONSUMER",
            Push = true,
            DeliverSubject = "deliver.hb",
            HeartbeatMs = 50, // 50ms heartbeat interval
        });

        var sent = new ConcurrentBag<(string Subject, string ReplyTo, byte[] Headers, byte[] Payload)>();

        ValueTask SendCapture(string subject, string replyTo, ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload, CancellationToken ct)
        {
            sent.Add((subject, replyTo, headers.ToArray(), payload.ToArray()));
            return ValueTask.CompletedTask;
        }

        using var cts = new CancellationTokenSource();

        engine.StartDeliveryLoop(consumer, SendCapture, cts.Token);

        // Wait long enough for at least one idle heartbeat to fire
        await Task.Delay(200);

        engine.StopDeliveryLoop();

        engine.IdleHeartbeatsSent.ShouldBeGreaterThan(0);

        // Verify the heartbeat messages were sent to the deliver subject
        var hbMessages = sent.Where(s =>
            Encoding.ASCII.GetString(s.Headers).Contains("Idle Heartbeat")).ToList();
        hbMessages.Count.ShouldBeGreaterThan(0);
        hbMessages.ShouldAllBe(m => m.Subject == "deliver.hb");
    }

    // -------------------------------------------------------------------------
    // Test 8 — Idle heartbeat resets on data delivery
    //
    // Go reference: consumer.go:5222 — the idle heartbeat timer is reset
    // whenever a data frame is delivered, so heartbeats only fire during
    // periods of inactivity.
    // -------------------------------------------------------------------------
    [Fact]
    public async Task IdleHeartbeat_ResetOnDataDelivery()
    {
        var engine = new PushConsumerEngine();
        var consumer = new ConsumerHandle("TEST-STREAM", new ConsumerConfig
        {
            DurableName = "HB-RESET",
            Push = true,
            DeliverSubject = "deliver.hbreset",
            HeartbeatMs = 100, // 100ms heartbeat interval
        });

        var dataFramesSent = new ConcurrentBag<string>();
        var heartbeatsSent = new ConcurrentBag<string>();

        ValueTask SendCapture(string subject, string replyTo, ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload, CancellationToken ct)
        {
            var headerStr = Encoding.ASCII.GetString(headers.Span);
            if (headerStr.Contains("Idle Heartbeat"))
                heartbeatsSent.Add(subject);
            else
                dataFramesSent.Add(subject);
            return ValueTask.CompletedTask;
        }

        using var cts = new CancellationTokenSource();

        engine.StartDeliveryLoop(consumer, SendCapture, cts.Token);

        // Continuously enqueue data messages faster than the heartbeat interval
        // to keep the timer resetting. Each data delivery resets the idle heartbeat.
        for (var i = 0; i < 5; i++)
        {
            engine.Enqueue(consumer, new StoredMessage
            {
                Sequence = (ulong)(i + 1),
                Subject = "test.data",
                Payload = Encoding.UTF8.GetBytes($"msg-{i}"),
                TimestampUtc = DateTime.UtcNow,
            });
            await Task.Delay(30); // 30ms between messages — well within 100ms heartbeat
        }

        // Wait a bit after last message for potential heartbeat
        await Task.Delay(50);

        engine.StopDeliveryLoop();

        // Data frames should have been sent
        dataFramesSent.Count.ShouldBeGreaterThan(0);

        // During continuous data delivery, idle heartbeats from the timer should
        // NOT have fired because the timer is reset on each data frame.
        // (The queue-based heartbeat frames still fire as part of Enqueue, but
        // the idle heartbeat timer counter should be 0 or very low since data
        // kept flowing within the heartbeat interval.)
        engine.IdleHeartbeatsSent.ShouldBe(0);
    }
}