import os
import logging
import h5py
import numpy as np
from tqdm import tqdm
from modules import mediapipe_generator

# Configuration
DATA_ROOT = "fsl-data/sentence_data" 
HDF5_LOCATION = "sentence_landmarks.h5"

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')

def main():
    if not os.path.exists(DATA_ROOT):
        logging.error(f"Directory {DATA_ROOT} not found!")
        return

    video_files = [f for f in os.listdir(DATA_ROOT) if f.lower().endswith(('.mp4', '.avi', '.mov'))]
    
    with h5py.File(HDF5_LOCATION, 'w') as f:
        for filename in tqdm(video_files, desc="Processing Videos"):
            video_path = os.path.join(DATA_ROOT, filename)
            
            try:
                data = mediapipe_generator.generate_mediapipe(filepath=video_path)
                if not data:
                    logging.warning(f"No landmarks detected for {filename}")
                    continue

                p_seq = np.array([mediapipe_generator.extract_to_array(r.pose_landmarks, 33, 4) for r in data])
                f_seq = np.array([mediapipe_generator.extract_to_array(r.face_landmarks, 468, 3) for r in data])
                lh_seq = np.array([mediapipe_generator.extract_to_array(r.left_hand_landmarks, 21, 3) for r in data])
                rh_seq = np.array([mediapipe_generator.extract_to_array(r.right_hand_landmarks, 21, 3) for r in data])

                group_key = filename.replace('/', '_')
                sample_grp = f.create_group(group_key)
                
                sample_grp.create_dataset('pose', data=p_seq, compression="gzip")
                sample_grp.create_dataset('face', data=f_seq, compression="gzip")
                sample_grp.create_dataset('left_hand', data=lh_seq, compression="gzip")
                sample_grp.create_dataset('right_hand', data=rh_seq, compression="gzip")
                
                sample_grp.attrs['frame_count'] = len(data)
                sample_grp.attrs['original_filename'] = filename

            except Exception as e:
                logging.error(f"Critical error processing {filename}: {e}")

    logging.info(f"HDF5 dataset created successfully at {HDF5_LOCATION}")

if __name__ == "__main__":
    main()