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

class DyMeshIterableDataset(IterableDataset):
    def __init__(self, data_dir, rank=0, world_size=1):
        super(DyMeshIterableDataset, self).__init__()
        # 获取所有分片文件
        self.tar_paths = sorted(glob.glob(os.path.join(data_dir, "shard_*.tar.zst")))
        # 根据 Rank 进行数据切分
        self.my_tar_paths = [p for i, p in enumerate(self.tar_paths) if i % world_size == rank]
        print(f"[Rank {rank}] 负责处理 {len(self.my_tar_paths)} 个分片")

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        # 根据 Worker 进行数据切分
        if worker_info is None:
            curr_paths = self.my_tar_paths
        else:
            per_worker = (len(self.my_tar_paths) + worker_info.num_workers - 1) // worker_info.num_workers
            wid = worker_info.id
            curr_paths = self.my_tar_paths[wid * per_worker : (wid + 1) * per_worker]

        dctx = zstandard.ZstdDecompressor()
        for tar_path in curr_paths:
            try:
                with open(tar_path, 'rb') as fh:
                    with dctx.stream_reader(fh) 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'],
                                                'faces': data['faces'],
                                                'caption': data.get('caption', '无文本描述')
                                            }
            except Exception as e:
                print(f"❌ 读取错误 {tar_path}: {e}")
                continue

def collate_fn(batch):
    """负责将数据组织成 Batch，格式与原始代码保持一致"""
    return {
        'vertices': [torch.tensor(item['vertices'], dtype=torch.float32) for item in batch],
        'faces': [torch.tensor(item['faces'], dtype=torch.long) for item in batch],
        'caption': [item['caption'] for item in batch]
    }

def get_mesh_dataloader(data_dir, batch_size=4, num_workers=2):
    """统一的 DataLoader 获取接口"""
    rank = int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    
    dataset = DyMeshIterableDataset(data_dir, rank=rank, world_size=world_size)
    
    return DataLoader(
        dataset, 
        batch_size=batch_size, 
        collate_fn=collate_fn,
        num_workers=num_workers, 
        pin_memory=True if torch.cuda.is_available() else False
    )