import os
import subprocess

def main():
    base_dir = r"D:\Users\roy\CV_final_project\champ-master"
    transferd_dir = os.path.join(base_dir, "transferd_result")
    blender_exe = r"D:\Program Files (x86)\blender-3.6.0-windows-x64\blender.exe"
    blend_file = os.path.join("scripts", "data_processors", "smpl", "blend", "smpl_rendering.blend")
    script_path = os.path.join("scripts", "data_processors", "smpl", "render_condition_maps.py")
    ref_img_path = os.path.join("reference_imgs", "images", "ref-01.png")
    
    if not os.path.exists(transferd_dir):
        print(f"錯誤：找不到轉換結果資料夾 {transferd_dir}")
        return

    # 取得所有的影片資料夾名稱
    video_names = [d for d in os.listdir(transferd_dir) if os.path.isdir(os.path.join(transferd_dir, d))]
    video_names.sort()
    
    print(f"總共找到 {len(video_names)} 支影片需要進行 3D 渲染...")

    for video_name in video_names:
        driving_path = os.path.join("transferd_result", video_name, "smpl_results")
        
        # 檢查該影片的 smpl_results 是否存在
        full_driving_path = os.path.join(base_dir, driving_path)
        if not os.path.exists(full_driving_path):
            print(f"[{video_name}] 找不到轉移後的骨架檔案 (跳過): {driving_path}")
            continue
            
        # 斷點續傳檢查：
        # 檢查深度圖 (depth) 資料夾中是否已經產生了所有對應的圖片
        depth_dir = os.path.join(base_dir, "transferd_result", video_name, "depth")
        if os.path.exists(depth_dir):
            depth_imgs = [f for f in os.listdir(depth_dir) if f.endswith('.png')]
            smpl_files = [f for f in os.listdir(full_driving_path) if f.endswith('.npy')]
            
            if len(depth_imgs) > 0 and len(depth_imgs) >= len(smpl_files):
                print(f"[{video_name}] 所有影格皆已渲染完畢，自動跳過...")
                continue
            elif len(depth_imgs) > 0:
                print(f"[{video_name}] 發現未完成的渲染進度，腳本內的 Blender 將會自動略過已存在的圖片繼續渲染...")
                
        print(f"\n========================================")
        print(f"正在渲染影片條件圖: {video_name}")
        print(f"========================================")
        
        cmd = [
            blender_exe, 
            blend_file,
            "--background", 
            "--python", script_path, 
            "--", 
            "--driving_path", driving_path,
            "--reference_path", ref_img_path,
            "--device", "0"
        ]
        
        result = subprocess.run(cmd, cwd=base_dir)
        
        if result.returncode != 0:
            print(f"[{video_name}] 執行 Blender 渲染時發生錯誤！")
        else:
            print(f"[{video_name}] 渲染完成！\n")

if __name__ == "__main__":
    main()