import io
import os
import tarfile
import pickle
import zstandard
import torch
from torch.utils.data import IterableDataset, DataLoader

class MultiFileWrapper(io.BufferedIOBase):
    """将多个文件对象包装成一个连续的缓冲流，提高读取性能"""
    def __init__(self, paths):
        self.files = [open(p, 'rb') for p in paths]
        self.current_idx = 0

    def read(self, size=-1):
        if self.current_idx >= len(self.files):
            return b''
        
        data = self.files[self.current_idx].read(size)
        while (not data or len(data) == 0) and self.current_idx < len(self.files) - 1:
            self.current_idx += 1
            data = self.files[self.current_idx].read(size)
        return data

    def readable(self):
        return True

    def close(self):
        for f in self.files:
            f.close()

class DyMeshIterableDataset(IterableDataset):
    def __init__(self, archive_paths=None):
        super(DyMeshIterableDataset, self).__init__()
        if archive_paths is None:
            base = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/"
            self.archive_paths = [os.path.join(base, f"DyMesh_50000v_16f_0000_part_{str(i).zfill(2)}") for i in range(4)]
        else:
            self.archive_paths = archive_paths

    def __iter__(self):
        wrapper = MultiFileWrapper(self.archive_paths)
        try:
            dctx = zstandard.ZstdDecompressor()
            with dctx.stream_reader(wrapper) as reader:
                with tarfile.open(fileobj=reader, mode='r|') as tar:
                    for member in tar:
                        if member.isfile():
                            f = tar.extractfile(member)
                            if f:
                                data = pickle.load(f)
                                if isinstance(data, dict) and 'vertices' in data:
                                    yield {
                                        'vertices': data['vertices'], # List or Array [16, V, 3]
                                        'faces': data['faces'],       # Array [F, 3]
                                        'caption': data.get('caption', '无文本描述')
                                    }
        except Exception as e:
            print(f"❌ 跨文件流读取异常: {e}")
        finally:
            wrapper.close()

def collate_fn(batch):
    """仅负责将数据组织成 Tensor，不再进行 VAE 推理"""
    vertices_list = [torch.tensor(item['vertices'], dtype=torch.float32) for item in batch]
    faces_list = [torch.tensor(item['faces'], dtype=torch.long) for item in batch]
    caption_list = [item['caption'] for item in batch]
    
    return {
        'vertices': [item['vertices'] for item in batch],
        'faces': [item['faces'] for item in batch],
        'caption': [item['caption'] for item in batch]
    }

def get_mesh_dataloader(batch_size=4, num_workers=0, archive_paths=None):
    """统一的 DataLoader 获取接口"""
    dataset = DyMeshIterableDataset(archive_paths=archive_paths)
    return DataLoader(
        dataset, 
        batch_size=batch_size, 
        collate_fn=collate_fn,
        num_workers=num_workers, # 现在可以安全地设置 num_workers > 0
        pin_memory=True if torch.cuda.is_available() else False
    )