import io
from typing import Iterator


class Blob:
    def __init__(self, name: str, chunk_iter: Iterator[bytes]):
        self.name = name
        self._original_chunk_iter = chunk_iter
        self._consumed = False

    def _ensure_not_consumed(self):
        if self._consumed:
            raise RuntimeError("Blob content has already been consumed. Use streaming APIs or reload the source.")

    def write_to_file(self, path: str) -> int:
        """Streams blob content to a file without loading it all into memory."""
        self._ensure_not_consumed()
        full_path = path + self.name
        written = 0
        with open(full_path, 'wb') as f:
            for chunk in self._original_chunk_iter:
                f.write(chunk)
                written += len(chunk)
        self._consumed = True
        return written

    def get_input_stream(self) -> io.IOBase:
        """Returns a streaming interface. If already consumed, raises error."""
        self._ensure_not_consumed()

        # Wrap the iterator in a file-like stream
        class IteratorStream(io.RawIOBase):
            def __init__(self, iterator):
                self.iterator = iterator
                self.buffer = b''

            def readable(self):
                return True

            def readinto(self, b):
                try:
                    while len(self.buffer) < len(b):
                        self.buffer += next(self.iterator)
                except StopIteration:
                    pass
                n = min(len(b), len(self.buffer))
                b[:n] = self.buffer[:n]
                self.buffer = self.buffer[n:]
                return n

        self._consumed = True
        return io.BufferedReader(IteratorStream(self._original_chunk_iter))

    def __repr__(self):
        return f"Blob(name='{self.name}')"


# 本地测试流式加载
def generate_chunks(data: bytes, chunk_size: int = 4) -> Iterator[bytes]:
    """模拟一个chunked迭代器，每次返回chunk_size大小的二进制块。"""
    for i in range(0, len(data), chunk_size):
        yield data[i:i + chunk_size]


def test_blob_streaming():
    data = b'This is a test of streaming blob data.'
    chunks = generate_chunks(data, chunk_size=5)

    # 初始化 Blob
    blob = Blob(name='test_blob.dat', chunk_iter=chunks)

    # 测试写入文件
    written = blob.write_to_file('./')
    print(f'Wrote {written} bytes to file.')

    # 尝试再次获取流（应该报错）
    try:
        blob.get_input_stream()
    except RuntimeError as e:
        print(f'Expected error after consumption: {e}')

    # 再创建一个 Blob 实例用于测试 get_input_stream
    blob2 = Blob(name='test_blob2.dat', chunk_iter=generate_chunks(data, chunk_size=6))
    stream = blob2.get_input_stream()

    # 从流中读取内容
    result = stream.read()
    print(f'Streamed content: {result}')
    assert result == data


if __name__ == "__main__":
    test_blob_streaming()
