import os
import glob
import pickle
import sys

# Add current dir to path
sys.path.append(os.getcwd())

from utils.render import render_smpl

DATA_DIR = "data"

def batch_render():
    if not os.path.exists(DATA_DIR):
        print(f"Data directory {DATA_DIR} not found.")
        return

    pkl_files = glob.glob(os.path.join(DATA_DIR, "*.pkl"))
    print(f"Found {len(pkl_files)} pickle files.")

    for pkl_path in pkl_files:
        base_name = os.path.splitext(os.path.basename(pkl_path))[0]
        mp4_path = os.path.join(DATA_DIR, f"{base_name}.mp4")

        # Skip if already exists (optional, but good for speed if re-running)
        # User requested render all, so maybe force? 
        # "Render all pkl files ... and save them" implies doing it. 
        # But if we want to update them with new rendering logic, we must overwrite.
        
        print(f"Processing {base_name}...")
        try:
            with open(pkl_path, 'rb') as f:
                data = pickle.load(f)
            
            # Data shape check
            if len(data.shape) != 3 or data.shape[1] != 24 or data.shape[2] != 3:
                print(f"Skipping {base_name}: Unexpected shape {data.shape}")
                continue

            render_smpl(data, mp4_path, fps=30)
            print(f"Saved {mp4_path}")
            
        except Exception as e:
            print(f"Failed to render {base_name}: {e}")

if __name__ == "__main__":
    batch_render()