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 VideoOutputStream(DataStreamWrapper):
    def __init__(self, basepath, dims, filenameFormat="{deviceId}_{timestamp}", fps=30, publish_hook=None):
        super().__init__()
        self.basepath = basepath
        self.dims = (int(dims[0]),int(dims[1]))
        self.fps = float(fps)
        self.filenameFormat = filenameFormat
        self.publish_hook = publish_hook
        self.four_cc = cv2.VideoWriter_fourcc(*'mp4v')
        self.out = None

    def publish(self, x):
        if self.publish_hook is not None:
            x = self.publish_hook(x)
        if len(x["metric"]) > 0:
            if self.out is None:
                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)
                self.out = cv2.VideoWriter("{}.mp4".format(os.path.join(path, self.filenameFormat.format(**x))),
                                           self.four_cc, self.fps, self.dims)
            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)
            self.out.write(frame)
        else:
            if self.out  is not None:
                self.out.release()
                self.out = None
        return True

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