using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;

namespace StellaOps.Scanner.Core.Security;

public sealed class InMemoryDpopReplayCache : IDpopReplayCache
{
    private readonly ConcurrentDictionary<string, DateTimeOffset> entries = new(StringComparer.Ordinal);
    private readonly TimeProvider timeProvider;

    public InMemoryDpopReplayCache(TimeProvider? timeProvider = null)
    {
        this.timeProvider = timeProvider ?? TimeProvider.System;
    }

    public ValueTask<bool> TryStoreAsync(string jwtId, DateTimeOffset expiresAt, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jwtId);

        var now = timeProvider.GetUtcNow();
        RemoveExpired(now);

        if (entries.TryAdd(jwtId, expiresAt))
        {
            return ValueTask.FromResult(true);
        }

        while (!cancellationToken.IsCancellationRequested)
        {
            if (!entries.TryGetValue(jwtId, out var existing))
            {
                if (entries.TryAdd(jwtId, expiresAt))
                {
                    return ValueTask.FromResult(true);
                }

                continue;
            }

            if (existing > now)
            {
                return ValueTask.FromResult(false);
            }

            if (entries.TryUpdate(jwtId, expiresAt, existing))
            {
                return ValueTask.FromResult(true);
            }
        }

        return ValueTask.FromResult(false);
    }

    private void RemoveExpired(DateTimeOffset now)
    {
        foreach (var entry in entries)
        {
            if (entry.Value <= now)
            {
                entries.TryRemove(entry.Key, out _);
            }
        }
    }
}
