import os
import torch
from utils.autodataloader import DyMeshIterableDataset, get_mesh_dataloader, MultiFileWrapper

def test_dataloader_full_flow():
    print("🚀 开始测试 Dataloader 数据流读取流程...")
    
    # 1. 实例化 DataLoader (无需传入 VAE)
    # 注意：如果文件不存在，请确保路径正确或有 mock 数据
    dataloader = get_mesh_dataloader(batch_size=2, num_workers=0)
    
    # 2. 遍历读取
    total_samples = 0
    try:
        for i, batch in enumerate(dataloader):
            # 这里的 batch 结构应为 collate_fn 返回的结构
            vertices = batch['vertices']
            faces = batch['faces']
            captions = batch['caption']
            
            print(f"✅ 第 {i+1} 个 Batch 读取成功:")
            print(f"   - Vertices shape: {vertices.shape}") # 应为 [B, 16, V, 3]
            print(f"   - Faces shape:    {faces.shape}")    # 应为 [B, F, 3]
            print(f"   - Caption count:  {len(captions)}")
            
            total_samples += len(captions)
            
            # 读满一定数量后停止，避免全量遍历耗时过长
            if total_samples >= 10:
                print("\n🎉 成功读取 10 个样本，测试通过！")
                break
                
    except Exception as e:
        print(f"\n❌ 测试过程中发生异常: {e}")
        import traceback
        traceback.print_exc()

def verify_file_wrapper_only():
    """验证 MultiFileWrapper 是否能连续读取字节流"""
    print("\n🔍 正在验证 MultiFileWrapper 连续性...")
    
    # 修改为实际数据集路径
    base = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/"
    paths = [os.path.join(base, f"DyMesh_50000v_16f_0000_part_{str(i).zfill(2)}") for i in range(4)]
    
    # 过滤掉不存在的文件，防止报错
    existing_paths = [p for p in paths if os.path.exists(p)]
    if not existing_paths:
        print("⚠️ 未找到任何数据文件，请检查路径。")
        return

    wrapper = MultiFileWrapper(existing_paths)
    data = wrapper.read(1024) # 读前 1024 字节
    
    if data and len(data) > 0:
        print(f"✅ 成功读取到数据流，前 10 字节: {data[:10]}")
    else:
        print("❌ 未能读取到数据，请检查路径。")
    wrapper.close()

if __name__ == "__main__":
    # 1. 检查底层文件流是否正常
    verify_file_wrapper_only()
    # 2. 检查 DataLoader 解析逻辑
    test_dataloader_full_flow()