// 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;
using System.Text;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
 
namespace Microsoft.AspNetCore.HttpLogging;
 
internal sealed class ResponseBufferingStream : BufferingStream, IHttpResponseBodyFeature
{
    private IHttpResponseBodyFeature _innerBodyFeature = null!;
    private int _limit;
    private PipeWriter? _pipeAdapter;
 
    private HttpLoggingInterceptorContext _logContext = null!;
    private HttpLoggingOptions _options = null!;
    private IHttpLoggingInterceptor[] _interceptors = [];
    private bool _logBody;
    private bool _hasLogged;
    private Encoding? _encoding;
    private string? _bodyBeforeClose;
 
    private static readonly StreamPipeWriterOptions _pipeWriterOptions = new StreamPipeWriterOptions(leaveOpen: true);
 
    /// <summary>
    /// Parameterless constructor for object pooling.
    /// </summary>
    internal ResponseBufferingStream()
        : base(Stream.Null, NullLogger.Instance)
    {
    }
 
    /// <summary>
    /// Initializes the stream for a new request. Call this after obtaining from the pool.
    /// </summary>
    internal void Initialize(
        IHttpResponseBodyFeature innerBodyFeature,
        ILogger logger,
        HttpLoggingInterceptorContext logContext,
        HttpLoggingOptions options,
        IHttpLoggingInterceptor[] interceptors)
    {
        _innerBodyFeature = innerBodyFeature;
        _innerStream = innerBodyFeature.Stream;
        _logger = logger;
        _logContext = logContext;
        _options = options;
        _interceptors = interceptors;
        // Reset transient state
        _limit = 0;
        _pipeAdapter = null;
        _logBody = false;
        _hasLogged = false;
        _encoding = null;
        _bodyBeforeClose = null;
        HeadersWritten = false;
    }
 
    /// <summary>
    /// Resets the stream state for returning to the pool.
    /// </summary>
    internal void ResetForPool()
    {
        _innerBodyFeature = null!;
        _innerStream = Stream.Null;
        _logger = NullLogger.Instance;
        _logContext = null!;
        _options = null!;
        _interceptors = [];
        _pipeAdapter = null;
        _encoding = null;
        _bodyBeforeClose = null;
        _limit = 0;
        _logBody = false;
        _hasLogged = false;
        HeadersWritten = false;
        // Reset base class buffer state
        Reset();
    }
 
    public bool HeadersWritten { get; private set; }
 
    public Stream Stream => this;
 
    public PipeWriter Writer => _pipeAdapter ??= PipeWriter.Create(Stream, _pipeWriterOptions);
 
    public override void Write(byte[] buffer, int offset, int count)
    {
        Write(buffer.AsSpan(offset, count));
    }
 
    public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
    {
        return TaskToApm.Begin(WriteAsync(buffer, offset, count), callback, state);
    }
 
    public override void EndWrite(IAsyncResult asyncResult)
    {
        TaskToApm.End(asyncResult);
    }
 
    public override void Write(ReadOnlySpan<byte> span)
    {
        OnFirstWriteSync();
        CommonWrite(span);
 
        _innerStream.Write(span);
    }
 
    public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
    {
        await WriteAsync(new Memory<byte>(buffer, offset, count), cancellationToken);
    }
 
    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
    {
        await OnFirstWriteAsync();
        CommonWrite(buffer.Span);
 
        await _innerStream.WriteAsync(buffer, cancellationToken);
    }
 
    private void CommonWrite(ReadOnlySpan<byte> span)
    {
        var remaining = _limit - _bytesBuffered;
        var innerCount = Math.Min(remaining, span.Length);
 
        if (_logBody && innerCount > 0)
        {
            var slice = span.Slice(0, innerCount);
            if (slice.TryCopyTo(_tailMemory.Span))
            {
                _tailBytesBuffered += innerCount;
                _bytesBuffered += innerCount;
                _tailMemory = _tailMemory.Slice(innerCount);
            }
            else
            {
                BuffersExtensions.Write(this, slice);
            }
        }
    }
 
    private void OnFirstWriteSync()
    {
        if (!HeadersWritten)
        {
            // Log headers as first write occurs (headers locked now)
            HttpLoggingMiddleware.LogResponseHeadersSync(_logContext, _options, _interceptors, _logger);
            OnFirstWriteCore();
        }
    }
 
    private async ValueTask OnFirstWriteAsync()
    {
        if (!HeadersWritten)
        {
            // Log headers as first write occurs (headers locked now)
            await HttpLoggingMiddleware.LogResponseHeadersAsync(_logContext, _options, _interceptors, _logger);
            OnFirstWriteCore();
        }
    }
 
    private void OnFirstWriteCore()
    {
        // The callback in LogResponseHeaders could disable body logging or adjust limits.
        if (_logContext.LoggingFields.HasFlag(HttpLoggingFields.ResponseBody) && _logContext.ResponseBodyLogLimit > 0)
        {
            if (MediaTypeHelpers.TryGetEncodingForMediaType(_logContext.HttpContext.Response.ContentType,
                _options.MediaTypeOptions.MediaTypeStates, out _encoding))
            {
                _logBody = true;
                _limit = _logContext.ResponseBodyLogLimit;
            }
            else
            {
                _logger.UnrecognizedMediaType("response");
            }
        }
 
        HeadersWritten = true;
    }
 
    public void DisableBuffering()
    {
        _innerBodyFeature.DisableBuffering();
    }
 
    public async Task SendFileAsync(string path, long offset, long? count, CancellationToken cancellation)
    {
        await OnFirstWriteAsync();
        await _innerBodyFeature.SendFileAsync(path, offset, count, cancellation);
    }
 
    public async Task StartAsync(CancellationToken token = default)
    {
        await OnFirstWriteAsync();
        await _innerBodyFeature.StartAsync(token);
    }
 
    public async Task CompleteAsync()
    {
        await OnFirstWriteAsync();
        await _innerBodyFeature.CompleteAsync();
    }
 
    public override void Flush()
    {
        OnFirstWriteSync();
        base.Flush();
    }
 
    public override async Task FlushAsync(CancellationToken cancellationToken)
    {
        await OnFirstWriteAsync();
        await base.FlushAsync(cancellationToken);
    }
 
    public void LogResponseBody()
    {
        if (_logBody)
        {
            var responseBody = GetStringInternal();
            _logger.ResponseBody(responseBody);
            _hasLogged = true;
        }
    }
 
    public void LogResponseBody(HttpLoggingInterceptorContext logContext)
    {
        if (_logBody)
        {
            logContext.AddParameter("ResponseBody", GetStringInternal());
            _hasLogged = true;
        }
    }
 
    private string GetStringInternal()
    {
        var result = _bodyBeforeClose ?? GetString(_encoding!);
        // Reset the value after its consumption to preserve GetString(encoding) behavior
        _bodyBeforeClose = null;
        return result;
    }
 
    public override void Close()
    {
        if (_logBody && !_hasLogged)
        {
            // Subsequent middleware can close the response stream after writing its body
            // Preserving the body for the final GetStringInternal() call.
            _bodyBeforeClose = GetString(_encoding!);
        }
 
        base.Close();
    }
}