import io
import pyperf

CHUNK_SIZE = 4096
SIZES = [1, 4, 8, 16, 32, 64, 128]


class ChunkedRaw(io.RawIOBase):
    def __init__(self, data, chunk_size):
        self._buf = memoryview(data)
        self._pos = 0
        self._chunk_size = chunk_size

    def readable(self):
        return True

    def read(self, n: int = -1):
        if self._pos >= len(self._buf):
            return b""

        to_read = (
            self._chunk_size
            if (n is None or n < 0)
            else min(n, self._chunk_size)
        )

        end = min(self._pos + to_read, len(self._buf))
        out = bytes(self._buf[self._pos : end])
        self._pos = end

        return out


def generate_bytes(total):
    block = b"abcdefghijklmnopqrstuvwxyz0123456789" * 128
    return (block * (total // len(block) + 1))[:total]


def _bench_readall(data, chunk_size):
    r = ChunkedRaw(data, chunk_size)
    out = r.readall()
    if len(out) != len(data):
        raise RuntimeError("what is going on...???")


def main():
    runner = pyperf.Runner()
    for size_mib in SIZES:
        total_bytes = size_mib * 1024 * 1024
        data = generate_bytes(total_bytes)
        name = f"rawiobase_readall_{size_mib}MB_chunk{CHUNK_SIZE}"
        runner.bench_func(name, _bench_readall, data, CHUNK_SIZE)


if __name__ == "__main__":
    main()