namespace ZymonicServices;

/// <summary>
/// Shares disposable resources, creates them on one background worker at a time,
/// and retains unused resources within count and optional weight budgets in LRU order.
/// Active leases keep their resource alive even after the cache is disposed.
/// </summary>
public sealed class ZymonicResourceCache<TKey, TResource> : IDisposable
    where TKey : notnull
    where TResource : class, IDisposable
{
    private readonly object _gate = new object();
    private readonly Dictionary<TKey, Entry> _entries = new();
    private readonly LinkedList<Entry> _unused = new();
    private readonly SemaphoreSlim _creationGate = new(1, 1);
    private readonly Func<TKey, TResource>? _create;
    private readonly Func<TResource, long> _getWeight;
    private readonly long _maxUnusedWeight;
    private long _unusedWeight;
    private readonly int _unusedCapacity;
    private bool _disposed;

    private sealed class Entry(TKey key)
    {
        public TKey Key { get; } = key;
        public Func<TResource>? Factory;
        public long Weight;
        public int References;
        public bool Removed;
        public TResource? Resource;
        public LinkedListNode<Entry>? UnusedNode;
        public TaskCompletionSource<TResource> Completion { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
    }

    public sealed class Lease : IDisposable
    {
        private Action? _release;
        public TResource Resource { get; }

        internal Lease(TResource resource, Action release)
        {
            Resource = resource;
            _release = release;
        }

        public void Dispose() => Interlocked.Exchange(ref _release, null)?.Invoke();
    }

    public ZymonicResourceCache(Func<TKey, TResource> create, int unusedCapacity = 32,
        Func<TResource, long>? getWeight = null, long maxUnusedWeight = long.MaxValue)
        : this(unusedCapacity, getWeight, maxUnusedWeight)
    {
        ArgumentNullException.ThrowIfNull(create);
        _create = create;
    }

    public ZymonicResourceCache(int unusedCapacity = 32,
        Func<TResource, long>? getWeight = null, long maxUnusedWeight = long.MaxValue)
    {
        ArgumentOutOfRangeException.ThrowIfNegative(unusedCapacity);
        ArgumentOutOfRangeException.ThrowIfNegative(maxUnusedWeight);
        _unusedCapacity = unusedCapacity;
        _getWeight = getWeight ?? (_ => 0);
        _maxUnusedWeight = maxUnusedWeight;
    }

    public Task<Lease> AcquireAsync(TKey key, CancellationToken cancellationToken = default)
    {
        if (_create is null) throw new InvalidOperationException("A creation factory must be supplied.");
        return AcquireAsync(key, () => _create(key), cancellationToken);
    }

    // The first factory for a key wins. Captured source data is released after
    // creation or abandonment, independently of the cached resource's lifetime.
    public async Task<Lease> AcquireAsync(TKey key, Func<TResource> create, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(create);
        cancellationToken.ThrowIfCancellationRequested();
        Entry entry;
        bool startCreation = false;
        lock (_gate)
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
            if (_entries.TryGetValue(key, out var existing))
            {
                entry = existing;
                RemoveUnusedNode(entry);
            }
            else
            {
                entry = new Entry(key) { Factory = create };
                _entries.Add(key, entry);
                startCreation = true;
            }
            entry.References++;
        }

        create = null!; // Do not retain an unused factory in an awaiting state machine.
        if (startCreation)
        {
            _ = Task.Run(() => CreateAsync(entry));
        }

        try
        {
            var resource = await entry.Completion.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
            cancellationToken.ThrowIfCancellationRequested();
            lock (_gate)
            {
                ObjectDisposedException.ThrowIf(_disposed, this);
            }
            return new Lease(resource, () => Release(entry));
        }
        catch
        {
            Release(entry);
            throw;
        }
    }

    private async Task CreateAsync(Entry entry)
    {
        await _creationGate.WaitAsync().ConfigureAwait(false);
        try
        {
            Func<TResource>? factory;
            lock (_gate)
            {
                factory = !_disposed && !entry.Removed && entry.References > 0 ? entry.Factory : null;
                entry.Factory = null;
            }
            if (factory is not null)
            {
                var resource = factory();
                factory = null;
                long weight;
                try
                {
                    weight = _getWeight(resource);
                    ArgumentOutOfRangeException.ThrowIfNegative(weight);
                }
                catch { resource.Dispose(); throw; }
                List<TResource> evicted;
                lock (_gate)
                {
                    entry.Resource = resource;
                    entry.Weight = weight;
                    evicted = RetainOrEvict(entry);
                    entry.Completion.TrySetResult(resource);
                }
                DisposeResources(evicted);
            }
            else
            {
                entry.Completion.TrySetCanceled();
            }
        }
        catch (Exception error)
        {
            lock (_gate)
            {
                RemoveEntry(entry);
            }
            entry.Completion.TrySetException(error);
            // All waiters may already have cancelled. Observe the failure in that case too.
            _ = entry.Completion.Task.Exception;
        }
        finally
        {
            _creationGate.Release();
        }
    }

    private void Release(Entry entry)
    {
        List<TResource> evicted;
        lock (_gate)
        {
            entry.References--;
            if (entry.References == 0 && entry.Resource is null)
            {
                // Queued work with no remaining consumers should not decode anything.
                RemoveEntry(entry);
            }
            evicted = RetainOrEvict(entry);
        }
        DisposeResources(evicted);
    }

    // Called under _gate; disposal itself happens outside the lock.
    private List<TResource> RetainOrEvict(Entry entry)
    {
        List<TResource> evicted = new();
        if (entry.References == 0 && entry.Resource is not null)
        {
            if (!_disposed && !entry.Removed && _unusedCapacity > 0 && entry.Weight <= _maxUnusedWeight)
            {
                // Evict before adding to avoid overflow in weight accounting.
                while (_unused.Count > 0 && (_unused.Count >= _unusedCapacity
                    || entry.Weight > _maxUnusedWeight - _unusedWeight))
                {
                    var oldest = _unused.First!.Value;
                    RemoveEntry(oldest);
                    evicted.Add(oldest.Resource!);
                    oldest.Resource = null;
                }
                entry.UnusedNode = _unused.AddLast(entry);
                _unusedWeight += entry.Weight;

            }
            else
            {
                RemoveEntry(entry);
                evicted.Add(entry.Resource);
                entry.Resource = null;
            }
        }
        return evicted;
    }

    private void RemoveUnusedNode(Entry entry)
    {
        if (entry.UnusedNode is not null)
        {
            _unusedWeight -= entry.Weight;
            _unused.Remove(entry.UnusedNode);
            entry.UnusedNode = null;
        }
    }

    private void RemoveEntry(Entry entry)
    {
        if (!entry.Removed)
        {
            entry.Factory = null;
            _entries.Remove(entry.Key);
            RemoveUnusedNode(entry);
            entry.Removed = true;
        }
    }

    private static void DisposeResources(IEnumerable<TResource> resources)
    {
        foreach (var resource in resources)
        {
            resource.Dispose();
        }
    }

    public void Dispose()
    {
        List<TResource> evicted = new();
        lock (_gate)
        {
            if (_disposed) return;
            _disposed = true;
            foreach (var entry in _entries.Values)
            {
                entry.Factory = null;
                entry.Removed = true;
                entry.UnusedNode = null;
                if (entry.References == 0 && entry.Resource is not null)
                {
                    evicted.Add(entry.Resource);
                    entry.Resource = null;
                }
            }
            _entries.Clear();
            _unused.Clear();
            _unusedWeight = 0;
        }
        DisposeResources(evicted);
        // Pending workers still use _creationGate; leave its managed lifetime to the GC.
    }
}