// QoS 1 outgoing message tracker for MQTT.
// Go reference: golang/nats-server/server/mqtt.go
//   QoS 1 outbound tracking — mqttProcessPub (~line 1200)

using System.Collections.Concurrent;

namespace NATS.Server.Mqtt;

/// <summary>
/// Tracks outgoing QoS 1 messages pending PUBACK from the client.
/// Messages are stored with their packet ID and can be redelivered on reconnect.
/// Go reference: server/mqtt.go — mqttProcessPub (QoS 1 outbound tracking).
/// </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>
    public ushort Register(string topic, byte[] payload)
    {
        var id = GetNextPacketId();
        _pending[id] = new QoS1PendingMessage
        {
            PacketId = id,
            Topic = topic,
            Payload = payload,
            SentAtUtc = DateTime.UtcNow,
            DeliveryCount = 1,
        };
        return id;
    }

    /// <summary>
    /// Acknowledges receipt of a PUBACK for the given packet ID.
    /// Returns true if the message was found and removed.
    /// </summary>
    public bool Acknowledge(ushort packetId)
    {
        return _pending.TryRemove(packetId, out _);
    }

    /// <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>
    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
{
    public ushort PacketId { get; init; }
    public string Topic { get; init; } = string.Empty;
    public byte[] Payload { get; init; } = [];
    public DateTime SentAtUtc { get; set; }
    public int DeliveryCount { get; set; } = 1;
}
