from edge_engine.common.logsetup import logger
from edge_engine.common.config import DEVICE_ID
from uuid import uuid4
import traceback


class FrameProcessor:
    def __init__(self, stream, model):
        self.model = model
        self.stream = stream
        logger.info("Setting up frame processor !!")

    def run_model(self):
        while self.stream.stream.isOpened():
            try:
                logger.debug("Getting frame mask_model")
                frame = self.stream.read()
                logger.debug("Running mask_model")
                data = {
                    "frame": frame,
                    "frameId": "{}".format(uuid4()),
                    "deviceId": "{}".format(DEVICE_ID),
                }
                self.model.predict(data)
                logger.debug("publishing mask_model output")
            except Exception as e:
                logger.error(e)
                logger.error(traceback.format_exc())


class FrameProcessorv2:
    def __init__(self, stream, model):
        self.model = model
        self.stream = stream
        self.count = 0
        self.skip_frame_every = 3  # 1 does not skip any frame (n-1 frames get skipped)

    def run_model(self):
        while self.stream.running():
            try:
                frame = self.stream.read()
                self.count += 1
                if frame is not None and self.count % self.skip_frame_every == 0:
                    fid = uuid4()
                    data = {
                        "frame": frame,
                        "frameId": "{}".format(fid),
                        "deviceId": "{}".format(DEVICE_ID),
                    }

                    self.model.predict(data)

            except Exception as e:
                logger.error(e)
                logger.error(traceback.format_exc())

        self.model.running = False
