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

namespace FinlyticFundamentals.Services;

public class KeyedLockPool
{
    private readonly ConcurrentDictionary<string, RefCountedSemaphore> _semaphores = new(StringComparer.OrdinalIgnoreCase);

    public async Task<IDisposable> LockAsync(string key, CancellationToken cancellationToken = default)
    {
        var item = _semaphores.AddOrUpdate(
            key,
            _ => new RefCountedSemaphore(),
            (_, existing) =>
            {
                Interlocked.Increment(ref existing.RefCount);
                return existing;
            });

        await item.Semaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
        return new Releaser(this, key, item);
    }

    private void Release(string key, RefCountedSemaphore item)
    {
        item.Semaphore.Release();
        if (Interlocked.Decrement(ref item.RefCount) <= 0)
        {
            _semaphores.TryRemove(new KeyValuePair<string, RefCountedSemaphore>(key, item));
        }
    }

    private sealed class RefCountedSemaphore
    {
        public int RefCount = 1;
        public readonly SemaphoreSlim Semaphore = new(1, 1);
    }

    private sealed class Releaser : IDisposable
    {
        private readonly KeyedLockPool _pool;
        private readonly string _key;
        private readonly RefCountedSemaphore _item;
        private bool _disposed;

        public Releaser(KeyedLockPool pool, string key, RefCountedSemaphore item)
        {
            _pool = pool;
            _key = key;
            _item = item;
        }

        public void Dispose()
        {
            if (!_disposed)
            {
                _disposed = true;
                _pool.Release(_key, _item);
            }
        }
    }
}