using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.JetStream.Consumers;

public sealed class PushConsumerEngine
{
    public void Enqueue(ConsumerHandle consumer, StoredMessage message)
    {
        if (message.Sequence <= consumer.AckProcessor.AckFloor)
            return;

        var availableAtUtc = DateTime.UtcNow;
        if (consumer.Config.RateLimitBps > 0)
        {
            if (consumer.NextPushDataAvailableAtUtc > availableAtUtc)
                availableAtUtc = consumer.NextPushDataAvailableAtUtc;

            var delayMs = (long)Math.Ceiling((double)message.Payload.Length * 1000 / consumer.Config.RateLimitBps);
            consumer.NextPushDataAvailableAtUtc = availableAtUtc.AddMilliseconds(Math.Max(delayMs, 1));
        }

        consumer.PushFrames.Enqueue(new PushFrame
        {
            IsData = true,
            Message = message,
            AvailableAtUtc = availableAtUtc,
        });

        if (consumer.Config.AckPolicy is AckPolicy.Explicit or AckPolicy.All)
            consumer.AckProcessor.Register(message.Sequence, consumer.Config.AckWaitMs);

        if (consumer.Config.FlowControl)
        {
            consumer.PushFrames.Enqueue(new PushFrame
            {
                IsFlowControl = true,
                AvailableAtUtc = availableAtUtc,
            });
        }

        if (consumer.Config.HeartbeatMs > 0)
        {
            consumer.PushFrames.Enqueue(new PushFrame
            {
                IsHeartbeat = true,
                AvailableAtUtc = availableAtUtc,
            });
        }
    }
}

public sealed class PushFrame
{
    public bool IsData { get; init; }
    public bool IsFlowControl { get; init; }
    public bool IsHeartbeat { get; init; }
    public StoredMessage? Message { get; init; }
    public DateTime AvailableAtUtc { get; init; } = DateTime.UtcNow;
}