// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.IO.Pipelines;
using System.Threading;
using System.Threading.Tasks;

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()
        {
            if (Disposed)
            {
                throw new ObjectDisposedException(nameof(PagedByteBuffer));
            }
        }
    }
}
