import os

import cv2


def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4', 'mkv', 'mov')):
    """
    Create a generator for unsupervised training on depth sequences from a video file or folder of video files
    :param video_path_or_folder: Video file or folder with list of video files to iterate through
    :param intrinsics: Intrinsics should be static for a single video
    :param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos
    :return:
    """
    if os.path.isfile(video_path_or_folder):
        # TODO: How to re-yield? Is this enough, since I'm just returning the actual generator?
        #  Or do I need to iterate like below?
        return _single_video_generator(video_path_or_folder)
    else:
        for root, dirs, files in os.walk(video_path_or_folder):
            for file in files:
                if os.path.splitext(file)[1] in allowed_extensions:
                    for frames in _single_video_generator(file):
                        yield frames


def _single_video_generator(video_file):
    # Single video file
    video = cv2.VideoCapture(video_file)

    try:
        # Buffer to store 3 frames, yield when this fills up
        current_frames = []
        while video.grab():
            # TODO: Should I be skipping frames or doing some magic to the frames? Or just leave that to some other
            #  function assuming this will be used in a tf.data object?
            current_frames.append(video.retrieve())
            if len(current_frames) == 3:
                temp_frames = current_frames
                current_frames = []
                # TODO: Convert to tensor or something notable (e.g. dict of 3 frames)
                yield temp_frames
    finally:
        video.release()


def image_generator(root_folder):
    """
    Create an image generator for unsupervised training
    :param root_folder:
    :return:
    """
    pass