import tarfile
import zstandard
import os

def count_samples_in_archive(archive_path):
    if not os.path.exists(archive_path):
        print(f"错误: 文件不存在 -> {archive_path}")
        return 0

    print(f"正在分析数据集: {archive_path} ... (这可能需要一些时间)")
    
    count = 0
    dctx = zstandard.ZstdDecompressor()
    
    try:
        with open(archive_path, 'rb') as fh:
            with dctx.stream_reader(fh) as reader:
                # 使用 'r|' 模式进行流式读取，不进行随机访问
                with tarfile.open(fileobj=reader, mode='r|') as tar:
                    for member in tar:
                        # 假设每个样本对应一个文件，且不是目录
                        if member.isfile():
                            count += 1
                            # 每计数 10000 个样本打印一次进度
                            if count % 10000 == 0:
                                print(f"已扫描样本数: {count}")
    except Exception as e:
        print(f"扫描过程中发生异常: {e}")
    
    return count

if __name__ == "__main__":
    archive_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/DyMesh_complete.tar.zst"
    total_samples = count_samples_in_archive(archive_path)
    print("="*40)
    print(f"统计完成！")
    print(f"该数据集总计包含: {total_samples} 个样本")
    print("="*40)