using AcDream.App.Rendering;

namespace AcDream.App.Rendering.Wb;

/// <summary>
/// A contiguous GPU-buffer suballocator whose released ranges become
/// reusable only after the frame-retirement queue says that older draws have
/// finished. This prevents <c>BufferSubData</c> from overwriting a range that
/// an in-flight draw still reads.
/// </summary>
internal sealed class GpuRetiredRangeAllocator
{
    private readonly ContiguousRangeAllocator _allocator;
    private readonly GpuRetirementLedger _retirementLedger;
    private readonly Dictionary<MeshBufferRange, RetryableGpuResourceRelease> _pendingReleases = [];
    private int _pendingReleaseCount;
    private int _pendingReleaseLength;

    public GpuRetiredRangeAllocator(int capacity, IGpuResourceRetirementQueue retirement)
    {
        _allocator = new ContiguousRangeAllocator(capacity);
        _retirementLedger = new GpuRetirementLedger(
            retirement ?? throw new ArgumentNullException(nameof(retirement)));
    }

    public int Capacity => _allocator.Capacity;
    public int Used => _allocator.Used;
    public int HighWaterMark => _allocator.HighWaterMark;
    public int LargestFreeRange => _allocator.LargestFreeRange;
    public int TrailingFreeLength => _allocator.TrailingFreeLength;
    public int PendingReleaseCount => _pendingReleaseCount;
    public int PendingReleaseLength => _pendingReleaseLength;

    public bool TryAllocate(int length, out MeshBufferRange allocation) =>
        _allocator.TryAllocate(length, out allocation);

    public void Grow(int newCapacity) => _allocator.Grow(newCapacity);

    public void Shrink(int newCapacity) => _allocator.Shrink(newCapacity);

    public void ReleaseAfterGpuUse(MeshBufferRange allocation)
    {
        if (_pendingReleases.TryGetValue(allocation, out RetryableGpuResourceRelease? pending))
        {
            // The caller retries this method when queue publication failed.
            // Re-publish the retained transaction instead of accounting for a
            // second logical release of the same physical range.
            try
            {
                _retirementLedger.RetryPendingPublication(pending);
            }
            catch when (pending.IsComplete)
            {
                // An immediate retirement queue may surface a wrapper failure
                // after the retained transaction itself fully committed.
            }
            return;
        }

        int nextPendingCount = checked(_pendingReleaseCount + 1);
        int nextPendingLength = checked(_pendingReleaseLength + allocation.Length);
        var release = new RetryableGpuResourceRelease(
            () => _allocator.Release(allocation),
            () => _pendingReleaseCount = checked(_pendingReleaseCount - 1),
            () => _pendingReleaseLength = checked(
                _pendingReleaseLength - allocation.Length),
            () =>
            {
                if (!_pendingReleases.Remove(allocation))
                {
                    throw new InvalidOperationException(
                        "GPU buffer range retirement lost its ownership record.");
                }
            });
        _pendingReleases.Add(allocation, release);
        _pendingReleaseCount = nextPendingCount;
        _pendingReleaseLength = nextPendingLength;

        try
        {
            _retirementLedger.Retire(release);
        }
        catch when (release.IsComplete)
        {
            // Completion is stronger than publication acknowledgement. The
            // physical range and its accounting already converged.
        }
    }

    /// <summary>
    /// Releases a range that failed before it could be submitted in a draw.
    /// </summary>
    public void ReleaseUnsubmitted(MeshBufferRange allocation)
    {
        if (_pendingReleases.ContainsKey(allocation))
        {
            throw new InvalidOperationException(
                "A GPU-submitted range cannot be released as unsubmitted.");
        }

        _allocator.Release(allocation);
    }
}