// 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.IO;
using System.Threading;
using System.Threading.Tasks;

namespace Microsoft.AspNetCore.WebUtilities
{
    /// <summary>
    /// HTTP extension methods for <see cref="Stream"/>.
    /// </summary>
    public static class StreamHelperExtensions
    {
        private const int _maxReadBufferSize = 1024 * 4;

        /// <summary>
        /// Reads the specified <paramref name="stream"/> to the end.
        /// <para>
        /// This API is effective when used in conjunction with buffering. It allows
        /// a buffered request stream to be synchronously read after it has been completely drained.
        /// </para>
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to completely read.</param>
        /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
        public static Task DrainAsync(this Stream stream, CancellationToken cancellationToken)
        {
            return stream.DrainAsync(ArrayPool<byte>.Shared, null, cancellationToken);
        }

        /// <summary>
        /// Reads the specified <paramref name="stream"/> to the end.
        /// <para>
        /// This API is effective when used in conjunction with buffering. It allows
        /// a buffered request stream to be synchronously read after it has been completely drained.
        /// </para>
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to completely read.</param>
        /// <param name="limit">The maximum number of bytes to read. Throws if the <see cref="Stream"/> is larger than this limit.</param>
        /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
        public static Task DrainAsync(this Stream stream, long? limit, CancellationToken cancellationToken)
        {
            return stream.DrainAsync(ArrayPool<byte>.Shared, limit, cancellationToken);
        }

        /// <summary>
        /// Reads the specified <paramref name="stream"/> to the end.
        /// <para>
        /// This API is effective when used in conjunction with buffering. It allows
        /// a buffered request stream to be synchronously read after it has been completely drained.
        /// </para>
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to completely read.</param>
        /// <param name="bytePool">The byte array pool to use.</param>
        /// <param name="limit">The maximum number of bytes to read. Throws if the <see cref="Stream"/> is larger than this limit.</param>
        /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
        public static async Task DrainAsync(this Stream stream, ArrayPool<byte> bytePool, long? limit, CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var buffer = bytePool.Rent(_maxReadBufferSize);
            long total = 0;
            try
            {
                var read = await stream.ReadAsync(buffer.AsMemory(), cancellationToken);
                while (read > 0)
                {
                    // Not all streams support cancellation directly.
                    cancellationToken.ThrowIfCancellationRequested();
                    if (limit.HasValue && limit.GetValueOrDefault() - total < read)
                    {
                        throw new InvalidDataException($"The stream exceeded the data limit {limit.GetValueOrDefault()}.");
                    }
                    total += read;
                    read = await stream.ReadAsync(buffer.AsMemory(), cancellationToken);
                }
            }
            finally
            {
                bytePool.Return(buffer);
            }
        }
    }
}
