// MQTT session persistence store.
// Go reference: golang/nats-server/server/mqtt.go:253-300
//   Session state management  — mqttInitSessionStore / mqttStoreSession
//   Flapper detection         — mqttCheckFlapper (lines ~300–360)

using System.Collections.Concurrent;

namespace NATS.Server.Mqtt;

/// <summary>
/// Serializable session data for an MQTT client.
/// Go reference: server/mqtt.go mqttSession struct ~line 253.
/// </summary>
public sealed record MqttSessionData
{
    public required string ClientId { get; init; }
    public Dictionary<string, int> Subscriptions { get; init; } = [];
    public List<MqttPendingPublish> PendingPublishes { get; init; } = [];
    public string? WillTopic { get; init; }
    public byte[]? WillPayload { get; init; }
    public int WillQoS { get; init; }
    public bool WillRetain { get; init; }
    public bool CleanSession { get; init; }
    public DateTime ConnectedAtUtc { get; init; } = DateTime.UtcNow;
    public DateTime LastActivityUtc { get; set; } = DateTime.UtcNow;
}

/// <summary>
/// In-memory MQTT session store with flapper detection.
/// The abstraction allows future JetStream backing.
/// Go reference: server/mqtt.go mqttInitSessionStore ~line 260.
/// </summary>
public sealed class MqttSessionStore
{
    private readonly ConcurrentDictionary<string, MqttSessionData> _sessions = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, List<DateTime>> _connectHistory = new(StringComparer.Ordinal);

    private readonly TimeSpan _flapWindow;
    private readonly int _flapThreshold;
    private readonly TimeSpan _flapBackoff;
    private readonly TimeProvider _timeProvider;

    /// <summary>
    /// Initializes a new session store.
    /// </summary>
    /// <param name="flapWindow">Window in which repeated connects trigger flap detection. Default 10 seconds.</param>
    /// <param name="flapThreshold">Number of connects within the window to trigger backoff. Default 3.</param>
    /// <param name="flapBackoff">Backoff delay to apply when flapping. Default 1 second.</param>
    /// <param name="timeProvider">Optional time provider for testing. Default uses system clock.</param>
    public MqttSessionStore(
        TimeSpan? flapWindow = null,
        int flapThreshold = 3,
        TimeSpan? flapBackoff = null,
        TimeProvider? timeProvider = null)
    {
        _flapWindow = flapWindow ?? TimeSpan.FromSeconds(10);
        _flapThreshold = flapThreshold;
        _flapBackoff = flapBackoff ?? TimeSpan.FromSeconds(1);
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    /// <summary>
    /// Saves (or overwrites) session data for the given client.
    /// Go reference: server/mqtt.go mqttStoreSession.
    /// </summary>
    public void SaveSession(MqttSessionData session)
    {
        ArgumentNullException.ThrowIfNull(session);
        _sessions[session.ClientId] = session;
    }

    /// <summary>
    /// Loads session data for the given client, or null if not found.
    /// Go reference: server/mqtt.go mqttLoadSession.
    /// </summary>
    public MqttSessionData? LoadSession(string clientId) =>
        _sessions.TryGetValue(clientId, out var session) ? session : null;

    /// <summary>
    /// Deletes the session for the given client. No-op if not found.
    /// Go reference: server/mqtt.go mqttDeleteSession.
    /// </summary>
    public void DeleteSession(string clientId) =>
        _sessions.TryRemove(clientId, out _);

    /// <summary>
    /// Returns all active sessions.
    /// </summary>
    public IReadOnlyList<MqttSessionData> ListSessions() =>
        _sessions.Values.ToList();

    /// <summary>
    /// Tracks a connect or disconnect event for flapper detection.
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 300.
    /// </summary>
    /// <param name="clientId">The MQTT client identifier.</param>
    /// <param name="connected">True for connect, false for disconnect.</param>
    public void TrackConnectDisconnect(string clientId, bool connected)
    {
        if (!connected)
            return;

        var now = _timeProvider.GetUtcNow().UtcDateTime;
        var history = _connectHistory.GetOrAdd(clientId, static _ => []);

        lock (history)
        {
            // Prune entries outside the flap window
            var cutoff = now - _flapWindow;
            history.RemoveAll(t => t < cutoff);
            history.Add(now);
        }
    }

    /// <summary>
    /// Returns the backoff delay if the client is flapping, otherwise <see cref="TimeSpan.Zero"/>.
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 320.
    /// </summary>
    public TimeSpan ShouldApplyBackoff(string clientId)
    {
        if (!_connectHistory.TryGetValue(clientId, out var history))
            return TimeSpan.Zero;

        var now = _timeProvider.GetUtcNow().UtcDateTime;

        lock (history)
        {
            var cutoff = now - _flapWindow;
            history.RemoveAll(t => t < cutoff);
            return history.Count >= _flapThreshold ? _flapBackoff : TimeSpan.Zero;
        }
    }
}
