using System.Collections.Concurrent;

namespace DysonNetwork.Pass.Storage;

public interface ICacheService
{
    Task<T?> GetAsync<T>(string key);
    Task<(bool found, T? value)> GetAsyncWithStatus<T>(string key);
    Task SetAsync<T>(string key, T value, TimeSpan? expiry = null);
    Task SetWithGroupsAsync<T>(string key, T value, string[] groups, TimeSpan? expiry = null);
    Task RemoveAsync(string key);
    Task RemoveGroupAsync(string groupName);
    Task<IAsyncDisposable?> AcquireLockAsync(string key, TimeSpan expiry, TimeSpan? wait = null);
}

public class InMemoryCacheService : ICacheService
{
    private readonly ConcurrentDictionary<string, object> _cache = new();
    private readonly ConcurrentDictionary<string, HashSet<string>> _groups = new();
    private readonly ConcurrentDictionary<string, DateTimeOffset> _expirations = new();

    public Task<T?> GetAsync<T>(string key)
    {
        if (_cache.TryGetValue(key, out var value) && !IsExpired(key))
        {
            return Task.FromResult((T?)value);
        }
        Remove(key);
        return Task.FromResult(default(T));
    }

    public Task<(bool found, T? value)> GetAsyncWithStatus<T>(string key)
    {
        if (_cache.TryGetValue(key, out var value) && !IsExpired(key))
        {
            return Task.FromResult((true, (T?)value));
        }
        Remove(key);
        return Task.FromResult((false, default(T)));
    }

    public Task SetAsync<T>(string key, T value, TimeSpan? expiry = null)
    {
        _cache[key] = value!;
        _expirations[key] = expiry.HasValue ? DateTimeOffset.UtcNow.Add(expiry.Value) : DateTimeOffset.MaxValue;
        return Task.CompletedTask;
    }

    public Task SetWithGroupsAsync<T>(string key, T value, string[] groups, TimeSpan? expiry = null)
    {
        _cache[key] = value!;
        _expirations[key] = expiry.HasValue ? DateTimeOffset.UtcNow.Add(expiry.Value) : DateTimeOffset.MaxValue;
        foreach (var group in groups)
        {
            _groups.GetOrAdd(group, _ => new HashSet<string>()).Add(key);
        }
        return Task.CompletedTask;
    }

    public Task RemoveAsync(string key)
    {
        Remove(key);
        return Task.CompletedTask;
    }

    public Task RemoveGroupAsync(string groupName)
    {
        if (_groups.TryRemove(groupName, out var keysToRemove))
        {
            foreach (var key in keysToRemove)
            {
                Remove(key);
            }
        }
        return Task.CompletedTask;
    }

    public async Task<IAsyncDisposable?> AcquireLockAsync(string key, TimeSpan expiry, TimeSpan? wait = null)
    {
        var startTime = DateTime.UtcNow;
        while (true)
        {
            if (_cache.TryAdd(key, new object()))
            {
                _expirations[key] = DateTimeOffset.UtcNow.Add(expiry);
                return new AsyncLockReleaser(() => Remove(key));
            }

            if (wait == null || DateTime.UtcNow - startTime > wait.Value)
            {
                return null; // Could not acquire lock within wait time
            }
            await Task.Delay(50); // Wait a bit before retrying
        }
    }

    private bool IsExpired(string key)
    {
        return _expirations.TryGetValue(key, out var expiration) && expiration <= DateTimeOffset.UtcNow;
    }

    private void Remove(string key)
    {
        _cache.TryRemove(key, out _);
        _expirations.TryRemove(key, out _);
        foreach (var group in _groups.Values)
        {
            group.Remove(key);
        }
    }

    private class AsyncLockReleaser : IAsyncDisposable
    {
        private readonly Action _releaseAction;

        public AsyncLockReleaser(Action releaseAction)
        {
            _releaseAction = releaseAction;
        }

        public ValueTask DisposeAsync()
        {
            _releaseAction();
            return ValueTask.CompletedTask;
        }
    }
}