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")
    
    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)} 支影片需要進行 DWPose 骨架特徵提取...")

    for video_name in video_names:
        input_normal_dir = os.path.join("transferd_result", video_name, "normal")
        output_dwpose_dir = os.path.join("transferd_result", video_name, "dwpose")
        
        # 檢查該影片的 normal 是否存在 (因為 generate_dwpose 是吃 normal 的影像)
        full_input_dir = os.path.join(base_dir, input_normal_dir)
        if not os.path.exists(full_input_dir):
            print(f"[{video_name}] 找不到 normal 影像資料夾 (跳過): {input_normal_dir}")
            continue
            
        normal_imgs = [f for f in os.listdir(full_input_dir) if f.endswith('.png') or f.endswith('.jpg')]
        if len(normal_imgs) == 0:
            print(f"[{video_name}] normal 資料夾內沒有圖片 (跳過): {input_normal_dir}")
            continue
            
        # 斷點續傳檢查：檢查 dwpose 是否已完成
        full_output_dir = os.path.join(base_dir, output_dwpose_dir)
        if os.path.exists(full_output_dir):
            dwpose_imgs = [f for f in os.listdir(full_output_dir) if f.endswith('.png') or f.endswith('.jpg')]
            
            if len(dwpose_imgs) >= len(normal_imgs):
                print(f"[{video_name}] 所有影格皆已完成 DWPose 提取，自動跳過...")
                continue
            elif len(dwpose_imgs) > 0:
                print(f"[{video_name}] 發現未完成的進度，腳本將自動略過已存在的圖片繼續執行...")
                
        print(f"\n========================================")
        print(f"正在提取骨架特徵圖: {video_name}")
        print(f"========================================")
        
        cmd = [
            "python", "-m", "scripts.data_processors.dwpose.generate_dwpose",
            "--input", input_normal_dir,
            "--output", output_dwpose_dir
        ]
        
        result = subprocess.run(cmd, cwd=base_dir)
        
        if result.returncode != 0:
            print(f"[{video_name}] 執行 DWPose 時發生錯誤！")
        else:
            print(f"[{video_name}] DWPose 提取完成！\n")

if __name__ == "__main__":
    main()