// Manages per-subscription JetStream consumers for MQTT QoS 1/2 delivery.
// Go reference: golang/nats-server/server/mqtt.go mqttAccountSessionManager ~line 600
//   Consumer creation per subscription, ack tracking, session resume.

using System.Collections.Concurrent;
using NATS.Server.JetStream;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.Mqtt;

/// <summary>
/// Tracks the mapping between an MQTT subscription and its JetStream consumer.
/// </summary>
public sealed record MqttConsumerBinding(string Stream, string DurableName, string FilterSubject);

/// <summary>
/// Manages per-subscription JetStream consumers for MQTT QoS 1/2 message delivery.
/// Each QoS > 0 subscription gets a durable consumer on $MQTT_msgs.
/// Go reference: server/mqtt.go — mqttAccountSessionManager consumer management.
/// </summary>
public sealed class MqttConsumerManager
{
    private readonly StreamManager _streamManager;
    private readonly ConsumerManager _consumerManager;
    private readonly ConcurrentDictionary<string, MqttConsumerBinding> _bindings = new(StringComparer.Ordinal);

    public MqttConsumerManager(StreamManager streamManager, ConsumerManager consumerManager)
    {
        _streamManager = streamManager;
        _consumerManager = consumerManager;
    }

    /// <summary>
    /// Creates a durable JetStream consumer for a QoS > 0 MQTT subscription.
    /// The consumer filters $MQTT_msgs by the translated NATS subject.
    /// Returns the binding, or null if creation failed.
    /// Go reference: server/mqtt.go mqttProcessSub consumer creation.
    /// </summary>
    public MqttConsumerBinding? CreateSubscriptionConsumer(string clientId, string natsSubject, int qos, int maxAckPending)
    {
        var durableName = $"$MQTT_{clientId}_{natsSubject.Replace('.', '_').Replace('*', 'W').Replace('>', 'G')}";
        var filterSubject = $"{MqttProtocolConstants.StreamSubjectPrefix}{natsSubject}";

        var response = _consumerManager.CreateOrUpdate(MqttProtocolConstants.StreamName, new ConsumerConfig
        {
            DurableName = durableName,
            FilterSubject = filterSubject,
            AckPolicy = AckPolicy.Explicit,
            DeliverPolicy = DeliverPolicy.All,
            MaxAckPending = maxAckPending,
            AckWaitMs = (int)MqttProtocolConstants.DefaultAckWait.TotalMilliseconds,
            MaxDeliver = -1,
        });

        if (response.Error != null)
            return null;

        var binding = new MqttConsumerBinding(MqttProtocolConstants.StreamName, durableName, filterSubject);
        _bindings[$"{clientId}:{natsSubject}"] = binding;
        return binding;
    }

    /// <summary>
    /// Removes the JetStream consumer for an MQTT subscription.
    /// Called on UNSUBSCRIBE or clean session disconnect.
    /// </summary>
    public void RemoveSubscriptionConsumer(string clientId, string natsSubject)
    {
        var key = $"{clientId}:{natsSubject}";
        if (_bindings.TryRemove(key, out var binding))
        {
            _consumerManager.Delete(binding.Stream, binding.DurableName);
        }
    }

    /// <summary>
    /// Removes all consumers for a client. Called on clean session disconnect.
    /// </summary>
    public void RemoveAllConsumers(string clientId)
    {
        var prefix = $"{clientId}:";
        var keysToRemove = _bindings.Keys.Where(k => k.StartsWith(prefix, StringComparison.Ordinal)).ToList();
        foreach (var key in keysToRemove)
        {
            if (_bindings.TryRemove(key, out var binding))
            {
                _consumerManager.Delete(binding.Stream, binding.DurableName);
            }
        }
    }

    /// <summary>
    /// Gets the binding for a subscription, or null if none exists.
    /// </summary>
    public MqttConsumerBinding? GetBinding(string clientId, string natsSubject)
    {
        return _bindings.TryGetValue($"{clientId}:{natsSubject}", out var binding) ? binding : null;
    }

    /// <summary>
    /// Gets all bindings for a client (for session persistence).
    /// </summary>
    public IReadOnlyDictionary<string, MqttConsumerBinding> GetClientBindings(string clientId)
    {
        var prefix = $"{clientId}:";
        return _bindings
            .Where(kvp => kvp.Key.StartsWith(prefix, StringComparison.Ordinal))
            .ToDictionary(kvp => kvp.Key[prefix.Length..], kvp => kvp.Value);
    }

    /// <summary>
    /// Publishes a message to the $MQTT_msgs stream for QoS delivery.
    /// Returns the sequence number, or 0 if publish failed.
    /// </summary>
    public ulong PublishToStream(string natsSubject, ReadOnlyMemory<byte> payload)
    {
        var subject = $"{MqttProtocolConstants.StreamSubjectPrefix}{natsSubject}";
        if (_streamManager.TryGet(MqttProtocolConstants.StreamName, out var handle))
        {
            var seq = handle.Store.AppendAsync(subject, payload, default).GetAwaiter().GetResult();
            return seq;
        }

        return 0;
    }

    /// <summary>
    /// Acknowledges a message in the stream by removing it (for interest-based retention).
    /// Called when PUBACK is received for QoS 1.
    /// </summary>
    public bool AcknowledgeMessage(ulong sequence)
    {
        if (_streamManager.TryGet(MqttProtocolConstants.StreamName, out var handle))
        {
            return handle.Store.RemoveAsync(sequence, default).GetAwaiter().GetResult();
        }

        return false;
    }

    /// <summary>
    /// Loads a message from the $MQTT_msgs stream by sequence.
    /// </summary>
    public async ValueTask<StoredMessage?> LoadMessageAsync(ulong sequence, CancellationToken ct = default)
    {
        if (_streamManager.TryGet(MqttProtocolConstants.StreamName, out var handle))
        {
            return await handle.Store.LoadAsync(sequence, ct);
        }

        return null;
    }

    /// <summary>
    /// Stores a QoS 2 incoming message for deduplication.
    /// Returns the sequence number, or 0 if failed.
    /// </summary>
    public ulong StoreQoS2Incoming(string clientId, ushort packetId, ReadOnlyMemory<byte> payload)
    {
        var subject = $"{MqttProtocolConstants.QoS2IncomingMsgsStreamSubjectPrefix}{clientId}.{packetId}";
        if (_streamManager.TryGet(MqttProtocolConstants.QoS2IncomingMsgsStreamName, out var handle))
        {
            return handle.Store.AppendAsync(subject, payload, default).GetAwaiter().GetResult();
        }

        return 0;
    }

    /// <summary>
    /// Loads a QoS 2 incoming message for delivery on PUBREL.
    /// </summary>
    public async ValueTask<StoredMessage?> LoadQoS2IncomingAsync(string clientId, ushort packetId, CancellationToken ct = default)
    {
        var subject = $"{MqttProtocolConstants.QoS2IncomingMsgsStreamSubjectPrefix}{clientId}.{packetId}";
        if (_streamManager.TryGet(MqttProtocolConstants.QoS2IncomingMsgsStreamName, out var handle))
        {
            return await handle.Store.LoadLastBySubjectAsync(subject, ct);
        }

        return null;
    }

    /// <summary>
    /// Removes a QoS 2 incoming message after PUBCOMP.
    /// </summary>
    public async ValueTask<bool> RemoveQoS2IncomingAsync(string clientId, ushort packetId, CancellationToken ct = default)
    {
        var subject = $"{MqttProtocolConstants.QoS2IncomingMsgsStreamSubjectPrefix}{clientId}.{packetId}";
        if (_streamManager.TryGet(MqttProtocolConstants.QoS2IncomingMsgsStreamName, out var handle))
        {
            var msg = await handle.Store.LoadLastBySubjectAsync(subject, ct);
            if (msg != null)
                return await handle.Store.RemoveAsync(msg.Sequence, ct);
        }

        return false;
    }
}
