// Go: consumer.go (sendIdleHeartbeat ~line 5222, sendFlowControl ~line 5495,
//               deliverMsg ~line 5364, dispatchToDeliver ~line 5040,
//               loopAndGatherMsgs ~line 1400)
using System.Text;
using System.Threading.Channels;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;
using NATS.Server.Subscriptions;

namespace NATS.Server.JetStream.Consumers;

// Go: consumer.go:1400 — signals sent to the gather loop to wake it early.
public enum ConsumerSignal
{
    /// <summary>Store has new message(s) available.</summary>
    NewMessage,
    /// <summary>An ack/nak/term was processed.</summary>
    AckEvent,
    /// <summary>Consumer config changed.</summary>
    ConfigChange,
    /// <summary>Stop the loop.</summary>
    Shutdown,
}

public sealed class PushConsumerEngine
{
    // Go: consumer.go — DeliverSubject routes push-mode messages (cfg.DeliverSubject)
    public string DeliverSubject { get; private set; } = string.Empty;

    private CancellationTokenSource? _cts;
    private Task? _deliveryTask;

    // Go: consumer.go:5222 — idle heartbeat timer state
    private Timer? _idleHeartbeatTimer;
    private Func<string, string, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>, CancellationToken, ValueTask>? _sendMessage;
    private CancellationToken _externalCt;

    // Go: consumer.go:5222 — reference to the consumer handle for pending count access
    private ConsumerHandle? _consumer;

    // Go: consumer.go:1400 — gather loop state
    private Channel<ConsumerSignal>? _signalChannel;
    private CancellationTokenSource? _gatherCts;
    private Task? _gatherTask;

    /// <summary>
    /// Tracks how many idle heartbeats have been sent since the last data delivery.
    /// Useful for testing that idle heartbeats fire and reset correctly.
    /// </summary>
    public int IdleHeartbeatsSent { get; private set; }

    /// <summary>
    /// Number of flow control frames sent but not yet acknowledged by the subscriber.
    /// Go reference: consumer.go:5495 flow control stall detection.
    /// </summary>
    public int FlowControlPendingCount { get; private set; }

    /// <summary>
    /// Maximum unacknowledged flow control frames before the consumer is considered stalled.
    /// Go reference: consumer.go flow control stall detection.
    /// </summary>
    public const int MaxFlowControlPending = 2;

    /// <summary>
    /// Returns true if the consumer appears stalled due to too many unacknowledged flow control messages.
    /// Go reference: consumer.go:5495 flow control stall detection.
    /// </summary>
    public bool IsFlowControlStalled => FlowControlPendingCount >= MaxFlowControlPending;

    /// <summary>
    /// Number of messages gathered and dispatched by the gather loop.
    /// Go reference: consumer.go:1400 loopAndGatherMsgs.
    /// </summary>
    public long GatheredCount { get; private set; }

    /// <summary>
    /// Decrements the flow control pending count when the subscriber acknowledges a flow control frame.
    /// Go reference: consumer.go:5495 flow control acknowledgement.
    /// </summary>
    public void AcknowledgeFlowControl()
    {
        if (FlowControlPendingCount > 0)
            FlowControlPendingCount--;
    }

    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,
            });
        }
    }

    // Go: consumer.go:1131 — dsubj is set from cfg.DeliverSubject at consumer creation.
    // StartDeliveryLoop wires the background pump that drains PushFrames and calls
    // sendMessage for each frame.  The delegate matches the wire-level send signature used
    // by NatsClient.SendMessage, mapped to an async ValueTask for testability.
    public void StartDeliveryLoop(
        ConsumerHandle consumer,
        Func<string, string, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>, CancellationToken, ValueTask> sendMessage,
        CancellationToken ct)
    {
        DeliverSubject = consumer.Config.DeliverSubject;

        _cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        var token = _cts.Token;

        _sendMessage = sendMessage;
        _externalCt = ct;
        _consumer = consumer;

        _deliveryTask = Task.Run(() => RunDeliveryLoopAsync(consumer, sendMessage, token), token);

        // Go: consumer.go:5222 — start idle heartbeat timer if configured
        if (consumer.Config.HeartbeatMs > 0)
        {
            StartIdleHeartbeatTimer(consumer.Config.HeartbeatMs);
        }
    }

    public void StopDeliveryLoop()
    {
        StopIdleHeartbeatTimer();
        _cts?.Cancel();
        _cts?.Dispose();
        _cts = null;
    }

    /// <summary>
    /// Starts the gather loop that polls the store for new messages.
    /// Go reference: consumer.go:1400 loopAndGatherMsgs.
    /// </summary>
    public void StartGatherLoop(
        ConsumerHandle consumer,
        IStreamStore store,
        Func<string, string, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>, CancellationToken, ValueTask> sendMessage,
        CancellationToken ct)
    {
        _signalChannel = Channel.CreateUnbounded<ConsumerSignal>();
        _gatherCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        _gatherTask = Task.Run(
            () => LoopAndGatherMsgsAsync(consumer, store, sendMessage, _gatherCts.Token),
            _gatherCts.Token);
    }

    /// <summary>
    /// Stops the gather loop and completes its signal channel.
    /// </summary>
    public void StopGatherLoop()
    {
        _signalChannel?.Writer.TryComplete();
        _gatherCts?.Cancel();
        _gatherCts?.Dispose();
        _gatherCts = null;
    }

    /// <summary>
    /// Signals the gather loop to wake up and re-poll the store.
    /// Go reference: consumer.go:1620 — channel send wakes the loop.
    /// </summary>
    public void Signal(ConsumerSignal signal)
    {
        _signalChannel?.Writer.TryWrite(signal);
    }

    /// <summary>
    /// Public test accessor for the filter predicate. Production code uses
    /// the private ShouldDeliver; this entry point avoids reflection in unit tests.
    /// </summary>
    public static bool ShouldDeliverPublic(ConsumerConfig config, string subject)
        => ShouldDeliver(config, subject);

    /// <summary>
    /// Reset the idle heartbeat timer. Called whenever a data frame is delivered
    /// so that the heartbeat only fires after a period of inactivity.
    /// </summary>
    public void ResetIdleHeartbeatTimer()
    {
        _idleHeartbeatTimer?.Change(Timeout.Infinite, Timeout.Infinite);
        if (_idleHeartbeatTimer != null)
        {
            // Re-arm the timer — we'll re-read HeartbeatMs from the captured period
            var state = _idleHeartbeatTimer;
            // The timer was created with the correct period; just restart it
        }
    }

    // Go: consumer.go:1400 loopAndGatherMsgs — background loop that polls the
    // store for new messages and dispatches them to the consumer, with redelivery
    // checking and signal-channel wake-up.
    private async Task LoopAndGatherMsgsAsync(
        ConsumerHandle consumer,
        IStreamStore store,
        Func<string, string, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>, CancellationToken, ValueTask> sendMessage,
        CancellationToken ct)
    {
        var nextSeq = consumer.NextSequence;

        while (!ct.IsCancellationRequested)
        {
            // Go: consumer.go:1544 — check redelivery tracker for expired pending entries
            while (consumer.AckProcessor.TryGetExpired(out var expiredSeq, out _))
            {
                var redelivered = await store.LoadAsync(expiredSeq, ct).ConfigureAwait(false);
                if (redelivered != null)
                {
                    Enqueue(consumer, redelivered);
                    consumer.AckProcessor.ScheduleRedelivery(expiredSeq, consumer.Config.AckWaitMs);
                    GatheredCount++;

                    var rHeaders = BuildDataHeaders(redelivered);
                    var rSubject = string.IsNullOrEmpty(consumer.Config.DeliverSubject)
                        ? redelivered.Subject
                        : consumer.Config.DeliverSubject;
                    try
                    {
                        await sendMessage(rSubject, redelivered.Subject, rHeaders, redelivered.Payload, ct)
                            .ConfigureAwait(false);
                    }
                    catch (OperationCanceledException) when (ct.IsCancellationRequested)
                    {
                        return;
                    }
                }
                else
                {
                    consumer.AckProcessor.Drop(expiredSeq);
                }
            }

            // Go: consumer.go:1560 — poll store for new messages from nextSeq to LastSeq
            var storeState = await store.GetStateAsync(ct).ConfigureAwait(false);
            while (nextSeq <= storeState.LastSeq && !ct.IsCancellationRequested)
            {
                var msg = await store.LoadAsync(nextSeq, ct).ConfigureAwait(false);

                if (msg != null && ShouldDeliver(consumer.Config, msg.Subject))
                {
                    Enqueue(consumer, msg);
                    GatheredCount++;

                    var headers = BuildDataHeaders(msg);
                    var subject = string.IsNullOrEmpty(consumer.Config.DeliverSubject)
                        ? msg.Subject
                        : consumer.Config.DeliverSubject;
                    try
                    {
                        await sendMessage(subject, msg.Subject, headers, msg.Payload, ct)
                            .ConfigureAwait(false);
                    }
                    catch (OperationCanceledException) when (ct.IsCancellationRequested)
                    {
                        return;
                    }
                }

                nextSeq++;
                consumer.NextSequence = nextSeq;
            }

            // Go: consumer.go:1620 — wait for a signal or the poll timeout before re-checking
            try
            {
                using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                timeoutCts.CancelAfter(250); // Poll every 250ms if no signal arrives
                await _signalChannel!.Reader.ReadAsync(timeoutCts.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (!ct.IsCancellationRequested)
            {
                // Timeout — loop again to re-poll the store
            }
            catch (ChannelClosedException)
            {
                // Signal channel closed — loop exits via ct check above
            }
        }
    }

    // Go: consumer.go — ShouldDeliver checks cfg.FilterSubject and cfg.FilterSubjects
    // against the message subject. An empty filter delivers everything.
    private static bool ShouldDeliver(ConsumerConfig config, string subject)
    {
        if (string.IsNullOrEmpty(config.FilterSubject) && config.FilterSubjects.Count == 0)
            return true;

        if (!string.IsNullOrEmpty(config.FilterSubject))
            return SubjectMatch.MatchLiteral(subject, config.FilterSubject);

        return config.FilterSubjects.Any(f => SubjectMatch.MatchLiteral(subject, f));
    }

    // Go: consumer.go:5040 — dispatchToDeliver drains the outbound message queue.
    // For push consumers the dsubj is cfg.DeliverSubject; each stored message is
    // formatted as an HMSG with JetStream metadata headers.
    private async Task RunDeliveryLoopAsync(
        ConsumerHandle consumer,
        Func<string, string, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>, CancellationToken, ValueTask> sendMessage,
        CancellationToken ct)
    {
        var deliverSubject = consumer.Config.DeliverSubject;
        var heartbeatMs = consumer.Config.HeartbeatMs;

        while (!ct.IsCancellationRequested)
        {
            if (consumer.PushFrames.Count == 0)
            {
                // Yield to avoid busy-spin when the queue is empty
                await Task.Delay(1, ct).ConfigureAwait(false);
                continue;
            }

            var frame = consumer.PushFrames.Peek();

            // Go: consumer.go — rate-limit by honouring AvailableAtUtc before dequeuing
            var now = DateTime.UtcNow;
            if (frame.AvailableAtUtc > now)
            {
                var wait = frame.AvailableAtUtc - now;
                try
                {
                    await Task.Delay(wait, ct).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
                continue;
            }

            consumer.PushFrames.Dequeue();

            try
            {
                if (frame.IsData && frame.Message is { } msg)
                {
                    // Go: consumer.go:5067 — build JetStream metadata headers
                    // Header format: NATS/1.0\r\nNats-Sequence: {seq}\r\nNats-Time-Stamp: {ts}\r\nNats-Subject: {subj}\r\n\r\n
                    var headers = BuildDataHeaders(msg);
                    var subject = string.IsNullOrEmpty(deliverSubject) ? msg.Subject : deliverSubject;
                    await sendMessage(subject, msg.Subject, headers, msg.Payload, ct).ConfigureAwait(false);

                    // Go: consumer.go:5222 — reset idle heartbeat timer on data delivery
                    if (heartbeatMs > 0)
                        ResetIdleHeartbeatTimer(heartbeatMs);
                }
                else if (frame.IsFlowControl)
                {
                    // Go: consumer.go:5501 — "NATS/1.0 100 FlowControl Request\r\n\r\n"
                    FlowControlPendingCount++;
                    var headers = "NATS/1.0 100 FlowControl Request\r\nNats-Flow-Control: \r\n\r\n"u8.ToArray();
                    var subject = string.IsNullOrEmpty(deliverSubject) ? "_fc_" : deliverSubject;
                    await sendMessage(subject, string.Empty, headers, ReadOnlyMemory<byte>.Empty, ct).ConfigureAwait(false);
                }
                else if (frame.IsHeartbeat)
                {
                    // Go: consumer.go:5223 — "NATS/1.0 100 Idle Heartbeat\r\n..."
                    var headers = "NATS/1.0 100 Idle Heartbeat\r\n\r\n"u8.ToArray();
                    var subject = string.IsNullOrEmpty(deliverSubject) ? "_hb_" : deliverSubject;
                    await sendMessage(subject, string.Empty, headers, ReadOnlyMemory<byte>.Empty, ct).ConfigureAwait(false);
                }
            }
            catch (OperationCanceledException)
            {
                break;
            }
        }
    }

    // Go: consumer.go:5222 — start the idle heartbeat background timer
    private void StartIdleHeartbeatTimer(int heartbeatMs)
    {
        _idleHeartbeatTimer = new Timer(
            SendIdleHeartbeatCallback,
            null,
            heartbeatMs,
            heartbeatMs);
    }

    // Go: consumer.go:5222 — reset idle heartbeat timer with the configured period
    private void ResetIdleHeartbeatTimer(int heartbeatMs)
    {
        _idleHeartbeatTimer?.Change(heartbeatMs, heartbeatMs);
    }

    private void StopIdleHeartbeatTimer()
    {
        _idleHeartbeatTimer?.Dispose();
        _idleHeartbeatTimer = null;
    }

    // Go: consumer.go:5222 — sendIdleHeartbeat callback
    private void SendIdleHeartbeatCallback(object? state)
    {
        if (_sendMessage is null || _externalCt.IsCancellationRequested)
            return;

        try
        {
            // Go: consumer.go:5222 — include Nats-Pending-Messages and Nats-Pending-Bytes headers
            var pendingMsgs = _consumer?.AckProcessor.PendingCount ?? 0;
            var pendingBytes = _consumer?.PendingBytes ?? 0;
            var header = $"NATS/1.0 100 Idle Heartbeat\r\nNats-Pending-Messages: {pendingMsgs}\r\nNats-Pending-Bytes: {pendingBytes}\r\n\r\n";
            var headers = System.Text.Encoding.ASCII.GetBytes(header);

            var subject = string.IsNullOrEmpty(DeliverSubject) ? "_hb_" : DeliverSubject;
            _sendMessage(subject, string.Empty, headers, ReadOnlyMemory<byte>.Empty, _externalCt)
                .AsTask()
                .GetAwaiter()
                .GetResult();
            IdleHeartbeatsSent++;
        }
        catch (OperationCanceledException)
        {
            // Shutting down — ignore
        }
    }

    // Go: stream.go:586 — JSSequence = "Nats-Sequence", JSTimeStamp = "Nats-Time-Stamp", JSSubject = "Nats-Subject"
    private static ReadOnlyMemory<byte> BuildDataHeaders(StoredMessage msg)
    {
        var ts = msg.TimestampUtc.ToString("O");   // ISO-8601 round-trip
        var header = $"NATS/1.0\r\nNats-Sequence: {msg.Sequence}\r\nNats-Time-Stamp: {ts}\r\nNats-Subject: {msg.Subject}\r\n\r\n";
        return Encoding.ASCII.GetBytes(header);
    }
}

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;

    /// <summary>
    /// The NATS subject of the delivered message. Populated for data frames;
    /// empty string for heartbeat and flow-control frames.
    /// Mirrors the Go server's deliver-subject routing (consumer.go).
    /// </summary>
    public string Subject => Message?.Subject ?? string.Empty;
}
