AyCode.Core/AyCode.Core/Compression/BrotliHelper.cs

using System.Buffers;
using System.IO.Compression;
using System.Runtime.CompilerServices;
using System.Text;

namespace AyCode.Core.Compression;

/// <summary>
/// Brotli compression/decompression helper for SignalR message transport.
/// Used when JSON serializer is configured to reduce payload size.
/// Optimized for zero-allocation scenarios with pooled buffers.
/// </summary>
public static class BrotliHelper
{
    private const int DefaultBufferSize = 4096;
    private const int MaxStackAllocSize = 1024;

    /// <summary>
    /// Compresses a string using Brotli compression with pooled buffers.
    /// </summary>
    public static byte[] Compress(string text, CompressionLevel compressionLevel = CompressionLevel.Optimal)
    {
        if (string.IsNullOrEmpty(text))
            return [];

        // Use stack allocation for small strings, pooled buffer for larger
        var maxByteCount = Encoding.UTF8.GetMaxByteCount(text.Length);
        
        if (maxByteCount <= MaxStackAllocSize)
        {
            Span<byte> utf8Bytes = stackalloc byte[maxByteCount];
            var actualLength = Encoding.UTF8.GetBytes(text.AsSpan(), utf8Bytes);
            return CompressSpan(utf8Bytes[..actualLength], compressionLevel);
        }
        
        var rentedBuffer = ArrayPool<byte>.Shared.Rent(maxByteCount);
        try
        {
            var actualLength = Encoding.UTF8.GetBytes(text.AsSpan(), rentedBuffer);
            return CompressSpan(rentedBuffer.AsSpan(0, actualLength), compressionLevel);
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(rentedBuffer);
        }
    }

    /// <summary>
    /// Compresses a byte array using Brotli compression.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static byte[] Compress(byte[] data, CompressionLevel compressionLevel = CompressionLevel.Optimal)
        => data == null || data.Length == 0 ? [] : CompressSpan(data.AsSpan(), compressionLevel);

    /// <summary>
    /// Compresses a ReadOnlySpan using Brotli compression with pooled output buffer.
    /// </summary>
    public static byte[] CompressSpan(ReadOnlySpan<byte> data, CompressionLevel compressionLevel = CompressionLevel.Optimal)
    {
        if (data.IsEmpty)
            return [];

        // Estimate compressed size (typically 10-30% of original for text)
        var estimatedSize = Math.Max(data.Length / 2, 64);
        var outputBuffer = ArrayPool<byte>.Shared.Rent(estimatedSize);
        
        try
        {
            using var outputStream = new PooledMemoryStream(outputBuffer);
            using (var brotliStream = new BrotliStream(outputStream, compressionLevel, leaveOpen: true))
            {
                brotliStream.Write(data);
            }
            return outputStream.ToArray();
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(outputBuffer);
        }
    }

    /// <summary>
    /// Compresses data directly to an IBufferWriter (zero intermediate allocation).
    /// </summary>
    public static void CompressTo(ReadOnlySpan<byte> data, IBufferWriter<byte> writer, CompressionLevel compressionLevel = CompressionLevel.Optimal)
    {
        if (data.IsEmpty)
            return;

        using var outputStream = new BufferWriterStream(writer);
        using var brotliStream = new BrotliStream(outputStream, compressionLevel, leaveOpen: true);
        brotliStream.Write(data);
    }

    /// <summary>
    /// Decompresses Brotli-compressed data to a string.
    /// Consider using Decompress + direct UTF-8 JsonTo for better performance.
    /// </summary>
    public static string DecompressToString(byte[] compressedData)
    {
        if (compressedData == null || compressedData.Length == 0)
            return string.Empty;

        var decompressedBytes = Decompress(compressedData);
        return Encoding.UTF8.GetString(decompressedBytes);
    }

    /// <summary>
    /// Decompresses Brotli-compressed data to a byte array.
    /// Uses pooled buffers internally for reduced GC pressure.
    /// </summary>
    public static byte[] Decompress(byte[] compressedData)
    {
        if (compressedData == null || compressedData.Length == 0)
            return [];

        return DecompressSpan(compressedData.AsSpan());
    }

    /// <summary>
    /// Decompresses Brotli-compressed data from a ReadOnlySpan.
    /// </summary>
    public static byte[] DecompressSpan(ReadOnlySpan<byte> compressedData)
    {
        if (compressedData.IsEmpty)
            return [];

        // Estimate decompressed size (typically 3-10x compressed for text)
        var estimatedSize = Math.Max(compressedData.Length * 4, DefaultBufferSize);
        var outputBuffer = ArrayPool<byte>.Shared.Rent(estimatedSize);
        
        try
        {
            using var inputStream = new ReadOnlySpanStream(compressedData);
            using var brotliStream = new BrotliStream(inputStream, CompressionMode.Decompress);
            using var outputStream = new PooledMemoryStream(outputBuffer);
            
            brotliStream.CopyTo(outputStream);
            return outputStream.ToArray();
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(outputBuffer);
        }
    }

    /// <summary>
    /// Decompresses directly into an ArrayBufferWriter for zero-copy deserialization.
    /// </summary>
    public static void DecompressTo(ReadOnlySpan<byte> compressedData, ArrayBufferWriter<byte> writer)
    {
        if (compressedData.IsEmpty)
            return;

        using var inputStream = new ReadOnlySpanStream(compressedData);
        using var brotliStream = new BrotliStream(inputStream, CompressionMode.Decompress);
        
        // Read in chunks directly to the writer
        int bytesRead;
        do
        {
            var buffer = writer.GetSpan(DefaultBufferSize);
            bytesRead = brotliStream.Read(buffer);
            if (bytesRead > 0)
                writer.Advance(bytesRead);
        } while (bytesRead > 0);
    }

    /// <summary>
    /// Decompresses to a rented buffer. Caller must return the buffer to ArrayPool.
    /// Returns the actual decompressed length.
    /// </summary>
    public static (byte[] Buffer, int Length) DecompressToRentedBuffer(ReadOnlySpan<byte> compressedData)
    {
        if (compressedData.IsEmpty)
            return ([], 0);

        var estimatedSize = Math.Max(compressedData.Length * 4, DefaultBufferSize);
        var outputBuffer = ArrayPool<byte>.Shared.Rent(estimatedSize);
        
        using var inputStream = new ReadOnlySpanStream(compressedData);
        using var brotliStream = new BrotliStream(inputStream, CompressionMode.Decompress);
        
        var totalRead = 0;
        int bytesRead;
        
        while ((bytesRead = brotliStream.Read(outputBuffer.AsSpan(totalRead))) > 0)
        {
            totalRead += bytesRead;
            
            // Need larger buffer
            if (totalRead == outputBuffer.Length)
            {
                var newBuffer = ArrayPool<byte>.Shared.Rent(outputBuffer.Length * 2);
                outputBuffer.AsSpan(0, totalRead).CopyTo(newBuffer);
                ArrayPool<byte>.Shared.Return(outputBuffer);
                outputBuffer = newBuffer;
            }
        }
        
        return (outputBuffer, totalRead);
    }

    /// <summary>
    /// Checks if the data appears to be Brotli compressed.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public static bool IsBrotliCompressed(byte[] data)
    {
        if (data == null || data.Length < 4)
            return false;

        var firstByte = data[0];
        
        // If it starts with '{' (0x7B) or '[' (0x5B), it's likely uncompressed JSON
        if (firstByte == 0x7B || firstByte == 0x5B)
            return false;

        // Try to decompress - if it fails, it's not Brotli
        try
        {
            using var inputStream = new MemoryStream(data, 0, Math.Min(data.Length, 64), writable: false);
            using var brotliStream = new BrotliStream(inputStream, CompressionMode.Decompress);
            Span<byte> buffer = stackalloc byte[1];
            return brotliStream.Read(buffer) >= 0;
        }
        catch
        {
            return false;
        }
    }

    #region Helper Stream Classes

    /// <summary>
    /// MemoryStream that uses a pre-allocated buffer and can expand using ArrayPool.
    /// </summary>
    private sealed class PooledMemoryStream : Stream
    {
        private byte[] _buffer;
        private int _position;
        private int _length;
        private bool _ownsBuffer;

        public PooledMemoryStream(byte[] initialBuffer)
        {
            _buffer = initialBuffer;
            _ownsBuffer = false;
        }

        public override bool CanRead => true;
        public override bool CanSeek => true;
        public override bool CanWrite => true;
        public override long Length => _length;
        public override long Position { get => _position; set => _position = (int)value; }

        public override void Write(byte[] buffer, int offset, int count) 
            => Write(buffer.AsSpan(offset, count));

        public override void Write(ReadOnlySpan<byte> buffer)
        {
            EnsureCapacity(_position + buffer.Length);
            buffer.CopyTo(_buffer.AsSpan(_position));
            _position += buffer.Length;
            if (_position > _length) _length = _position;
        }

        public override int Read(byte[] buffer, int offset, int count)
        {
            var bytesToRead = Math.Min(count, _length - _position);
            if (bytesToRead <= 0) return 0;
            _buffer.AsSpan(_position, bytesToRead).CopyTo(buffer.AsSpan(offset));
            _position += bytesToRead;
            return bytesToRead;
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            _position = origin switch
            {
                SeekOrigin.Begin => (int)offset,
                SeekOrigin.Current => _position + (int)offset,
                SeekOrigin.End => _length + (int)offset,
                _ => _position
            };
            return _position;
        }

        public override void SetLength(long value) => _length = (int)value;
        public override void Flush() { }

        public byte[] ToArray()
        {
            var result = new byte[_length];
            _buffer.AsSpan(0, _length).CopyTo(result);
            return result;
        }

        private void EnsureCapacity(int required)
        {
            if (required <= _buffer.Length) return;
            
            var newSize = Math.Max(_buffer.Length * 2, required);
            var newBuffer = ArrayPool<byte>.Shared.Rent(newSize);
            _buffer.AsSpan(0, _length).CopyTo(newBuffer);
            
            if (_ownsBuffer) ArrayPool<byte>.Shared.Return(_buffer);
            
            _buffer = newBuffer;
            _ownsBuffer = true;
        }

        protected override void Dispose(bool disposing)
        {
            if (_ownsBuffer && _buffer != null)
            {
                ArrayPool<byte>.Shared.Return(_buffer);
                _buffer = null!;
            }
            base.Dispose(disposing);
        }
    }

    /// <summary>
    /// Read-only stream wrapper for ReadOnlySpan.
    /// </summary>
    private sealed class ReadOnlySpanStream : Stream
    {
        private readonly ReadOnlyMemory<byte> _data;
        private int _position;

        public ReadOnlySpanStream(ReadOnlySpan<byte> data)
        {
            _data = data.ToArray(); // Must copy for stream usage
        }

        public override bool CanRead => true;
        public override bool CanSeek => true;
        public override bool CanWrite => false;
        public override long Length => _data.Length;
        public override long Position { get => _position; set => _position = (int)value; }

        public override int Read(byte[] buffer, int offset, int count)
        {
            var bytesToRead = Math.Min(count, _data.Length - _position);
            if (bytesToRead <= 0) return 0;
            _data.Span.Slice(_position, bytesToRead).CopyTo(buffer.AsSpan(offset));
            _position += bytesToRead;
            return bytesToRead;
        }

        public override int Read(Span<byte> buffer)
        {
            var bytesToRead = Math.Min(buffer.Length, _data.Length - _position);
            if (bytesToRead <= 0) return 0;
            _data.Span.Slice(_position, bytesToRead).CopyTo(buffer);
            _position += bytesToRead;
            return bytesToRead;
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            _position = origin switch
            {
                SeekOrigin.Begin => (int)offset,
                SeekOrigin.Current => _position + (int)offset,
                SeekOrigin.End => _data.Length + (int)offset,
                _ => _position
            };
            return _position;
        }

        public override void SetLength(long value) => throw new NotSupportedException();
        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
        public override void Flush() { }
    }

    /// <summary>
    /// Stream that writes directly to an IBufferWriter.
    /// </summary>
    private sealed class BufferWriterStream : Stream
    {
        private readonly IBufferWriter<byte> _writer;

        public BufferWriterStream(IBufferWriter<byte> writer) => _writer = writer;

        public override bool CanRead => false;
        public override bool CanSeek => false;
        public override bool CanWrite => true;
        public override long Length => throw new NotSupportedException();
        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }

        public override void Write(byte[] buffer, int offset, int count)
        {
            var span = _writer.GetSpan(count);
            buffer.AsSpan(offset, count).CopyTo(span);
            _writer.Advance(count);
        }

        public override void Write(ReadOnlySpan<byte> buffer)
        {
            var span = _writer.GetSpan(buffer.Length);
            buffer.CopyTo(span);
            _writer.Advance(buffer.Length);
        }

        public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
        public override void SetLength(long value) => throw new NotSupportedException();
        public override void Flush() { }
    }

    #endregion
}