// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Buffers;
using System.IO.Pipelines;
 
namespace Microsoft.AspNetCore.WebUtilities;
 
internal sealed class PagedByteBuffer : IDisposable
{
    internal const int PageSize = 1024;
    private readonly ArrayPool<byte> _arrayPool;
    private byte[]? _currentPage;
    private int _currentPageIndex;
 
    public PagedByteBuffer(ArrayPool<byte> arrayPool)
    {
        _arrayPool = arrayPool;
        Pages = new List<byte[]>();
    }
 
    public int Length { get; private set; }
 
    internal bool Disposed { get; private set; }
 
    internal List<byte[]> Pages { get; }
 
    private byte[] CurrentPage
    {
        get
        {
            if (_currentPage == null || _currentPageIndex == _currentPage.Length)
            {
                _currentPage = _arrayPool.Rent(PageSize);
                Pages.Add(_currentPage);
                _currentPageIndex = 0;
            }
 
            return _currentPage;
        }
    }
 
    public void Add(byte[] buffer, int offset, int count)
        => Add(buffer.AsMemory(offset, count));
 
    public void Add(ReadOnlyMemory<byte> memory)
    {
        ThrowIfDisposed();
 
        while (!memory.IsEmpty)
        {
            var currentPage = CurrentPage;
            var copyLength = Math.Min(memory.Length, currentPage.Length - _currentPageIndex);
 
            memory.Slice(0, copyLength).CopyTo(currentPage.AsMemory(_currentPageIndex, copyLength));
 
            Length += copyLength;
            _currentPageIndex += copyLength;
 
            memory = memory.Slice(copyLength);
        }
    }
 
    public void MoveTo(Stream stream)
    {
        ThrowIfDisposed();
 
        for (var i = 0; i < Pages.Count; i++)
        {
            var page = Pages[i];
            var length = (i == Pages.Count - 1) ?
                _currentPageIndex :
                page.Length;
 
            stream.Write(page, 0, length);
        }
 
        ClearBuffers();
    }
 
    public async Task MoveToAsync(PipeWriter writer, CancellationToken cancellationToken)
    {
        ThrowIfDisposed();
 
        for (var i = 0; i < Pages.Count; i++)
        {
            var page = Pages[i];
            var length = (i == Pages.Count - 1) ?
                _currentPageIndex :
                page.Length;
 
            await writer.WriteAsync(page.AsMemory(0, length), cancellationToken);
        }
 
        ClearBuffers();
    }
 
    public async Task MoveToAsync(Stream stream, CancellationToken cancellationToken)
    {
        ThrowIfDisposed();
 
        for (var i = 0; i < Pages.Count; i++)
        {
            var page = Pages[i];
            var length = (i == Pages.Count - 1) ?
                _currentPageIndex :
                page.Length;
 
            await stream.WriteAsync(page.AsMemory(0, length), cancellationToken);
        }
 
        ClearBuffers();
    }
 
    public void Dispose()
    {
        if (!Disposed)
        {
            Disposed = true;
            ClearBuffers();
        }
    }
 
    private void ClearBuffers()
    {
        for (var i = 0; i < Pages.Count; i++)
        {
            _arrayPool.Return(Pages[i]);
        }
 
        Pages.Clear();
        _currentPage = null;
        Length = 0;
        _currentPageIndex = 0;
    }
 
    private void ThrowIfDisposed()
    {
        ObjectDisposedException.ThrowIf(Disposed, this);
    }
}