#nullable enable
using System;
using System.Buffers;

namespace ShrinkNetwork
{
    /// <summary>Owns rented memory until Dispose. A sender must await completion before disposing.</summary>
    public sealed class ShrinkBufferWriter : IBufferWriter<byte>, IDisposable
    {
        private byte[]? _buffer;
        public ShrinkBufferWriter(int initialCapacity = 256) => _buffer = ArrayPool<byte>.Shared.Rent(Math.Max(1, initialCapacity));
        public int WrittenCount { get; private set; }
        public ReadOnlyMemory<byte> WrittenMemory => Buffer.AsMemory(0, WrittenCount);
        internal Span<byte> WrittenSpan => Buffer.AsSpan(0, WrittenCount);
        private byte[] Buffer => _buffer ?? throw new ObjectDisposedException(nameof(ShrinkBufferWriter));
        public void Advance(int count)
        {
            if (count < 0 || count > Buffer.Length - WrittenCount) throw new ArgumentOutOfRangeException(nameof(count));
            WrittenCount += count;
        }
        public Memory<byte> GetMemory(int sizeHint = 0) { Ensure(sizeHint); return Buffer.AsMemory(WrittenCount); }
        public Span<byte> GetSpan(int sizeHint = 0) { Ensure(sizeHint); return Buffer.AsSpan(WrittenCount); }
        private void Ensure(int sizeHint)
        {
            if (sizeHint < 0) throw new ArgumentOutOfRangeException(nameof(sizeHint));
            sizeHint = Math.Max(1, sizeHint);
            if (sizeHint <= Buffer.Length - WrittenCount) return;
            var next = ArrayPool<byte>.Shared.Rent(checked(Math.Max(Buffer.Length * 2, WrittenCount + sizeHint)));
            Buffer.AsSpan(0, WrittenCount).CopyTo(next);
            ArrayPool<byte>.Shared.Return(Buffer, clearArray: true);
            _buffer = next;
        }
        public void Dispose()
        {
            var buffer = _buffer;
            _buffer = null;
            if (buffer != null) ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
        }
    }
}