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

# --- Configuration ---
DATA_ROOT = "fsl-data/sliced_dataset" 
HDF5_LOCATION = "fsl-data/multitask_mediapipe.h5"
SIGN_MAP_PATH = "metadata/sign_dict.json"
EMOTION_MAP_PATH = "metadata/emotion_dict.json"

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

def get_mappings(root_path):
    signs = sorted([d for d in os.listdir(root_path) if os.path.isdir(os.path.join(root_path, d))])
    first_sign = os.path.join(root_path, signs[0])
    emotions = sorted([d for d in os.listdir(first_sign) if os.path.isdir(os.path.join(first_sign, d))])
    
    return {name: i for i, name in enumerate(signs)}, {name: i for i, name in enumerate(emotions)}

def save_json_mapping(mapping, filepath):
    inverted_map = {int(v): k for k, v in mapping.items()}
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(inverted_map, f, indent=4, ensure_ascii=False)
    logging.info(f"Mapping saved to {filepath}")

def main():
    sign_map, emotion_map = get_mappings(DATA_ROOT)
    
    save_json_mapping(sign_map, SIGN_MAP_PATH)
    save_json_mapping(emotion_map, EMOTION_MAP_PATH)
    
    tasks = []
    for sign in sign_map:
        for emotion in emotion_map:
            folder = os.path.join(DATA_ROOT, sign, emotion)
            if os.path.exists(folder):
                for video in os.listdir(folder):
                    tasks.append({
                        'path': os.path.join(folder, video),
                        'sign_id': sign_map[sign],
                        'emotion_id': emotion_map[emotion],
                        'name': video
                    })

    num_samples = len(tasks)
    str_dt = h5py.string_dtype(encoding='utf-8')

    with h5py.File(HDF5_LOCATION, 'w') as f:
        ds_l_sign = f.create_dataset('label_sign', (num_samples,), dtype='i')
        ds_l_emotion = f.create_dataset('label_emotion', (num_samples,), dtype='i')
        ds_l_path = f.create_dataset('label_filepath', (num_samples,), dtype=str_dt)
        
        data_group = f.create_group('data')

        for i, task in enumerate(tqdm(tasks, desc="Processing Videos")):
            try:
                data = mediapipe_generator.generate_mediapipe(filepath=task['path'])
                if not data: 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])

                sample_grp = data_group.create_group(str(i))
                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")
                
                ds_l_sign[i] = task['sign_id']
                ds_l_emotion[i] = task['emotion_id']
                ds_l_path[i] = task['path']
                
                sample_grp.attrs['video_name'] = task['name']
                sample_grp.attrs['frame_count'] = len(data)

            except Exception as e:
                logging.error(f"Error at {task['name']}: {e}")

        f.attrs['sign_map_ref'] = SIGN_MAP_PATH
        f.attrs['emotion_map_ref'] = EMOTION_MAP_PATH

if __name__ == "__main__":
    main()