// QoS 1 outgoing message tracker for MQTT with JetStream ack integration.
// Go reference: golang/nats-server/server/mqtt.go
//   QoS 1 outbound tracking — mqttProcessPub (~line 1200)
//   trackPublish — maps packet IDs to stream sequences for ack tracking.

using System.Collections.Concurrent;

namespace NATS.Server.Mqtt;

/// <summary>
/// Tracks outgoing QoS 1 messages pending PUBACK from the client.
/// Maps packet IDs to JetStream stream sequences for ack-based cleanup.
/// Go reference: server/mqtt.go — mqttProcessPub, trackPublish.
/// </summary>
public sealed class MqttQoS1Tracker
{
    private readonly ConcurrentDictionary<ushort, QoS1PendingMessage> _pending = new();
    private ushort _nextPacketId;
    private readonly Lock _lock = new();

    /// <summary>Number of messages pending PUBACK.</summary>
    public int PendingCount => _pending.Count;

    /// <summary>
    /// Registers an outgoing QoS 1 message and assigns a packet ID.
    /// Returns the assigned packet ID.
    /// </summary>
    /// <param name="topic">MQTT topic for the outbound message.</param>
    /// <param name="payload">Outbound payload bytes.</param>
    /// <param name="streamSequence">Optional JetStream stream sequence tied to this delivery.</param>
    public ushort Register(string topic, byte[] payload, ulong streamSequence = 0)
    {
        var id = GetNextPacketId();
        _pending[id] = new QoS1PendingMessage
        {
            PacketId = id,
            Topic = topic,
            Payload = payload,
            SentAtUtc = DateTime.UtcNow,
            DeliveryCount = 1,
            StreamSequence = streamSequence,
        };
        return id;
    }

    /// <summary>
    /// Acknowledges receipt of a PUBACK for the given packet ID.
    /// Returns the pending message if found, or null.
    /// </summary>
    /// <param name="packetId">MQTT packet identifier from PUBACK.</param>
    public QoS1PendingMessage? Acknowledge(ushort packetId)
    {
        return _pending.TryRemove(packetId, out var msg) ? msg : null;
    }

    /// <summary>
    /// Returns all pending messages for redelivery (e.g., on reconnect).
    /// Increments their delivery count.
    /// </summary>
    public IReadOnlyList<QoS1PendingMessage> GetPendingForRedelivery()
    {
        var result = new List<QoS1PendingMessage>();
        foreach (var kvp in _pending)
        {
            var msg = kvp.Value;
            msg.DeliveryCount++;
            msg.SentAtUtc = DateTime.UtcNow;
            result.Add(msg);
        }
        return result;
    }

    /// <summary>
    /// Checks if a packet ID is pending acknowledgment.
    /// </summary>
    /// <param name="packetId">MQTT packet identifier to check.</param>
    public bool IsPending(ushort packetId) => _pending.ContainsKey(packetId);

    /// <summary>Clears all pending messages.</summary>
    public void Clear() => _pending.Clear();

    private ushort GetNextPacketId()
    {
        lock (_lock)
        {
            _nextPacketId++;
            if (_nextPacketId == 0) _nextPacketId = 1; // MQTT spec: packet ID must be non-zero
            return _nextPacketId;
        }
    }
}

/// <summary>
/// A QoS 1 message pending PUBACK from the client.
/// </summary>
public sealed class QoS1PendingMessage
{
    /// <summary>MQTT packet identifier assigned to this outbound QoS1 message.</summary>
    public ushort PacketId { get; init; }
    /// <summary>MQTT topic associated with the message.</summary>
    public string Topic { get; init; } = string.Empty;
    /// <summary>Outbound payload bytes awaiting PUBACK.</summary>
    public byte[] Payload { get; init; } = [];
    /// <summary>UTC timestamp when this message was last sent.</summary>
    public DateTime SentAtUtc { get; set; }
    /// <summary>Number of delivery attempts for this packet.</summary>
    public int DeliveryCount { get; set; } = 1;

    /// <summary>
    /// JetStream stream sequence for this message. 0 if not backed by JetStream.
    /// Used to ack the message in the stream on PUBACK.
    /// Go reference: server/mqtt.go trackPublish — maps packet ID → stream sequence.
    /// </summary>
    public ulong StreamSequence { get; init; }
}