import io
import os
import resource
from typing import Iterator
import psutil


def monitor_memory(threshold_mb=100):
    process = psutil.Process(os.getpid())
    mem_mb = process.memory_info().rss / (1024 * 1024)
    print(f"[监控] 当前内存使用：{mem_mb:.2f} MB")
    if mem_mb > threshold_mb:
        raise MemoryError(f"内存超出限制！当前：{mem_mb:.2f} MB，限制：{threshold_mb} MB")


# 模拟大数据（比如 1GB）
def generate_large_chunks(chunk_size=1024 * 1024, total_size=1024 * 1024 * 1024) -> Iterator[bytes]:
    sent = 0
    while sent < total_size:
        yield b'x' * chunk_size  # 每次产生1MB的假数据
        sent += chunk_size
        print(f"[生成] 已生成 {sent // (1024*1024)} MB")

# Blob 类
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.")

    def write_to_file(self, path: str) -> int:
        self._ensure_not_consumed()
        full_path = os.path.join(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)
                print(f"[写入] 已写入 {written // (1024*1024)} MB")
        self._consumed = True
        return written

    def get_input_stream(self) -> io.IOBase:
        self._ensure_not_consumed()
        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 test_streaming_write():
    try:
        data_iter = generate_large_chunks(chunk_size=1024 * 1024, total_size=1024 * 1024 * 1024)  # 1GB
        blob = Blob("test_bigfile.bin", data_iter)

        full_path = os.path.join("./", blob.name)
        written = 0
        with open(full_path, 'wb') as f:
            for chunk in blob._original_chunk_iter:
                monitor_memory(threshold_mb=100)  # 每写一块检查内存
                f.write(chunk)
                written += len(chunk)
                print(f"[写入] 已写入 {written // (1024 * 1024)} MB")

        print(f"[完成] 成功写入 {written // (1024*1024)} MB 到文件")
    except MemoryError as e:
        print(f"[内存警告] {e}")
    except Exception as e:
        print(f"[异常] {e}")


if __name__ == "__main__":
    test_streaming_write()
