using System.Collections.Concurrent;

namespace Mbproxy.Proxy.Multiplexing;

/// <summary>
/// Maps a proxy-assigned MBAP TxId → <see cref="InFlightRequest"/>. The multiplexer's
/// per-upstream <c>OnFrame</c> path adds entries; the backend reader task removes them
/// when the matching response arrives.
///
/// <para>Backed by <see cref="ConcurrentDictionary{TKey, TValue}"/>. The single-writer /
/// single-remover pattern does not strictly require it — but cascade-on-disconnect walks
/// the map from a separate task and the coalescing path adds upstream-side cancellation
/// paths, so the safer primitive is worth the negligible cost.</para>
/// </summary>
internal sealed class CorrelationMap
{
    private readonly ConcurrentDictionary<ushort, InFlightRequest> _entries = new();

    /// <summary>
    /// Adds <paramref name="req"/> under <paramref name="proxyTxId"/>. Returns <c>false</c>
    /// if a request was already stored under that key — which would be a programming
    /// error (the allocator should never hand out the same key twice while it is still
    /// in flight). Callers should treat <c>false</c> as a fatal contract violation and
    /// drop the upstream connection.
    /// </summary>
    public bool TryAdd(ushort proxyTxId, InFlightRequest req)
        => _entries.TryAdd(proxyTxId, req);

    /// <summary>
    /// Removes the entry under <paramref name="proxyTxId"/>. Returns <c>false</c> when
    /// no entry exists (which is normal for cascade cleanup and for stale-response paths).
    /// </summary>
    public bool TryRemove(ushort proxyTxId, out InFlightRequest req)
        => _entries.TryRemove(proxyTxId, out req!);

    /// <summary>Number of currently-in-flight requests.</summary>
    public int Count => _entries.Count;

    /// <summary>
    /// Returns a point-in-time copy of all in-flight requests. Allocates a list; intended
    /// for diagnostics (cascade walk on backend disconnect; future drain-on-shutdown).
    /// </summary>
    public IReadOnlyCollection<InFlightRequest> Snapshot()
    {
        // ConcurrentDictionary.Values is a snapshot-safe enumerable; materialise to
        // detach from the live dictionary and give callers a stable view.
        return _entries.Values.ToArray();
    }

    /// <summary>
    /// Returns and removes every entry. Used by the multiplexer's cascade path when the
    /// backend socket dies — the multiplexer must close every interested upstream pipe
    /// and free every allocated proxy TxId.
    /// </summary>
    public IReadOnlyList<KeyValuePair<ushort, InFlightRequest>> DrainAll()
    {
        var drained = new List<KeyValuePair<ushort, InFlightRequest>>(_entries.Count);
        foreach (var kvp in _entries)
        {
            if (_entries.TryRemove(kvp.Key, out var req))
                drained.Add(new KeyValuePair<ushort, InFlightRequest>(kvp.Key, req));
        }
        return drained;
    }

    /// <summary>
    /// Returns a snapshot of (proxyTxId, InFlightRequest) pairs whose <see cref="InFlightRequest.SentAtUtc"/>
    /// is older than <paramref name="threshold"/>. Allocates a list; intended for the
    /// periodic per-request timeout watchdog only. The entries are NOT removed by this
    /// call — the caller decides which to time out.
    /// </summary>
    public IReadOnlyList<KeyValuePair<ushort, InFlightRequest>> SnapshotOlderThan(DateTimeOffset threshold)
    {
        var stale = new List<KeyValuePair<ushort, InFlightRequest>>();
        foreach (var kvp in _entries)
        {
            if (kvp.Value.SentAtUtc <= threshold)
                stale.Add(new KeyValuePair<ushort, InFlightRequest>(kvp.Key, kvp.Value));
        }
        return stale;
    }
}
