// 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)
//   Will message delivery     — mqttDeliverWill (lines ~490–530)

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

namespace NATS.Server.Mqtt;

/// <summary>
/// Per-client flapper detection state tracking connect/disconnect cycling.
/// Go reference: server/mqtt.go mqttCheckFlapper ~line 300.
/// </summary>
public sealed class FlapperState
{
    /// <summary>Total number of connect/disconnect events tracked in the current window.</summary>
    public int ConnectDisconnectCount { get; set; }

    /// <summary>Start of the current detection window.</summary>
    public DateTime WindowStart { get; set; }

    /// <summary>When the backoff expires, or null if not currently backing off.</summary>
    public DateTime? BackoffUntil { get; set; }

    /// <summary>Current exponential backoff level (0 = 1s, 1 = 2s, 2 = 4s, …, capped at 60s).</summary>
    public int BackoffLevel { get; set; }
}

/// <summary>
/// Will message to be published on abnormal client disconnection.
/// Go reference: server/mqtt.go mqttWill struct ~line 270.
/// </summary>
public sealed class WillMessage
{
    public string Topic { get; init; } = string.Empty;
    public byte[] Payload { get; init; } = [];
    public byte QoS { get; init; }
    public bool Retain { get; init; }
    public int DelayIntervalSeconds { get; init; }
}

/// <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 ConcurrentDictionary<string, FlapperState> _flapperStates = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, WillMessage> _wills = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, (WillMessage Will, DateTime ScheduledAt)> _delayedWills = new(StringComparer.Ordinal);

    private readonly TimeSpan _flapWindow;
    private readonly int _flapThreshold;
    private readonly TimeSpan _flapBackoff;
    private readonly TimeProvider _timeProvider;
    private readonly IStreamStore? _backingStore;

    /// <summary>Backing store for JetStream persistence. Null for in-memory only.</summary>
    public IStreamStore? BackingStore => _backingStore;

    /// <summary>
    /// Delegate invoked when a will message is published.
    /// Parameters: topic, payload, qos, retain.
    /// Go reference: server/mqtt.go mqttDeliverWill ~line 495.
    /// </summary>
    public Action<string, byte[], byte, bool>? OnPublish { get; set; }

    /// <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>
    /// Initializes a new session store with an optional JetStream backing store.
    /// </summary>
    /// <param name="backingStore">Optional JetStream stream store for persistence.</param>
    /// <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(
        IStreamStore? backingStore,
        TimeSpan? flapWindow = null,
        int flapThreshold = 3,
        TimeSpan? flapBackoff = null,
        TimeProvider? timeProvider = null)
        : this(flapWindow, flapThreshold, flapBackoff, timeProvider)
    {
        _backingStore = backingStore;
    }

    /// <summary>
    /// Sets the will message for the given client, replacing any existing will.
    /// Called when a CONNECT packet with will flag is received.
    /// Go reference: server/mqtt.go mqttSession will field ~line 270.
    /// </summary>
    public void SetWill(string clientId, WillMessage will)
    {
        ArgumentNullException.ThrowIfNull(will);
        _wills[clientId] = will;
    }

    /// <summary>
    /// Clears the will message for the given client.
    /// Called on a clean DISCONNECT (no will should be sent).
    /// Go reference: server/mqtt.go mqttDeliverWill — will is cleared on graceful disconnect.
    /// </summary>
    public void ClearWill(string clientId)
    {
        _wills.TryRemove(clientId, out _);
        _delayedWills.TryRemove(clientId, out _);
    }

    /// <summary>
    /// Returns the current will message for the given client, or null if none.
    /// </summary>
    public WillMessage? GetWill(string clientId) =>
        _wills.TryGetValue(clientId, out var will) ? will : null;

    /// <summary>
    /// Publishes the will message for the given client on abnormal disconnection.
    /// If the will has a delay > 0, the will is recorded as delayed and not immediately published.
    /// Clears the will after publishing (or scheduling).
    /// Returns true if a will was found, false if none was registered.
    /// Go reference: server/mqtt.go mqttDeliverWill ~line 490.
    /// </summary>
    public bool PublishWillMessage(string clientId)
    {
        if (!_wills.TryRemove(clientId, out var will))
            return false;

        if (will.DelayIntervalSeconds > 0)
        {
            // Track as delayed — not immediately published.
            // A full implementation would schedule via a timer; for now we record it.
            _delayedWills[clientId] = (will, _timeProvider.GetUtcNow().UtcDateTime);
            return true;
        }

        OnPublish?.Invoke(will.Topic, will.Payload, will.QoS, will.Retain);
        return true;
    }

    /// <summary>
    /// Returns the delayed will entry for the given client if one exists,
    /// or null if the client has no pending delayed will.
    /// </summary>
    public (WillMessage Will, DateTime ScheduledAt)? GetDelayedWill(string clientId) =>
        _delayedWills.TryGetValue(clientId, out var entry) ? entry : null;

    /// <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>
    /// Backward-compatible overload: tracks a connect event (disconnect is ignored).
    /// Delegates to <see cref="TrackConnectDisconnect(string)"/> when <paramref name="connected"/> is true.
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 300.
    /// </summary>
    public void TrackConnectDisconnect(string clientId, bool connected)
    {
        if (connected)
            TrackConnectDisconnect(clientId);

        // Also maintain the legacy _connectHistory for ShouldApplyBackoff callers
        if (!connected)
            return;

        var now = _timeProvider.GetUtcNow().UtcDateTime;
        var history = _connectHistory.GetOrAdd(clientId, static _ => []);
        lock (history)
        {
            var cutoff = now - _flapWindow;
            history.RemoveAll(t => t < cutoff);
            history.Add(now);
        }
    }

    /// <summary>
    /// Tracks a connect or disconnect event for flapper detection.
    /// Increments the count for the client within the detection window. If 3+ events
    /// occur within 10 seconds the client is marked as a flapper and exponential
    /// backoff is applied: min(2^backoffLevel * 1000, 60000) ms.
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 300.
    /// </summary>
    /// <param name="clientId">The MQTT client identifier.</param>
    /// <returns>The updated <see cref="FlapperState"/> for the client.</returns>
    public FlapperState TrackConnectDisconnect(string clientId)
    {
        var now = _timeProvider.GetUtcNow().UtcDateTime;

        var state = _flapperStates.GetOrAdd(clientId, static _ => new FlapperState
        {
            WindowStart = DateTime.UtcNow,
        });

        lock (state)
        {
            // Reset window if we're past the flap window duration
            if (now - state.WindowStart > _flapWindow)
            {
                state.ConnectDisconnectCount = 0;
                state.WindowStart = now;
            }

            state.ConnectDisconnectCount++;

            if (state.ConnectDisconnectCount >= _flapThreshold)
            {
                // Exponential backoff: min(2^backoffLevel * 1000 ms, 60000 ms)
                var backoffMs = Math.Min((int)Math.Pow(2, state.BackoffLevel) * 1000, 60_000);
                state.BackoffUntil = now + TimeSpan.FromMilliseconds(backoffMs);
                state.BackoffLevel++;
            }
        }

        return state;
    }

    /// <summary>
    /// Returns true if the client is currently in a backoff period (is a flapper).
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 320.
    /// </summary>
    public bool IsFlapper(string clientId)
    {
        if (!_flapperStates.TryGetValue(clientId, out var state))
            return false;

        var now = _timeProvider.GetUtcNow().UtcDateTime;
        lock (state)
        {
            return state.BackoffUntil.HasValue && state.BackoffUntil.Value > now;
        }
    }

    /// <summary>
    /// Returns the remaining backoff in milliseconds, or 0 if the client is not flapping.
    /// Go reference: server/mqtt.go mqttCheckFlapper ~line 325.
    /// </summary>
    public long GetBackoffMs(string clientId)
    {
        if (!_flapperStates.TryGetValue(clientId, out var state))
            return 0;

        var now = _timeProvider.GetUtcNow().UtcDateTime;
        lock (state)
        {
            if (!state.BackoffUntil.HasValue || state.BackoffUntil.Value <= now)
                return 0;

            return (long)(state.BackoffUntil.Value - now).TotalMilliseconds;
        }
    }

    /// <summary>
    /// Removes all flapper tracking state for the given client.
    /// Called when stability is restored or the client is cleanly disconnected.
    /// </summary>
    public void ClearFlapperState(string clientId) =>
        _flapperStates.TryRemove(clientId, out _);

    /// <summary>
    /// Clears flapper states for clients whose <see cref="FlapperState.BackoffUntil"/> has
    /// expired by at least <paramref name="stableThreshold"/> ago, indicating the client
    /// has been stable for the given duration.
    /// </summary>
    /// <param name="stableThreshold">How long past the BackoffUntil expiry before clearing.</param>
    public void CheckAndClearStableClients(TimeSpan stableThreshold)
    {
        var now = _timeProvider.GetUtcNow().UtcDateTime;
        foreach (var (clientId, state) in _flapperStates)
        {
            lock (state)
            {
                if (state.BackoffUntil.HasValue && now - state.BackoffUntil.Value >= stableThreshold)
                {
                    _flapperStates.TryRemove(clientId, out _);
                }
            }
        }
    }

    /// <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;
        }
    }

    /// <summary>
    /// Connects a client session. If cleanSession is false, loads existing session from backing store.
    /// If cleanSession is true, deletes existing session data.
    /// Go reference: server/mqtt.go mqttInitSessionStore.
    /// </summary>
    public async Task ConnectAsync(string clientId, bool cleanSession, CancellationToken ct = default)
    {
        if (cleanSession)
        {
            DeleteSession(clientId);

            // Remove from JetStream backing store
            if (_backingStore is not null)
            {
                var subject = $"{MqttProtocolConstants.SessStreamSubjectPrefix}{clientId}";
                var msg = await _backingStore.LoadLastBySubjectAsync(subject, ct);
                if (msg is not null)
                    await _backingStore.RemoveAsync(msg.Sequence, ct);
            }

            return;
        }

        // Try to load from backing store using subject-keyed lookup
        if (_backingStore is not null)
        {
            var subject = $"{MqttProtocolConstants.SessStreamSubjectPrefix}{clientId}";
            var msg = await _backingStore.LoadLastBySubjectAsync(subject, ct);
            if (msg is not null)
            {
                var data = System.Text.Json.JsonSerializer.Deserialize<MqttSessionData>(msg.Payload.Span);
                if (data is not null)
                {
                    SaveSession(data);
                }
            }
        }
    }

    /// <summary>
    /// Adds a subscription to the client's session.
    /// </summary>
    public void AddSubscription(string clientId, string topic, int qos)
    {
        var session = LoadSession(clientId);
        if (session is null)
        {
            session = new MqttSessionData { ClientId = clientId };
        }
        session.Subscriptions[topic] = qos;
        SaveSession(session);
    }

    /// <summary>
    /// Saves the session to the backing JetStream store if available.
    /// Uses the $MQTT_sess stream with MaxMsgsPer=1 for idempotent per-subject writes.
    /// Go reference: server/mqtt.go mqttStoreSession.
    /// </summary>
    public async Task SaveSessionAsync(string clientId, CancellationToken ct = default)
    {
        var session = LoadSession(clientId);
        if (session is null || _backingStore is null)
            return;

        var json = System.Text.Json.JsonSerializer.SerializeToUtf8Bytes(session);
        await _backingStore.AppendAsync($"{MqttProtocolConstants.SessStreamSubjectPrefix}{clientId}", json, ct);
    }

    /// <summary>
    /// Returns subscriptions for the given client, or an empty dictionary.
    /// </summary>
    public IReadOnlyDictionary<string, int> GetSubscriptions(string clientId)
    {
        var session = LoadSession(clientId);
        return session?.Subscriptions ?? new Dictionary<string, int>();
    }
}
