import cv2
import numpy as np
import random
import os

def extract_video_frames(video_path, n_frames=30, frame_size=(96, 96)):
    """
    Extracts frames from a video, handling various lengths and potential errors.

    Args:
        video_path (str): Path to the video file.
        n_frames (int): The target number of frames to extract.
        frame_size (tuple): The target (width, height) for each frame.

    Returns:
        np.ndarray: An array of shape (n_frames, height, width, 3) with normalized
                    pixel values (0-1), or None if extraction fails critically.
                    Frames will be padded if the video is too short or has read errors.
    """
    if not os.path.exists(video_path):
        print(f"Error: Video file not found at {video_path}")
        return None

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video file {video_path}")
        return None

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # Basic validation
    if total_frames < 1:
        print(f"Warning: Video has {total_frames} frames. Cannot extract.")
        cap.release()
        # Return array of zeros matching the expected shape
        return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
    if fps < 1:
        print(f"Warning: Video has invalid FPS ({fps}). Proceeding, but timing might be off.")
        # Use a default assumption if FPS is invalid but frames exist
        fps = 30.0 # Or another sensible default

    frames = []
    extracted_count = 0
    last_good_frame_processed = None # Store the last successfully processed frame

    # Calculate indices of frames to attempt extraction (evenly spaced)
    # Ensure indices are within the valid range [0, total_frames - 1]
    indices = np.linspace(0, total_frames - 1, n_frames, dtype=int)

    for i, frame_index in enumerate(indices):
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
        ret, frame = cap.read()

        processed_frame = None
        if ret and frame is not None:
            try:
                # Process valid frame
                frame_resized = cv2.resize(frame, frame_size)
                frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
                processed_frame = frame_rgb.astype(np.float32) / 255.0
                last_good_frame_processed = processed_frame # Update last good frame
                extracted_count += 1
            except cv2.error as e:
                print(f"Warning: OpenCV error processing frame {frame_index}: {e}")
                # Fallback to last good frame if available
                if last_good_frame_processed is not None:
                    processed_frame = last_good_frame_processed.copy()
                else: # If no good frame seen yet, create a placeholder
                    processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)
            except Exception as e:
                 print(f"Warning: Unexpected error processing frame {frame_index}: {e}")
                 if last_good_frame_processed is not None:
                    processed_frame = last_good_frame_processed.copy()
                 else:
                    processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)

        else:
            # Handle read failure (e.g., end of video reached early, corrupted frame)
            print(f"Warning: Failed to read frame at index {frame_index}. Using fallback.")
            if last_good_frame_processed is not None:
                processed_frame = last_good_frame_processed.copy()
            else:
                # If read fails and no previous frame exists, use a zero frame
                processed_frame = np.zeros((*frame_size[::-1], 3), dtype=np.float32)

        frames.append(processed_frame)

    cap.release()

    if extracted_count == 0 and total_frames > 0:
         print("Warning: Failed to extract or process any valid frames, returning array of zeros.")
         # This case should ideally be covered by fallbacks, but as a safeguard:
         return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)

    # Ensure the final output always has n_frames by padding if necessary
    # (This should technically be handled by the loop logic now, but double-check)
    final_frames = np.array(frames)
    if final_frames.shape[0] < n_frames:
        print(f"Warning: Padding needed, final array shape {final_frames.shape} vs target {n_frames}")
        if final_frames.shape[0] == 0: # If somehow array is empty
             padding = np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)
        else:
            padding_needed = n_frames - final_frames.shape[0]
            # Use the very last frame in the list (could be a fallback frame) for padding
            last_frame_for_padding = final_frames[-1][np.newaxis, ...]
            padding = np.repeat(last_frame_for_padding, padding_needed, axis=0)
        final_frames = np.concatenate((final_frames, padding), axis=0)
    elif final_frames.shape[0] > n_frames:
        # Should not happen with linspace logic, but truncate if it does
        print(f"Warning: More frames than expected ({final_frames.shape[0]}), truncating to {n_frames}")
        final_frames = final_frames[:n_frames]


    # Final check of output shape
    if final_frames.shape != (n_frames, frame_size[1], frame_size[0], 3):
         print(f"Error: Final frame array shape mismatch! Expected {(n_frames, frame_size[1], frame_size[0], 3)}, Got {final_frames.shape}")
         # Attempt to reshape or return None/zeros? Returning zeros is safer.
         return np.zeros((n_frames, *frame_size[::-1], 3), dtype=np.float32)


    return final_frames