import os

import cv2
from tqdm import tqdm


class ImageCreator:
    def __init__(self, filename, imgs_dir, image_start=0, image_end=0, pbar=True):
        """
        :param str filename: The name of the video's filename.
        :param str imgs_dir: The directory where to store the image files.
        :param int image_start: The first image to be extracted.
        :param int image_end: The last image to be extracted, 0 if full video.
        :param bool pbar: Whether to display a progress bar.
        """

        self.filename = filename
        self.imgs_dir = imgs_dir
        self.image_start = image_start
        self.image_end = image_end
        self.pbar = pbar
        if not os.path.exists(imgs_dir):
            os.makedirs(imgs_dir)

    def get_images(self):
        vid = cv2.VideoCapture(self.filename)
        total_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
        success, image = vid.read()
        count = 0
        if self.image_end == 0:
            self.image_end = total_frames
        zfill_max = len(str(total_frames))
        ok_count = 0
        print("Writing images...")
        if self.pbar:
            pb = tqdm(total=total_frames)
        while success:
            if count >= self.image_start and count <= self.image_end:
                cv2.imwrite(
                    f"{self.imgs_dir}/frame_{str(ok_count).zfill(zfill_max)}.png", image
                )
                ok_count += 1
            success, image = vid.read()
            if self.pbar:
                pb.update(1)
            count += 1
        if self.pbar:
            pb.close()
        print("Wrote {} image files.".format(ok_count))