import os
import shutil
import subprocess
import argparse

def main():
    parser = argparse.ArgumentParser(description="Batch process driving videos for generate_smpls")
    parser.add_argument('--max_videos', type=int, default=-1, help='Number of videos to process (integer). -1 means process all.')
    args = parser.parse_args()

    base_dir = r"D:\Users\roy\CV_final_project\champ-master"
    driving_videos_dir = os.path.join(base_dir, "driving_videos")
    guidance_motion_dir = os.path.join(base_dir, "guidance_motion")
    source_ref_base = os.path.join(base_dir, "reference_imgs")
    
    if not os.path.exists(source_ref_base):
        print(f"錯誤：找不到原始的參考圖片根目錄 {source_ref_base}")
        return

    # 取得所有的影片資料夾名稱
    video_names = [d for d in os.listdir(driving_videos_dir) if os.path.isdir(os.path.join(driving_videos_dir, d))]
    video_names.sort()
    
    if args.max_videos > 0:
        video_names = video_names[:args.max_videos]
        print(f"已指定處理前 {args.max_videos} 支影片。")
        
    print(f"總共找到 {len(video_names)} 支影片需要處理。")

    for video_name in video_names:
        target_ref_folder = os.path.join(guidance_motion_dir, video_name, "reference_imgs")
        video_smpl_output = os.path.join(driving_videos_dir, video_name, "smpl_results", "smpls_group.npz")
        
        # 檢查該影片的動作捕捉是否已完成 (smpls_group.npz 是所有幀處理完最後才產生的檔案)
        if os.path.exists(video_smpl_output):
            print(f"[{video_name}] 動作捕捉 (Task 2) 已全數完成，自動跳過...")
            continue
            
        print(f"\n========================================")
        print(f"正在處理影片: {video_name}")
        print(f"========================================")
        
        # 無論 reference_imgs 是否存在，都執行 copytree (搭配 dirs_exist_ok=True)
        # 這樣如果之前中斷，也能確保缺漏的參考圖片快取被補齊
        print(f"[{video_name}] 準備參考圖片快取 (Task 1)...")
        shutil.copytree(source_ref_base, target_ref_folder, dirs_exist_ok=True)
                
        # 執行 generate_smpls 指令 (若 Task 1 的快取都齊全，generate_smpls 會瞬間跳過 Task 1)
        cmd = [
            "python", "-m", "scripts.data_processors.smpl.generate_smpls",
            "--reference_imgs_folder", target_ref_folder,
            "--driving_video_path", os.path.join("driving_videos", video_name)
        ]
        
        print(f"[{video_name}] 啟動動作捕捉程式...")
        print(f"執行指令: {' '.join(cmd)}")
        result = subprocess.run(cmd, cwd=base_dir)
        
        if result.returncode != 0:
            print(f"[{video_name}] 執行時發生錯誤！")
        else:
            print(f"[{video_name}] 處理完成！\n")

if __name__ == "__main__":
    main()