import cv2
import os
import ntpath

class Writer:
    """The base class for writer to process data. 
    Extend the class to implement different writers based on requirement. 
    """

    def _generate_filename_from_filepath(self, filepath):
        """Retreives filename from its path on disk

        :param filepath: Path of file on disk
        :type filepath: str
        :return: name of the file
        :rtype: str
        """
        input_filepath = ntpath.basename(filepath)
        input_filename = input_filepath.split(".")[0]

        return input_filename

    def generate_output_filename(self, filepath, keyframe_number):
        """Generates the output filename given the filepath and keyframe number.

        :param filepath: filepath of the original input file
        :type filepath: str
        :param keyframe_number: keyframe number
        :type keyframe_number: int
        :return: filename
        :rtype: str
        """

        return self._generate_filename_from_filepath(filepath)

    
    def write(self, filepath, data):
        """The writer logic for processing the keyframes for the file @ filepath

        :param filepath: Absolute file path
        :type filepath: str
        :param data: Data generated by Katna library for the file @ filepath
        :type data: list
        """
        pass

    def _create_dir(self, dir_path):
        """Creates directory given the directory path.

        :param path: path to output directory
        :type path: str
        """

        if not os.path.isdir(os.path.join(".", dir_path)):
            os.mkdir(os.path.join(".", dir_path))
    


class KeyFrameDiskWriter(Writer):
    """Writes the keyframe data to the disk

    :param BaseWriter: [description]
    :type BaseWriter: [type]
    """
    def __init__(self, location, file_ext=".jpeg"):
        """Initalizes the file extension of the output file and location of output directory.

        :param location: [description]
        :type location: [type]
        :param file_ext: [description], defaults to ""
        :type file_ext: str, optional
        """
        self.output_dir_path = location
        self.file_ext = file_ext
        self._create_dir(location)

    def generate_output_filename(self, filepath, keyframe_number):
        """Generates the filename of output data file.

        :param filepath: path of the file
        :type filepath: str
        :param keyframe_number: keyframe number
        :type keyframe_number: int
        :return: name of the output file 
        :rtype: str
        """

        file_name_arr = []
        input_file_name = self._generate_filename_from_filepath(filepath)

        # append the input filename
        file_name_arr.append(input_file_name)

        # append the keyframe number
        file_name_arr.append(str(keyframe_number))
        
        # join the name using _ 
        filename = "_".join(file_name_arr)
        
        return filename


    def save_frame_data_to_disk(self, frame, file_name):
        """Saves the keyframe data to the disk at ouutput directory path set at
        initialization.

        :param frame: frame data to be written
        :type frame: np.array
        :param file_name: Name of the file
        :type file_name: str
        """

        file_full_path = os.path.join(self.output_dir_path, file_name + self.file_ext)
        cv2.imwrite(file_full_path, frame)

    def write(self, filepath, data):
        """Write the data generated by Katna library for file @ filepath on to the disk.

        :param filepath: path of the input file processed by Katna
        :type filepath: str
        :param data: data generated using Katna library
        :type data: list
        """

        for counter, img in enumerate(data):
            output_filename = self.generate_output_filename(filepath, keyframe_number=counter)
            self.save_frame_data_to_disk(img, file_name=output_filename)


class  ImageCropDiskWriter(Writer):
    """DiskWriter for Image Crop

    :param Writer: Base class for writer
    :type Writer: [type]
    """

    def __init__(self, location, file_ext=".jpeg"):
        """Initializes the location where the cropped images needs to be saved.

        :param location: [description]
        :type location: [type]
        :param file_ext: [description]
        :type file_ext: [type]
        """

        self.output_dir_path = location
        self.file_ext = file_ext
        self._create_dir(location)

    def generate_output_filename(self, filepath, crop_number):
        """Generates the filename of output data file.

        :param filepath: path of the file
        :type filepath: str
        :param crop_number: crop number
        :type crop_number: int
        :return: name of the output file 
        :rtype: str
        """

        file_name_arr = []
        input_file_name = self._generate_filename_from_filepath(filepath)

        # append the input filename
        file_name_arr.append(input_file_name)

        # append the keyframe number
        file_name_arr.append(str(crop_number))
        
        # join the name using _ 
        filename = "_".join(file_name_arr)
        
        return filename

    def save_crop_to_disk(self, crop_rect, frame, file_name):
        """Saves the image crop data to the disk at output directory path set at
        initialization.

        :param crop_rect: In-memory crop_rect.
        :type crop_rect: crop_rect, required
        :param frame: In-memory input image.
        :type frame: numpy.ndarray, required
        :param file_name: name of the image.
        :type file_name: str, required
        """

        cropped_img = crop_rect.get_image_crop(frame)
        file_full_path = os.path.join(self.output_dir_path, file_name + self.file_ext)
        cv2.imwrite(file_full_path, cropped_img)

    def write(self, filepath, data):
        """Write the data generated by Katna library for file @ filepath on to the disk.

        :param filepath: path of the input file processed by Katna
        :type filepath: str
        :param data: data generated using Katna library
        :type data: list
        """
        if len(data) > 0:
            for counter, crop in enumerate(data):
                img = cv2.imread(filepath)
                output_filename = self.generate_output_filename(filepath, crop_number=counter)
                self.save_crop_to_disk(crop, img, file_name=output_filename)
        else:
            print("ERROR : No Perfect crop found for Image %s" %(filepath))

    

    

    