otio_burnin

class ModifiedBurnins(ffmpeg_burnins.Burnins):
    '''
    This is modification of OTIO FFmpeg Burnin adapter.
    - requires FFmpeg in PATH

    Offers 6 positions for burnin text. Each can be set with:
    - static text
    - frames
    - timecode

    Options - dictionary which sets the final look.
    - Datatypes explanation:
    <color> string format must be supported by FFmpeg.
        Examples: "#000000", "0x000000", "black"
    <font> must be accesible by ffmpeg = name of registered Font in system
        or path to font file.
        Examples: "Arial", "C:/Windows/Fonts/arial.ttf"

    - Possible keys:
    "opacity" - Opacity of text - <float, Range:0-1>
    "bg_opacity" - Opacity of background (box around text) - <float, Range:0-1>
    "bg_color" - Background color - <color>
    "bg_padding" - Background padding in pixels - <int>
    "x_offset" - offsets burnin vertically by entered pixels
        from border - <int>
    "y_offset" - offsets burnin horizontally by entered pixels
        from border - <int>
    - x_offset & y_offset should be set at least to same value as bg_padding!!
    "font" - Font Family for text - <font>
    "font_size" - Font size in pixels - <int>
    "font_color" - Color of text - <color>
    "frame_offset" - Default start frame - <int>
        - required IF start frame is not set when using frames
          or timecode burnins

    On initializing class can be set General options through
        "options_init" arg.
    General options can be overridden when adding burnin.

    '''
    TOP_CENTERED = ffmpeg_burnins.TOP_CENTERED
    BOTTOM_CENTERED = ffmpeg_burnins.BOTTOM_CENTERED
    TOP_LEFT = ffmpeg_burnins.TOP_LEFT
    BOTTOM_LEFT = ffmpeg_burnins.BOTTOM_LEFT
    TOP_RIGHT = ffmpeg_burnins.TOP_RIGHT
    BOTTOM_RIGHT = ffmpeg_burnins.BOTTOM_RIGHT

    options_init = {
        'opacity': 1,
        'x_offset': 5,
        'y_offset': 5,
        'bg_padding': 5,
        'bg_opacity': 0.5,
        'font_size': 42
    }

    def __init__(
        self, source, ffprobe_data=None, options_init=None, first_frame=None
    ):
        if not ffprobe_data:
            ffprobe_data = _get_ffprobe_data(source)

        # Validate 'streams' before calling super to raise more specific
        #   error
        source_streams = ffprobe_data.get("streams")
        if not source_streams:
            raise ValueError((
                "Input file \"{}\" does not contain any streams"
                " with image/video content."
            ).format(source))

        self.ffprobe_data = ffprobe_data
        self.first_frame = first_frame
        self.input_args = []
        self.cleanup_paths = []

        super().__init__(source, source_streams)

        if options_init:
            self.options_init.update(options_init)

    def add_text(
        self,
        text,
        align,
        frame_start=None,
        frame_end=None,
        options=None,
    ):
        """
        Adding static text to a filter.

        :param str text: text to apply to the drawtext
        :param enum align: alignment, must use provided enum flags
        :param int frame_start: starting frame for burnins current frame
        :param dict options: recommended to use TextOptions
        """
        if not text:
            # do not add empty text, eg.: if it was using an optional part
            return

        if not options:
            options = ffmpeg_burnins.TextOptions(**self.options_init)

        options = options.copy()
        if frame_start is not None:
            options["frame_offset"] = frame_start

        # `frame_end` is only for meassurements of text position
        if frame_end is not None:
            options["frame_end"] = frame_end

        options["label"] = align
        self._add_burnin(text, align, options, DRAWTEXT)

    def add_timecode(
        self, align, frame_start=None, frame_end=None, frame_start_tc=None,
        text=None, options=None
    ):
        """
        Convenience method to create the frame number expression.

        :param enum align: alignment, must use provided enum flags
        :param int frame_start:  starting frame for burnins current frame
        :param int frame_start_tc:  starting frame for burnins timecode
        :param str text: text that will be before timecode
        :param dict options: recommended to use TimeCodeOptions
        """
        if not options:
            options = ffmpeg_burnins.TimeCodeOptions(**self.options_init)

        options = options.copy()
        if frame_start is not None:
            options["frame_offset"] = frame_start

        # `frame_end` is only for meassurements of text position
        if frame_end is not None:
            options["frame_end"] = frame_end

        if not frame_start_tc:
            frame_start_tc = options["frame_offset"]

        if not text:
            text = ""

        if not options.get("fps"):
            options["fps"] = self.frame_rate

        if isinstance(frame_start_tc, str):
            options["timecode"] = frame_start_tc
        else:
            options["timecode"] = ffmpeg_burnins._frames_to_timecode(
                frame_start_tc,
                self.frame_rate
            )

        self._add_burnin(text, align, options, TIMECODE)

    def add_per_frame_text(
        self,
        text,
        align,
        frame_start,
        frame_end,
        listed_keys,
        options=None
    ):
        """Add text that changes per frame.

        Args:
            text (str): Template string with unfilled keys that are changed
                per frame.
            align (str): Alignment of text.
            frame_start (int): Starting frame for burnins current frame.
            frame_end (int): Ending frame for burnins current frame.
            listed_keys (list): List of keys that are changed per frame.
            options (Optional[dict]): Options to affect style of burnin.
        """

        if not options:
            options = ffmpeg_burnins.TimeCodeOptions(**self.options_init)

        options = options.copy()
        if frame_start is None:
            frame_start = options["frame_offset"]

        # `frame_end` is only for meassurements of text position
        if frame_end is None:
            frame_end = options["frame_end"]

        fps = options.get("fps")
        if not fps:
            fps = self.frame_rate

        text_for_size = text
        if CURRENT_FRAME_SPLITTER in text:
            expr = self._get_current_frame_expression(frame_start, frame_end)
            if expr is None:
                expr = MISSING_KEY_VALUE
                text_for_size = text_for_size.replace(
                    CURRENT_FRAME_SPLITTER, MISSING_KEY_VALUE)
            text = text.replace(CURRENT_FRAME_SPLITTER, expr)

        # Find longest list with values
        longest_list_len = max(
            len(item["values"]) for item in listed_keys.values()
        )
        # Where to store formatted values per frame by key
        new_listed_keys = [{} for _ in range(longest_list_len)]
        # Find the longest value per fill key.
        #   The longest value is used to determine size of burnin box.
        longest_value_by_key = {}
        for key, item in listed_keys.items():
            values = item["values"]
            # Fill the missing values from the longest list with the last
            #   value to make sure all values have same "frame count"
            last_value = values[-1] if values else ""
            for _ in range(longest_list_len - len(values)):
                values.append(last_value)

            # Prepare dictionary structure for nestes values
            # - last key is overriden on each frame loop
            item_keys = list(item["keys"])
            fill_data = {}
            sub_value = fill_data
            last_item_key = item_keys.pop(-1)
            for item_key in item_keys:
                sub_value[item_key] = {}
                sub_value = sub_value[item_key]

            # Fill value per frame
            key_max_len = 0
            key_max_value = ""
            for value, new_values in zip(values, new_listed_keys):
                sub_value[last_item_key] = value
                try:
                    value = key.format(**sub_value)
                except (TypeError, KeyError, ValueError):
                    value = MISSING_KEY_VALUE
                new_values[key] = value

                value_len = len(value)
                if value_len > key_max_len:
                    key_max_value = value
                    key_max_len = value_len

            # Store the longest value
            longest_value_by_key[key] = key_max_value

        # Make sure the longest value of each key is replaced for text size
        #   calculation
        for key, value in longest_value_by_key.items():
            text_for_size = text_for_size.replace(key, value)

        # Create temp file with instructions for each frame of text
        lines = []
        for frame, value in enumerate(new_listed_keys):
            seconds = float(frame) / fps
            # Escape special character
            new_text = text
            for _key, _value in value.items():
                _value = str(_value)
                new_text = new_text.replace(_key, str(_value))

            new_text = (
                str(new_text)
                .replace("\\", "\\\\")
                .replace(",", "\\,")
                .replace(":", "\\:")
            )
            lines.append(
                f"{seconds} drawtext@{align} reinit text='{new_text}';")

        with tempfile.NamedTemporaryFile(
            prefix="ayon_burnin_", mode="w", delete=False,
        ) as temp:
            path = temp.name
            temp.write("\n".join(lines))

        self.cleanup_paths.append(path)
        self.filters["drawtext"].append("sendcmd=f='{}'".format(
            path.replace("\\", "/").replace(":", "\\:")
        ))
        self.add_text(text_for_size, align, frame_start, frame_end, options)

    def _get_current_frame_expression(self, frame_start, frame_end):
        if frame_start is None:
            return None
        return (
            "%{eif:n+" + str(frame_start)
            + ":d:" + str(len(str(frame_end))) + "}"
        )

    def _add_burnin(self, text, align, options, draw):
        """
        Generic method for building the filter flags.
        :param str text: text to apply to the drawtext
        :param enum align: alignment, must use provided enum flags
        :param dict options:
        """

        final_text = text
        text_for_size = text
        if CURRENT_FRAME_SPLITTER in text:
            frame_start = options["frame_offset"]
            frame_end = options.get("frame_end", frame_start)
            expr = self._get_current_frame_expression(frame_start, frame_end)
            if expr is not None:
                max_length = len(str(frame_end))
                # Use number '8' length times for replacement
                size_replacement = max_length * "8"
            else:
                expr = size_replacement = MISSING_KEY_VALUE

            final_text = final_text.replace(
                CURRENT_FRAME_SPLITTER, expr
            )
            text_for_size = text_for_size.replace(
                CURRENT_FRAME_SPLITTER, size_replacement
            )

        timecode_text = options.get("timecode") or ""
        text_for_size += timecode_text

        font_path = options.get("font")
        if not font_path or not os.path.exists(font_path):
            options["font"] = ffmpeg_burnins.FONT

        drawtext_kwargs = get_drawtext_kwargs(
            align,
            self.resolution,
            text_for_size,
            options
        )

        final_text = final_text.replace(",", r"\,").replace(':', r'\:')
        drawtext_kwargs["text"] = f"'{final_text}'"

        args = ":".join(f"{k}={v}" for k, v in drawtext_kwargs.items())
        drawtext = draw % {
            **options,
            "args": args,
            "label": align,
        }

        self.filters["drawtext"].append(drawtext)

    def command(self, output=None, args=None, overwrite=False):
        """
        Generate the entire FFMPEG command.

        :param str output: output file
        :param str args: additional FFMPEG arguments
        :param bool overwrite: overwrite the output if it exists
        :returns: completed command
        :rtype: str
        """
        output = '"{}"'.format(output or '')
        if overwrite:
            output = '-y {}'.format(output)

        filters = ""
        filter_string = self.filter_string
        if filter_string:
            with tempfile.NamedTemporaryFile(
                prefix="ayon_burnin_", mode="w", delete=False
            ) as temp:
                temp.write(filter_string)
                filters_path = temp.name
            filters = '-filter_script:v "{}"'.format(filters_path)
            print("Filters:", filter_string)
            self.cleanup_paths.append(filters_path)

        if self.first_frame is not None:
            start_number_arg = "-start_number {}".format(self.first_frame)
            self.input_args.append(start_number_arg)
            if "start_number" not in args:
                if not args:
                    args = start_number_arg
                else:
                    args = " ".join((start_number_arg, args))

        input_args = ""
        if self.input_args:
            input_args = " {}".format(" ".join(self.input_args))

        return (FFMPEG % {
            'input_args': input_args,
            'input': self.source,
            'output': output,
            'args': '%s ' % args if args else '',
            'filters': filters
        }).strip()

    def render(self, output, args=None, overwrite=False, **kwargs):
        """
        Render the media to a specified destination.

        :param str output: output file
        :param str args: additional FFMPEG arguments
        :param bool overwrite: overwrite the output if it exists
        """
        if not overwrite and os.path.exists(output):
            raise RuntimeError("Destination '%s' exists, please "
                               "use overwrite" % output)

        is_sequence = "%" in output

        command = self.command(
            output=output,
            args=args,
            overwrite=overwrite
        )
        print("Launching command: {}".format(command))

        use_shell = True
        try:
            test_proc = subprocess.Popen(
                f"{FFMPEG_EXE_COMMAND} --help", shell=True
            )
            test_proc.wait()
        except BaseException:
            use_shell = False

        kwargs = {
            "stdout": subprocess.PIPE,
            "stderr": subprocess.PIPE,
            "shell": use_shell,
        }
        proc = subprocess.Popen(command, **kwargs)

        _stdout, _stderr = proc.communicate()
        if _stdout:
            print(_stdout.decode("utf-8", errors="backslashreplace"))

        # This will probably never happen as ffmpeg use stdout
        if _stderr:
            print(_stderr.decode("utf-8", errors="backslashreplace"))

        if proc.returncode != 0:
            raise RuntimeError(
                "Failed to render '{}': {}'".format(output, command)
            )
        if is_sequence:
            output = output % kwargs.get("duration")

        if not os.path.exists(output):
            raise RuntimeError(
                "Failed to generate this f*cking file '%s'" % output
            )

        for path in self.cleanup_paths:
            if os.path.exists(path):
                os.remove(path)