import cv2 as cv
import numpy as np
import mediapipe as mp
import h5py
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_holistic = mp.solutions.holistic

def extract_to_array(landmark_list, num_landmarks, dims=3):
    if not landmark_list:
        return np.zeros((num_landmarks, dims))
    
    return np.array([[lm.x, lm.y, lm.z, getattr(lm, 'visibility', 0.0)] 
                     if dims == 4 else [lm.x, lm.y, lm.z] 
                     for lm in landmark_list.landmark])

def save_holistic_sample(h5_path, class_id, sample_name, mediapipe_results):
    pose_seq = np.array([extract_to_array(r.pose_landmarks, 33, 4) for r in mediapipe_results])
    face_seq = np.array([extract_to_array(r.face_landmarks, 468, 3) for r in mediapipe_results])
    lh_seq = np.array([extract_to_array(r.left_hand_landmarks, 21, 3) for r in mediapipe_results])
    rh_seq = np.array([extract_to_array(r.right_hand_landmarks, 21, 3) for r in mediapipe_results])

    with h5py.File(h5_path, 'a') as f:
        grp = f.require_group(f"{class_id}/{sample_name}")
        
        datasets = {
            'pose': pose_seq, 
            'face': face_seq, 
            'left_hand': lh_seq, 
            'right_hand': rh_seq
        }
        
        for name, data in datasets.items():
            if name in grp: del grp[name]
            ds = grp.create_dataset(name, data=data, compression="gzip", chunks=True)
            
            ds.attrs['dims'] = "x, y, z, visibility" if name == 'pose' else "x, y, z"
            
        grp.attrs['frame_count'] = len(mediapipe_results)

def generate_mediapipe(filepath):
    video_list = []
    frame_count = 0
    cap = cv.VideoCapture(filepath)

    with mp_holistic.Holistic(
        model_complexity=2
    ) as holistic:    
        while cap.isOpened():
            ret, frame = cap.read()
            frame_count += 1

            if not ret:
                break

            results = holistic.process(cv.cvtColor(frame, cv.COLOR_BGR2RGB))
            video_list.append(results)

        cap.release()
        cv.destroyAllWindows()
    return video_list

def generate_mediapipe_gradio(filepath):
    video_list = []
    cap = cv.VideoCapture(filepath)

    target_fps = 30
    frame_interval_ms = 1000 / target_fps
    
    max_w, max_h = 640, 360
    
    orig_w = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
    orig_h = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
    
    scale = min(max_w / orig_w, max_h / orig_h)
    new_dims = (int(orig_w * scale), int(orig_h * scale))

    next_timestamp_to_capture = 0

    with mp_holistic.Holistic(
        model_complexity=2,
        static_image_mode=False
    ) as holistic:    
        while cap.isOpened():
            cap.set(cv.CAP_PROP_POS_MSEC, next_timestamp_to_capture)
            ret, frame = cap.read()

            if not ret:
                break

            resized_frame = cv.resize(frame, new_dims, interpolation=cv.INTER_AREA)

            rgb_frame = cv.cvtColor(resized_frame, cv.COLOR_BGR2RGB)
            results = holistic.process(rgb_frame)
            
            video_list.append(results)

            next_timestamp_to_capture += frame_interval_ms

        cap.release()
        cv.destroyAllWindows()
        
    return video_list