from edge_engine.streamio.datastream.datastreamwrapper import DataStreamWrapper
import cv2
import base64
import numpy as np
import os
from edge_engine.common.logsetup import logger
from datetime import datetime


class FrameOutputStream(DataStreamWrapper):
    def __init__(self, basepath, iformat="jpg", filenameFormat="{deviceId}_{frameId}_{timestamp}", publish_hook=None):
        super().__init__()
        self.basepath = basepath
        self.iformat = iformat
        self.filenameFormat = filenameFormat
        self.publish_hook = publish_hook


    def publish(self, x):
        if self.publish_hook is not None:
            x= self.publish_hook(x)
        frame = x["frame"]
        frame = base64.b64decode(frame.split("data:image/jpeg;base64,")[1])
        frame = np.fromstring(frame, np.uint8)
        frame = cv2.imdecode(frame, cv2.IMREAD_COLOR)
        path = os.path.join(self.basepath, datetime.now().date().isoformat())
        if not os.path.isdir(path):
            logger.info("Creating {}  \n".format(path))
            os.mkdir(path)
        cv2.imwrite("{path}.{iformat}".format(path=os.path.join(path, self.filenameFormat.format(**x)),
                                              iformat=self.iformat), frame)
        return True

    def subscribe(self, hook):
        super().subscribe(hook)
