from edge_engine.streamio.datastream.datastreamwrapper import DataStreamWrapper
import subprocess as sp


class FFMPEGOutputStream(DataStreamWrapper):

    def __init__(self, ffmpeg_cmd, rtp_endpoint, publish_hook=None):
        super().__init__()
        self.ffmpeg_cmd = ffmpeg_cmd
        self.rtp_endpoint = rtp_endpoint
        self.ffmpeg_cmd.append(self.rtp_endpoint[0])
        self.proc = sp.Popen(self.ffmpeg_cmd, stdin=sp.PIPE, shell=False)
        self.publish_hook = publish_hook

    def publish(self, x):
        if self.publish_hook is not None:
            x = self.publish_hook(x)
        frame = x["frame"]
        self.proc.stdin.write(frame.tostring())
        self.proc.stdin.flush()
