import ffmpeg
import numpy as np


class FfmpegReader:
    def __init__(self, path):
        self.process1 = (
            ffmpeg
                .input(path, **{'hwaccel': "cuda"})
                .output('pipe:', format='rawvideo', pix_fmt='rgb24')
                .run_async(pipe_stdout=True)
        )
        self.width=1280
        self.height=720
        # probe = ffmpeg.probe(path)
        # video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        # self.width = int(video_stream['width'])
        # self.height = int(video_stream['height'])
        # print("")

    def read(self):
        in_bytes = self.process1.stdout.read(self.width * self.height * 3)
        if not in_bytes:
            return False, None
        in_frame = (
            np
                .frombuffer(in_bytes, np.uint8)
                .reshape([self.width, self.height, 3])
        )
        return True, in_frame


class FfmpegWriter:
    def __init__(self, output_path, width, height, fps, gpu_on=False):

        self.width = width
        self.height = height
        self.fps = fps
        if gpu_on:
            self.process2 = (
                ffmpeg
                    .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(self.height, self.width),
                           **{"hwaccel_output_format": "cuda"})
                    .output(output_path, pix_fmt='yuv420p', **{"c:v": "h264_nvenc"})
                    .overwrite_output()
                    .run_async(pipe_stdin=True)
            )
        else:
            self.process2 = (
                ffmpeg
                    .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(self.height, self.width))
                    .output(output_path, pix_fmt='yuv420p', **{"r":self.fps, "c:v": "h264"})
                    .overwrite_output()
                    .run_async(pipe_stdin=True)
            )

    def write(self, image):
        self.process2.stdin.write(
            image
                .astype(np.uint8)
                .tobytes()
        )


if __name__ == "__main__":
    input = FfmpegReader("/home/sergej/Downloads/borya.mp4")
    output = None
    while True:
        ret, frame = input.read()

        # print("test")
        if not ret:
            break

        if output is None:
            output = FfmpegWriter("ffmpeg_write_dec9.mp4", frame.shape[1], frame.shape[0], 30, gpu_on=False)
        output.write(frame)
    print("a")