from json import dump, load
from os import path
from time import sleep

from pycs import app
from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaStorage import MediaStorage
from pycs.interfaces.Pipeline import Pipeline as Interface


class Pipeline(Interface):
    def __init__(self, root_folder, distribution):
        app.logger.debug('fmv1 init')
        self.root_folder = root_folder

    def close(self):
        app.logger.debug('fmv1 close')

    def execute(self, storage: MediaStorage, file: MediaFile):
        app.logger.debug('fmv1 execute')

        data_file = path.join(self.root_folder, 'data.json')
        if path.exists(data_file):
            with open(data_file, 'r') as f:
                result = load(f)
        else:
            result = {}

        if file.path in result:
            for r in result[file.path]:
                if r['type'] == 'MediaBoundingBox':
                    file.add_bounding_box(r['x'], r['y'], r['w'], r['h'], r['label'], r['frame'])
                if r['type'] == 'MediaImageLabel':
                    file.set_image_label(r['label'], r['frame'])

    def fit(self, storage: MediaStorage):
        app.logger.debug('fmv1 fit')

        for i in range(10):
            yield i / 10
            sleep(1)

        result = {}
        for f in storage.files().iter():
            result[f.path] = list(map(lambda r: dict(r.__dict__, **{'type': type(r).__name__}),
                                      f.results()))

        data_file = path.join(self.root_folder, 'data.json')
        with open(data_file, 'w') as file:
            dump(result, file, indent=4)