import os
import shutil
import glob

def copy_checkpoints(indices, src_root, dst_root):
    for idx in indices:
        idx_str = f"{idx:04d}"
        # 1. 체크포인트 폴더 찾기
        pattern = os.path.join(src_root, f"{idx_str}_HYVideo*")
        matches = glob.glob(pattern)
        if len(matches) == 0:
            print(f"[ERROR] No folder found for index {idx_str}")
            continue
        if len(matches) > 1:
            print(f"[ERROR] Multiple folders found for index {idx_str}: {matches}")
            continue
        src_folder = matches[0]
        dst_folder = os.path.join(dst_root, os.path.basename(src_folder))
        os.makedirs(dst_folder, exist_ok=True)

        # 2. checkpoints 폴더 처리
        checkpoints_dir = os.path.join(src_folder, "checkpoints")
        dst_checkpoints_dir = os.path.join(dst_folder, "checkpoints")
        if os.path.isdir(checkpoints_dir):
            os.makedirs(dst_checkpoints_dir, exist_ok=True)
            # 가장 숫자가 높은 하위폴더 찾기
            subfolders = [f for f in os.listdir(checkpoints_dir) if os.path.isdir(os.path.join(checkpoints_dir, f))]
            max_num = -1
            max_folder = None
            for f in subfolders:
                try:
                    num = int(f)
                    if num > max_num:
                        max_num = num
                        max_folder = f
                except ValueError:
                    continue
            if max_folder is None:
                print(f"[ERROR] No valid subfolder in {checkpoints_dir}")
            else:
                src_max_folder = os.path.join(checkpoints_dir, max_folder)
                dst_max_folder = os.path.join(dst_checkpoints_dir, max_folder)
                os.makedirs(dst_max_folder, exist_ok=True)
                src_file = os.path.join(src_max_folder, "mp_rank_00_model_states.pt")
                dst_file = os.path.join(dst_max_folder, "mp_rank_00_model_states.pt")
                if os.path.isfile(src_file):
                    shutil.copy2(src_file, dst_file)
                    print(f"[OK] Copied {src_file} -> {dst_file}")
                else:
                    print(f"[ERROR] File not found: {src_file}")

        # 3. checkpoints 폴더 외의 파일/폴더 복사
        for item in os.listdir(src_folder):
            src_item = os.path.join(src_folder, item)
            dst_item = os.path.join(dst_folder, item)
            if item == "checkpoints":
                continue  # 이미 처리함
            if os.path.isdir(src_item):
                shutil.copytree(src_item, dst_item, dirs_exist_ok=True)
            else:
                shutil.copy2(src_item, dst_item)
        print(f"[OK] Finished copying for index {idx_str}")

# 사용 예시
if __name__ == "__main__":
    indices = [6, 9, 10, 12, 15, 17, 25, 33, 34, 42, 54, 64, 65, 66, 67, 72, 75, 76, 80, 89, 90, 92, 93, 109, 110]  # 복사할 인덱스 리스트
    src_root = "/fsx/junyoung/my_checkpoints/hunyuan/log_EXP"  # 체크포인트 폴더들이 있는 경로
    dst_root = "/fsx/junyoung/my_checkpoints/important_checkpoints"       # 복사할 대상 경로
    copy_checkpoints(indices, src_root, dst_root)