using System.Collections.Concurrent;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;
using ZB.MOM.WW.MxGateway.Server.Configuration;

namespace ZB.MOM.WW.MxGateway.Server.Security.Authentication;

/// <summary>
/// An <see cref="IApiKeyStore"/> decorator that coalesces <see cref="MarkUsedAsync"/> writes to at
/// most one per key per <see cref="SecurityOptions.ApiKeyLastUsedCoalesceSeconds"/> window. Lookups
/// pass through unchanged.
/// </summary>
/// <remarks>
/// The library verifier couples the <c>last_used_utc</c> write into every successful verification
/// by calling <see cref="IApiKeyStore.MarkUsedAsync"/>. On the bulk-read hot path that turns a
/// per-RPC database write into the throughput ceiling. This decorator sits under the library
/// verifier: it forwards the first mark for a key, then drops subsequent marks until the window
/// elapses, so <c>last_used_utc</c> is refreshed at most once per key per window while the read
/// path stays correct. <c>last_used_utc</c> is a coarse staleness indicator, not an audit record
/// (audit rows are written separately), so bounded staleness of up to one window is acceptable.
/// </remarks>
public sealed class CoalescingMarkApiKeyStore : IApiKeyStore
{
    private readonly IApiKeyStore _inner;
    private readonly TimeSpan _window;
    private readonly TimeProvider _clock;

    // keyId -> UtcTicks of the last forwarded mark. Bounded by the number of API keys, which is
    // small; no eviction needed.
    private readonly ConcurrentDictionary<string, long> _lastMarkedTicks = new(StringComparer.Ordinal);

    /// <summary>Initializes a new instance of the <see cref="CoalescingMarkApiKeyStore"/> class.</summary>
    /// <param name="inner">The wrapped store.</param>
    /// <param name="security">Security options carrying the coalescing window.</param>
    /// <param name="clock">The time provider.</param>
    public CoalescingMarkApiKeyStore(IApiKeyStore inner, SecurityOptions security, TimeProvider clock)
        : this(
            inner,
            TimeSpan.FromSeconds((security ?? throw new ArgumentNullException(nameof(security)))
                .ApiKeyLastUsedCoalesceSeconds),
            clock)
    {
    }

    /// <summary>Initializes a new instance of the <see cref="CoalescingMarkApiKeyStore"/> class with an explicit coalescing window (test/explicit-window seam).</summary>
    /// <param name="inner">The wrapped store.</param>
    /// <param name="window">The coalescing window; marks within this window of the last forwarded mark for a key are dropped.</param>
    /// <param name="clock">The time provider.</param>
    internal CoalescingMarkApiKeyStore(IApiKeyStore inner, TimeSpan window, TimeProvider clock)
    {
        ArgumentNullException.ThrowIfNull(inner);
        ArgumentNullException.ThrowIfNull(clock);
        _inner = inner;
        _window = window;
        _clock = clock;
    }

    /// <summary>Looks up an API key record by key id, delegating to the wrapped store unchanged.</summary>
    /// <param name="keyId">The API key id to look up.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>The matching <see cref="ApiKeyRecord"/>, or <see langword="null"/> if none exists.</returns>
    public Task<ApiKeyRecord?> FindByKeyIdAsync(string keyId, CancellationToken ct)
        => _inner.FindByKeyIdAsync(keyId, ct);

    /// <summary>Looks up an active API key record by key id, delegating to the wrapped store unchanged.</summary>
    /// <param name="keyId">The API key id to look up.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>The matching active <see cref="ApiKeyRecord"/>, or <see langword="null"/> if none exists or is inactive.</returns>
    public Task<ApiKeyRecord?> FindActiveByKeyIdAsync(string keyId, CancellationToken ct)
        => _inner.FindActiveByKeyIdAsync(keyId, ct);

    /// <summary>Marks the key as used, coalescing writes so at most one reaches the wrapped store per key per window.</summary>
    /// <param name="keyId">The API key id being marked as used.</param>
    /// <param name="whenUtc">The UTC timestamp of the use.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public Task MarkUsedAsync(string keyId, DateTimeOffset whenUtc, CancellationToken ct)
    {
        ArgumentNullException.ThrowIfNull(keyId);

        if (_window <= TimeSpan.Zero)
        {
            return _inner.MarkUsedAsync(keyId, whenUtc, ct);
        }

        long now = _clock.GetUtcNow().UtcTicks;

        while (true)
        {
            if (_lastMarkedTicks.TryGetValue(keyId, out long last))
            {
                if (now - last < _window.Ticks)
                {
                    // Within the coalescing window — drop the write.
                    return Task.CompletedTask;
                }

                if (_lastMarkedTicks.TryUpdate(keyId, now, last))
                {
                    return _inner.MarkUsedAsync(keyId, whenUtc, ct);
                }

                // Lost the race to another writer; re-evaluate.
                continue;
            }

            if (_lastMarkedTicks.TryAdd(keyId, now))
            {
                return _inner.MarkUsedAsync(keyId, whenUtc, ct);
            }

            // Lost the race to another writer; re-evaluate.
        }
    }
}
