import os
import glob
import subprocess
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

def extract_frames(mp4_path, target_dir):
    os.makedirs(target_dir, exist_ok=True)
    out_pattern = os.path.join(target_dir, "%04d.png")
    
    cmd = [
        "ffmpeg", "-y", "-ss", "8", "-i", mp4_path,
        "-vf", "fps=24,scale=640:-1",
        out_pattern
    ]
    # 執行 ffmpeg 並隱藏輸出以保持進度條乾淨
    try:
        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"\nError processing {mp4_path}: {e}")
        return False

def main():
    src_dir = r"D:\Users\roy\CV_final_project\video_dataset"
    base_dir = r"D:\Users\roy\CV_final_project\champ-master"
    new_driving_dir = os.path.join(base_dir, "driving_videos_new")
    old_driving_dir = os.path.join(base_dir, "driving_videos")
    backup_dir = os.path.join(base_dir, "driving_videos_old")
    
    mp4_files = glob.glob(os.path.join(src_dir, "*.mp4"))
    print(f"找到 {len(mp4_files)} 部 MP4 影片，準備用 ffmpeg 進行裁切與抽幀...")
    
    success_count = 0
    # 使用 8 個執行緒平行處理
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = {}
        for mp4_path in mp4_files:
            vname = os.path.splitext(os.path.basename(mp4_path))[0]
            target_dir = os.path.join(new_driving_dir, vname, "images")
            futures[executor.submit(extract_frames, mp4_path, target_dir)] = vname
            
        for future in tqdm(as_completed(futures), total=len(futures), desc="處理進度"):
            if future.result():
                success_count += 1
                
    print(f"\n成功處理了 {success_count} / {len(mp4_files)} 部影片！")
    
    # 替換資料夾
    if os.path.exists(backup_dir):
        print("正在刪除舊的備份資料夾...")
        shutil.rmtree(backup_dir)
        
    if os.path.exists(old_driving_dir):
        print("正在將舊有的 driving_videos 備份為 driving_videos_old...")
        os.rename(old_driving_dir, backup_dir)
        
    print("正在將新生成的圖片套用至 driving_videos...")
    os.rename(new_driving_dir, old_driving_dir)
    print("\n替換完成！您現在可以重新執行 `batch_run.py` 了。")

if __name__ == "__main__":
    main()