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

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

/// <summary>
/// Cheap, in-process sliding-window failure counter for the gRPC auth path. It is checked BEFORE the
/// API-key verification store read, so online secret guessing cannot spend a SQLite read per attempt,
/// and it admits a periodic probe so a throttle can never deny a key to the holder of the correct
/// secret.
/// </summary>
/// <remarks>
/// <para>
/// Two layers share one window mechanism. <b>Layer 1</b> partitions on the composite
/// <c>(transport peer, key id)</c>: an attacker's failures bind to the address that produced them,
/// so a peer spraying guesses at a key id it does not own throttles only itself, while co-located
/// clients using other key ids are untouched. <b>Layer 2</b> counts failures for a key id across all
/// peers (<see cref="SecurityOptions.ApiKeyFailureAggregateLimit"/>), which bounds a distributed or
/// source-rotating sprayer that never trips any single partition.
/// </para>
/// <para>
/// An over-limit state is a valve, not a wall (SEC-31): at most one request per
/// <see cref="SecurityOptions.ApiKeyFailureProbeIntervalSeconds"/> is admitted through to the real
/// verifier, and a successful verification resets both layers. The reset path therefore stays
/// reachable while throttled — the earlier design keyed solely on the attacker-supplied key id and
/// rejected before verification, so the legitimate holder could never clear the block.
/// </para>
/// <para>
/// The tracked partition set is a bounded LRU (<see cref="SecurityOptions.ApiKeyFailureTrackedPeers"/>).
/// Because only a validly shaped token mints a key-id partition and each transport peer may mint at
/// most <see cref="MaxKeyIdPartitionsPerPeer"/> of them (the overflow collapses onto that address's
/// fallback partition), the map bounds memory <em>and</em> resists a junk-token flush (SEC-32):
/// eviction prefers fully expired windows, never removes an over-limit partition below a transient
/// overshoot ceiling of twice the cap, and only then falls back to least-recently-active. Successful
/// peers are removed on reset, so in steady state the map holds only partitions with recent failures
/// — the common success path is a lock-free dictionary miss.
/// </para>
/// <para>
/// Eviction is deliberately a best-effort full scan of the over-capacity map, run only on the
/// failure path and only once the cap is exceeded. Under a sustained overflow that scan is O(n) per
/// recorded failure, which is accepted: n is bounded by the cap, the work lands on attack traffic
/// rather than on authenticated calls, and an exact ordered structure would need a second index kept
/// consistent with the per-state locks for no security gain.
/// </para>
/// </remarks>
public sealed class ApiKeyFailureLimiter
{
    /// <summary>
    /// Maximum distinct key-id partitions one transport peer may mint. Not configurable: no
    /// legitimate address fails against dozens of distinct keys inside one window, and the cap is
    /// what bounds an attacker's total partitions to this value plus its fallback partition.
    /// </summary>
    internal const int MaxKeyIdPartitionsPerPeer = 32;

    // NUL appears in neither a gRPC peer string ("ipv4:10.0.0.1:5000") nor an ASCII metadata
    // value, so a composite key can never collide with a transport-peer-only fallback key.
    // A collision would merge two partitions, never widen an admission.
    private const char PartitionSeparator = '\0';

    private readonly int _limit;
    private readonly int _aggregateLimit;
    private readonly long _windowTicks;
    private readonly long _probeIntervalTicks;
    private readonly int _maxPartitions;
    private readonly TimeProvider _clock;

    private readonly ConcurrentDictionary<string, WindowState> _partitions = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, WindowState> _aggregates = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, PeerKeyIds> _peerKeyIds = new(StringComparer.Ordinal);

    /// <summary>Initializes a new instance of the <see cref="ApiKeyFailureLimiter"/> class.</summary>
    /// <param name="security">Security options carrying the failure-limit knobs.</param>
    /// <param name="clock">The time provider.</param>
    public ApiKeyFailureLimiter(SecurityOptions security, TimeProvider clock)
        : this(
            (security ?? throw new ArgumentNullException(nameof(security))).ApiKeyFailureLimit,
            TimeSpan.FromSeconds(security.ApiKeyFailureWindowSeconds),
            security.ApiKeyFailureTrackedPeers,
            security.ApiKeyFailureAggregateLimit,
            TimeSpan.FromSeconds(security.ApiKeyFailureProbeIntervalSeconds),
            clock)
    {
    }

    /// <summary>Initializes a new instance of the <see cref="ApiKeyFailureLimiter"/> class. Test/explicit seam.</summary>
    /// <param name="limit">The maximum number of failures allowed per partition within <paramref name="window"/>.</param>
    /// <param name="window">The sliding window over which failures are counted.</param>
    /// <param name="maxPartitions">The maximum number of tracked partitions before eviction kicks in.</param>
    /// <param name="aggregateLimit">The per-key-id cross-peer failure limit; <c>0</c> disables the aggregate layer.</param>
    /// <param name="probeInterval">The minimum interval between probe admissions; <see cref="TimeSpan.Zero"/> blocks absolutely.</param>
    /// <param name="clock">The time provider.</param>
    internal ApiKeyFailureLimiter(
        int limit,
        TimeSpan window,
        int maxPartitions,
        int aggregateLimit,
        TimeSpan probeInterval,
        TimeProvider clock)
    {
        ArgumentNullException.ThrowIfNull(clock);
        _limit = limit;
        _windowTicks = window.Ticks;
        _maxPartitions = maxPartitions;
        _aggregateLimit = aggregateLimit;
        _probeIntervalTicks = probeInterval.Ticks;
        _clock = clock;
    }

    /// <summary>Gets the number of tracked <c>(peer, key id)</c> partitions. Test seam.</summary>
    internal int TrackedPartitionCount => _partitions.Count;

    /// <summary>Gets the number of tracked per-key-id aggregates. Test seam.</summary>
    internal int TrackedAggregateCount => _aggregates.Count;

    /// <summary>
    /// Test seam invoked between a probe claim and its compensating release; <see langword="null"/>
    /// in production, where it costs one null check on the already-refused path.
    /// </summary>
    /// <remarks>
    /// The window this straddles is a few nanoseconds wide, so the interleaving it exists to cover —
    /// a concurrent failure re-arming the same state between the two — cannot be produced reliably
    /// by racing threads. Without this seam the version guard in <see cref="ReleaseProbe"/> would be
    /// verifiable only by inspection.
    /// </remarks>
    internal Action? ProbeReleaseInterleaveHook { get; set; }

    /// <summary>Decides whether an authentication attempt may reach the verifier.</summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    /// <returns>The admission decision for this attempt.</returns>
    public ApiKeyThrottleDecision Check(ApiKeyThrottlePartition partition) => Check(partition, out _);

    /// <summary>
    /// Decides whether an authentication attempt may reach the verifier, handing back the storage
    /// key it resolved so a caller that goes on to <see cref="Reset(ApiKeyThrottlePartition, PartitionResolution)"/>
    /// the same request does not resolve — and rebuild the composite key string — a second time.
    /// </summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    /// <param name="resolution">
    /// The resolved storage key, or <see cref="PartitionResolution.Unresolved"/> when the limiter is
    /// disabled and never resolved one.
    /// </param>
    /// <returns>The admission decision for this attempt.</returns>
    internal ApiKeyThrottleDecision Check(ApiKeyThrottlePartition partition, out PartitionResolution resolution)
    {
        string peer = RequirePeer(partition);
        if (_limit <= 0)
        {
            resolution = PartitionResolution.Unresolved;
            return ApiKeyThrottleDecision.Allowed;
        }

        long now = _clock.GetUtcNow().UtcTicks;
        (string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: false);
        resolution = new PartitionResolution(partitionKey, effectiveKeyId);

        WindowState? peerState = _partitions.TryGetValue(partitionKey, out WindowState? tracked) ? tracked : null;
        WindowState? aggregateState = null;
        if (effectiveKeyId is not null
            && _aggregateLimit > 0
            && _aggregates.TryGetValue(effectiveKeyId, out WindowState? aggregate))
        {
            aggregateState = aggregate;
        }

        bool peerOver = peerState is not null && IsOverLimit(peerState, now, _limit);
        bool aggregateOver = aggregateState is not null && IsOverLimit(aggregateState, now, _aggregateLimit);
        if (!peerOver && !aggregateOver)
        {
            return ApiKeyThrottleDecision.Allowed;
        }

        if (_probeIntervalTicks <= 0)
        {
            return peerOver ? ApiKeyThrottleDecision.ThrottledByPeer : ApiKeyThrottleDecision.ThrottledByAggregate;
        }

        // Every over-limit layer must yield its probe slot for the request to pass. The slots are
        // claimed one at a time (holding two per-state locks at once would need a global ordering to
        // stay deadlock-free), so a claim is reserved and then compensated if a later layer refuses.
        ProbeClaim peerClaim = default;
        bool peerProbeClaimed = false;
        if (peerOver)
        {
            if (!TryConsumeProbe(peerState!, now, out peerClaim))
            {
                return ApiKeyThrottleDecision.ThrottledByPeer;
            }

            peerProbeClaimed = true;
        }

        if (aggregateOver && !TryConsumeProbe(aggregateState!, now, out _))
        {
            if (peerProbeClaimed)
            {
                ProbeReleaseInterleaveHook?.Invoke();
                ReleaseProbe(peerState!, peerClaim);
            }

            return ApiKeyThrottleDecision.ThrottledByAggregate;
        }

        return ApiKeyThrottleDecision.ProbeAdmitted;
    }

    /// <summary>Records a failed verification attempt against both limiter layers.</summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    public void RecordFailure(ApiKeyThrottlePartition partition)
    {
        string peer = RequirePeer(partition);
        if (_limit <= 0)
        {
            return;
        }

        long now = _clock.GetUtcNow().UtcTicks;
        (string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: true);

        WindowState state = _partitions.GetOrAdd(
            partitionKey,
            static (_, owner) => new WindowState(owner.Peer, owner.KeyId),
            (Peer: peer, KeyId: effectiveKeyId));
        RecordInto(state, now, _limit);

        // Only a key id that earned its own partition feeds the aggregate: an id squeezed out by the
        // per-peer cap is spray, and letting it through would make aggregate cardinality unbounded.
        if (effectiveKeyId is not null && _aggregateLimit > 0)
        {
            WindowState aggregate = _aggregates.GetOrAdd(effectiveKeyId, static _ => new WindowState(null, null));
            RecordInto(aggregate, now, _aggregateLimit);
        }

        EvictIfOverCapacity(_partitions, now, _limit, releaseKeyIds: true);
        EvictIfOverCapacity(_aggregates, now, _aggregateLimit, releaseKeyIds: false);
    }

    /// <summary>Clears both limiter layers for the partition after a successful verification.</summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    public void Reset(ApiKeyThrottlePartition partition) => Reset(partition, PartitionResolution.Unresolved);

    /// <summary>
    /// Clears both limiter layers for the partition after a successful verification, reusing the
    /// storage key <see cref="Check(ApiKeyThrottlePartition, out PartitionResolution)"/> already
    /// resolved for this request.
    /// </summary>
    /// <param name="partition">The throttle partition derived from the request.</param>
    /// <param name="resolution">
    /// The resolution handed back by <c>Check</c>; <see cref="PartitionResolution.Unresolved"/>
    /// resolves here instead. Reusing the check-time resolution is deliberate: it is the partition
    /// this request was admitted against, so the reset clears exactly what the check consulted.
    /// </param>
    internal void Reset(ApiKeyThrottlePartition partition, PartitionResolution resolution)
    {
        string peer = RequirePeer(partition);
        (string partitionKey, string? effectiveKeyId) = resolution.IsResolved
            ? (resolution.PartitionKey!, resolution.EffectiveKeyId)
            : ResolvePartitionKey(peer, partition.KeyId, mint: false);

        // Clear only a partition this caller actually owns. When its key id was squeezed into the
        // address's shared fallback bucket by the per-peer cap, that bucket also holds failures
        // contributed by other key ids — and by junk-shaped tokens — from the same address, so
        // removing it would let one successful authentication wipe an in-progress spray. The shared
        // bucket decays by window expiry instead; the caller still recovers through probe admission.
        if (string.Equals(effectiveKeyId, partition.KeyId, StringComparison.Ordinal))
        {
            RemovePartition(partitionKey);
        }

        // The aggregate is cleared on the presented key id even when the composite partition
        // collapsed to the fallback: a verified secret is proof the key is not under successful
        // attack, and clearing is what makes the block recoverable.
        if (partition.KeyId is not null)
        {
            _aggregates.TryRemove(partition.KeyId, out _);
        }
    }

    /// <summary>Returns whether the partition currently has tracked failure state. Test seam.</summary>
    /// <param name="partition">The throttle partition to look up.</param>
    /// <returns><see langword="true"/> when the partition has a tracked window.</returns>
    internal bool IsTracked(ApiKeyThrottlePartition partition)
    {
        string peer = RequirePeer(partition);
        (string partitionKey, _) = ResolvePartitionKey(peer, partition.KeyId, mint: false);

        return _partitions.ContainsKey(partitionKey);
    }

    private static string RequirePeer(ApiKeyThrottlePartition partition)
    {
        if (string.IsNullOrEmpty(partition.TransportPeer))
        {
            throw new ArgumentException("A throttle partition requires a transport peer.", nameof(partition));
        }

        return partition.TransportPeer;
    }

    private static string Composite(string peer, string keyId) => peer + PartitionSeparator + keyId;

    private void RecordInto(WindowState state, long now, int limit)
    {
        lock (state)
        {
            Prune(state, now);
            state.FailureTicks.Enqueue(now);
            state.LastActivityTicks = now;

            // Arm (or push out) the probe slot whenever the state is at or over its limit, so the
            // attempt that trips the limit is not itself followed by an immediate free probe. This
            // is a write of NextProbeAtTicks, so it bumps the version that identifies a probe claim
            // — otherwise a release could mistake this re-arm for its own reservation.
            if (limit > 0 && state.FailureTicks.Count >= limit)
            {
                state.NextProbeAtTicks = now + _probeIntervalTicks;
                state.ProbeVersion++;
            }
        }
    }

    private bool IsOverLimit(WindowState state, long now, int limit)
    {
        if (limit <= 0)
        {
            return false;
        }

        lock (state)
        {
            Prune(state, now);
            return state.FailureTicks.Count >= limit;
        }
    }

    /// <summary>
    /// Atomically claims this interval's probe slot. Checking and advancing must happen in one
    /// critical section: a due-check that released the lock before advancing would let every request
    /// arriving at the interval boundary observe "due" and all be admitted, which is exactly the
    /// unbounded-guessing burst the probe interval exists to prevent.
    /// </summary>
    private bool TryConsumeProbe(WindowState state, long now, out ProbeClaim claim)
    {
        lock (state)
        {
            long previousProbeAtTicks = state.NextProbeAtTicks;
            if (now < previousProbeAtTicks)
            {
                claim = default;
                return false;
            }

            state.NextProbeAtTicks = now + _probeIntervalTicks;
            state.LastActivityTicks = now;
            claim = new ProbeClaim(previousProbeAtTicks, ++state.ProbeVersion);
            return true;
        }
    }

    /// <summary>
    /// Returns a probe slot claimed for a request that a later layer then refused, so the wasted
    /// reservation does not cost the next arrival its slot.
    /// </summary>
    /// <remarks>
    /// The claim is identified by the per-state version stamped when it was made, never by the
    /// timestamp it wrote. Every writer of <c>NextProbeAtTicks</c> bumps that version, so a re-arm
    /// from a concurrent failure — which writes the identical <c>now + interval</c> expression, and
    /// at ~1 ms clock resolution routinely lands on the same tick — cannot be mistaken for the
    /// caller's own reservation and stomped back to a stale, already-due value. Restoring bumps the
    /// version again so no other stale release can match either.
    /// </remarks>
    private static void ReleaseProbe(WindowState state, ProbeClaim claim)
    {
        lock (state)
        {
            if (state.ProbeVersion != claim.Version)
            {
                return;
            }

            state.NextProbeAtTicks = claim.PreviousProbeAtTicks;
            state.ProbeVersion++;
        }
    }

    private void Prune(WindowState state, long now)
    {
        while (state.FailureTicks.Count > 0 && now - state.FailureTicks.Peek() >= _windowTicks)
        {
            state.FailureTicks.Dequeue();
        }
    }

    /// <summary>
    /// Maps a partition onto its storage key, enforcing the per-peer key-id cap. Returns the
    /// transport-peer fallback key (and a null effective key id) when the token carried no key id or
    /// when this address has already minted <see cref="MaxKeyIdPartitionsPerPeer"/> of them.
    /// </summary>
    private (string PartitionKey, string? EffectiveKeyId) ResolvePartitionKey(string peer, string? keyId, bool mint)
    {
        if (keyId is null)
        {
            return (peer, null);
        }

        while (true)
        {
            if (!_peerKeyIds.TryGetValue(peer, out PeerKeyIds? keyIds))
            {
                if (!mint)
                {
                    // Nothing tracked for this address yet, so the composite lookup simply misses.
                    return (Composite(peer, keyId), keyId);
                }

                keyIds = _peerKeyIds.GetOrAdd(peer, static _ => new PeerKeyIds());
            }

            lock (keyIds)
            {
                if (keyIds.Removed)
                {
                    // Lost a race with cleanup; re-read the dictionary and try again.
                    continue;
                }

                if (keyIds.KeyIds.Contains(keyId))
                {
                    return (Composite(peer, keyId), keyId);
                }

                if (keyIds.KeyIds.Count >= MaxKeyIdPartitionsPerPeer)
                {
                    return (peer, null);
                }

                if (mint)
                {
                    keyIds.KeyIds.Add(keyId);
                }

                return (Composite(peer, keyId), keyId);
            }
        }
    }

    private bool RemovePartition(string partitionKey)
    {
        if (!_partitions.TryRemove(partitionKey, out WindowState? state))
        {
            return false;
        }

        if (state.TransportPeer is not null && state.KeyId is not null)
        {
            ReleaseKeyId(state.TransportPeer, state.KeyId);
        }

        return true;
    }

    private void ReleaseKeyId(string peer, string keyId)
    {
        if (!_peerKeyIds.TryGetValue(peer, out PeerKeyIds? keyIds))
        {
            return;
        }

        lock (keyIds)
        {
            keyIds.KeyIds.Remove(keyId);
            if (keyIds.KeyIds.Count == 0)
            {
                keyIds.Removed = true;
                _peerKeyIds.TryRemove(peer, out _);
            }
        }
    }

    /// <summary>
    /// Best-effort eviction, run only when a map exceeds the cap. Preference order: fully expired
    /// windows, then least-recently-active entries that are still under their limit, and only above
    /// the 2x transient-overshoot ceiling the oldest over-limit entry. Racy under concurrency, which
    /// is acceptable for a bound rather than an exact policy.
    /// </summary>
    private void EvictIfOverCapacity(
        ConcurrentDictionary<string, WindowState> map,
        long now,
        int limit,
        bool releaseKeyIds)
    {
        while (map.Count > _maxPartitions)
        {
            bool overHardCeiling = map.Count > (long)_maxPartitions * 2;
            string? expired = null;
            string? leastRecent = null;
            string? oldestOverLimit = null;
            long leastRecentTicks = long.MaxValue;
            long oldestOverLimitTicks = long.MaxValue;

            foreach (KeyValuePair<string, WindowState> entry in map)
            {
                int count;
                long activity;
                lock (entry.Value)
                {
                    Prune(entry.Value, now);
                    count = entry.Value.FailureTicks.Count;
                    activity = entry.Value.LastActivityTicks;
                }

                if (count == 0)
                {
                    expired = entry.Key;
                    break;
                }

                if (limit > 0 && count >= limit)
                {
                    if (activity < oldestOverLimitTicks)
                    {
                        oldestOverLimitTicks = activity;
                        oldestOverLimit = entry.Key;
                    }

                    continue;
                }

                if (activity < leastRecentTicks)
                {
                    leastRecentTicks = activity;
                    leastRecent = entry.Key;
                }
            }

            string? victim = expired ?? leastRecent ?? (overHardCeiling ? oldestOverLimit : null);
            if (victim is null)
            {
                // Everything left is load-bearing and the ceiling is not breached: accept the
                // documented transient overshoot rather than clearing an active block.
                return;
            }

            bool removed = releaseKeyIds ? RemovePartition(victim) : map.TryRemove(victim, out _);
            if (!removed)
            {
                return;
            }
        }
    }

    private sealed class WindowState(string? transportPeer, string? keyId)
    {
        /// <summary>Gets the transport peer this partition belongs to; null for key-id aggregates.</summary>
        public string? TransportPeer { get; } = transportPeer;

        /// <summary>Gets the key id this partition belongs to; null for fallback partitions and aggregates.</summary>
        public string? KeyId { get; } = keyId;

        /// <summary>Timestamps (in ticks) of failures still within the sliding window.</summary>
        public Queue<long> FailureTicks { get; } = new();

        public long LastActivityTicks;

        public long NextProbeAtTicks;

        /// <summary>
        /// Monotonic stamp bumped by every writer of <see cref="NextProbeAtTicks"/> (probe claim and
        /// failure re-arm alike). It is what lets a compensating release recognise its own
        /// reservation without comparing timestamps, which collide whenever two writers share a
        /// clock tick.
        /// </summary>
        public long ProbeVersion;
    }

    /// <summary>
    /// A partition's resolved storage key, carried from the check to the reset of the same request.
    /// The composite key is a fresh string per build, so resolving once per RPC rather than once per
    /// call keeps the successful auth path (check, then reset) to a single allocation.
    /// </summary>
    /// <param name="PartitionKey">The storage key, or <see langword="null"/> when unresolved.</param>
    /// <param name="EffectiveKeyId">
    /// The key id that actually earned a partition, or <see langword="null"/> when the token carried
    /// none or the per-peer cap collapsed it onto the transport-peer fallback.
    /// </param>
    internal readonly record struct PartitionResolution(string? PartitionKey, string? EffectiveKeyId)
    {
        /// <summary>Gets the sentinel for "not resolved yet"; the receiving call resolves it itself.</summary>
        internal static PartitionResolution Unresolved => default;

        /// <summary>Gets a value indicating whether this carries a resolved storage key.</summary>
        internal bool IsResolved => PartitionKey is not null;
    }

    /// <summary>A probe slot reservation: what to restore, and the stamp proving it is still ours.</summary>
    /// <param name="PreviousProbeAtTicks">The slot value replaced when the claim was made.</param>
    /// <param name="Version">The <see cref="WindowState.ProbeVersion"/> stamped by this claim.</param>
    private readonly record struct ProbeClaim(long PreviousProbeAtTicks, long Version);

    private sealed class PeerKeyIds
    {
        /// <summary>Key ids this transport peer has minted a partition for.</summary>
        public HashSet<string> KeyIds { get; } = new(StringComparer.Ordinal);

        /// <summary>Set once the entry has been dropped from the peer map, so racing minters retry.</summary>
        public bool Removed { get; set; }
    }
}
