import cv2
import numpy as np
import random

def extract_video_frames(video_path, n_frames=30, frame_size=(96, 96)):
    """
    Simplified robust frame extractor for short videos (2-10 sec)
    - Automatically handles varying video lengths
    - Ensures consistent output shape
    - Optimized for MP4/MPEG
    """
    # Open video
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video {video_path}")
        return None
    
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    # Basic video validation
    if total_frames < 1 or fps < 1:
        print(f"Error: Invalid video (frames:{total_frames}, fps:{fps})")
        cap.release()
        return None
    
    # Calculate how many frames to skip (adaptive based on video length)
    video_length = total_frames / fps
    frame_step = max(1, int(total_frames / n_frames))
    
    frames = []
    last_good_frame = None
    
    for i in range(n_frames):
        # Calculate position to read (spread evenly across video)
        pos = min(int(i * (total_frames / n_frames)), total_frames - 1)
        cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
        
        ret, frame = cap.read()
        
        # Fallback strategies if read fails
        if not ret or frame is None:
            if last_good_frame is not None:
                frame = last_good_frame.copy()
            else:
                # Generate placeholder frame (light gray)
                frame = np.full((*frame_size[::-1], 3), 0.8, dtype=np.float32)
        else:
            # Process valid frame
            frame = cv2.resize(frame, frame_size)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = frame.astype(np.float32) / 255.0
            last_good_frame = frame
        
        frames.append(frame)
    
    cap.release()
    return np.array(frames)