using System.Collections.Concurrent;

namespace NATS.Server.Mqtt;

/// <summary>
/// Flow controller for MQTT QoS 1/2 messages, enforcing MaxAckPending limits.
/// Uses SemaphoreSlim for async-compatible blocking when the limit is reached.
/// Go reference: server/mqtt.go — mqttMaxAckPending, flow control logic.
/// </summary>
public sealed class MqttFlowController : IDisposable
{
    private readonly ConcurrentDictionary<string, SubscriptionFlowState> _subscriptions = new(StringComparer.Ordinal);
    private int _defaultMaxAckPending;

    public MqttFlowController(int defaultMaxAckPending = 1024)
    {
        _defaultMaxAckPending = defaultMaxAckPending;
    }

    /// <summary>Default MaxAckPending limit for new subscriptions.</summary>
    public int DefaultMaxAckPending => _defaultMaxAckPending;

    /// <summary>
    /// Tries to acquire a slot for sending a QoS message on the given subscription.
    /// Returns true if a slot was acquired, false if the limit would be exceeded.
    /// </summary>
    public async ValueTask<bool> TryAcquireAsync(string subscriptionId, CancellationToken ct = default)
    {
        var state = GetOrCreate(subscriptionId);
        return await state.Semaphore.WaitAsync(0, ct) || false;
    }

    /// <summary>
    /// Waits for a slot to become available. Blocks until one is released or cancelled.
    /// </summary>
    public async ValueTask AcquireAsync(string subscriptionId, CancellationToken ct = default)
    {
        var state = GetOrCreate(subscriptionId);
        await state.Semaphore.WaitAsync(ct);
    }

    /// <summary>
    /// Releases a slot after receiving PUBACK/PUBCOMP.
    /// If the semaphore is already at max (duplicate or spurious ack), the release is a no-op.
    /// </summary>
    public void Release(string subscriptionId)
    {
        if (_subscriptions.TryGetValue(subscriptionId, out var state))
        {
            // Guard against releasing more than the max (e.g. duplicate PUBACK).
            // CurrentCount == MaxAckPending means nothing is pending, so there is nothing to release.
            if (state.Semaphore.CurrentCount < state.MaxAckPending)
                state.Semaphore.Release();
        }
    }

    /// <summary>
    /// Returns the current pending count for a subscription.
    /// </summary>
    public int GetPendingCount(string subscriptionId)
    {
        if (!_subscriptions.TryGetValue(subscriptionId, out var state))
            return 0;
        return state.MaxAckPending - state.Semaphore.CurrentCount;
    }

    /// <summary>
    /// Updates the MaxAckPending limit (e.g., on config reload).
    /// Creates a new semaphore with the updated limit.
    /// </summary>
    public void UpdateLimit(int newLimit)
    {
        _defaultMaxAckPending = newLimit;
        // Note: existing subscriptions keep their old limit until re-created
    }

    /// <summary>
    /// Removes tracking for a subscription.
    /// </summary>
    public void RemoveSubscription(string subscriptionId)
    {
        if (_subscriptions.TryRemove(subscriptionId, out var state))
            state.Semaphore.Dispose();
    }

    /// <summary>Number of tracked subscriptions.</summary>
    public int SubscriptionCount => _subscriptions.Count;

    public void Dispose()
    {
        foreach (var kvp in _subscriptions)
            kvp.Value.Semaphore.Dispose();
        _subscriptions.Clear();
    }

    private SubscriptionFlowState GetOrCreate(string subscriptionId)
    {
        return _subscriptions.GetOrAdd(subscriptionId, _ => new SubscriptionFlowState
        {
            MaxAckPending = _defaultMaxAckPending,
            Semaphore = new SemaphoreSlim(_defaultMaxAckPending, _defaultMaxAckPending),
        });
    }

    private sealed class SubscriptionFlowState
    {
        public int MaxAckPending { get; init; }
        public required SemaphoreSlim Semaphore { get; init; }
    }
}
