namespace NATS.Server.Tls;

public sealed class PeekableStream : Stream
{
    private readonly Stream _inner;
    private byte[]? _peekedBytes;
    private int _peekedOffset;
    private int _peekedCount;

    public PeekableStream(Stream inner) => _inner = inner;

    public async Task<byte[]> PeekAsync(int count, CancellationToken ct = default)
    {
        var buf = new byte[count];
        int read = await _inner.ReadAsync(buf.AsMemory(0, count), ct);
        if (read < count) Array.Resize(ref buf, read);
        _peekedBytes = buf;
        _peekedOffset = 0;
        _peekedCount = read;
        return buf;
    }

    public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct = default)
    {
        if (_peekedBytes != null && _peekedOffset < _peekedCount)
        {
            int available = _peekedCount - _peekedOffset;
            int toCopy = Math.Min(available, buffer.Length);
            _peekedBytes.AsMemory(_peekedOffset, toCopy).CopyTo(buffer);
            _peekedOffset += toCopy;
            if (_peekedOffset >= _peekedCount) _peekedBytes = null;
            return toCopy;
        }
        return await _inner.ReadAsync(buffer, ct);
    }

    public override int Read(byte[] buffer, int offset, int count)
    {
        if (_peekedBytes != null && _peekedOffset < _peekedCount)
        {
            int available = _peekedCount - _peekedOffset;
            int toCopy = Math.Min(available, count);
            Array.Copy(_peekedBytes, _peekedOffset, buffer, offset, toCopy);
            _peekedOffset += toCopy;
            if (_peekedOffset >= _peekedCount) _peekedBytes = null;
            return toCopy;
        }
        return _inner.Read(buffer, offset, count);
    }

    public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken ct)
        => ReadAsync(buffer.AsMemory(offset, count), ct).AsTask();

    // Write passthrough
    public override void Write(byte[] buffer, int offset, int count) => _inner.Write(buffer, offset, count);
    public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken ct) => _inner.WriteAsync(buffer, offset, count, ct);
    public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct = default) => _inner.WriteAsync(buffer, ct);
    public override void Flush() => _inner.Flush();
    public override Task FlushAsync(CancellationToken ct) => _inner.FlushAsync(ct);

    // Required Stream overrides
    public override bool CanRead => _inner.CanRead;
    public override bool CanSeek => false;
    public override bool CanWrite => _inner.CanWrite;
    public override long Length => throw new NotSupportedException();
    public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
    public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
    public override void SetLength(long value) => throw new NotSupportedException();

    protected override void Dispose(bool disposing) { if (disposing) _inner.Dispose(); base.Dispose(disposing); }
}