1# -*- coding: utf-8 -*-
2
3# Copyright 2018-2021 Mike Fährmann
4#
5# This program is free software; you can redistribute it and/or modify
6# it under the terms of the GNU General Public License version 2 as
7# published by the Free Software Foundation.
8
9"""Convert Pixiv Ugoira to WebM"""
10
11from .common import PostProcessor
12from .. import util
13import collections
14import subprocess
15import tempfile
16import zipfile
17import os
18
19
20class UgoiraPP(PostProcessor):
21
22    def __init__(self, job, options):
23        PostProcessor.__init__(self, job)
24        self.extension = options.get("extension") or "webm"
25        self.args = options.get("ffmpeg-args") or ()
26        self.twopass = options.get("ffmpeg-twopass", False)
27        self.output = options.get("ffmpeg-output", True)
28        self.delete = not options.get("keep-files", False)
29        self.repeat = options.get("repeat-last-frame", True)
30
31        ffmpeg = options.get("ffmpeg-location")
32        self.ffmpeg = util.expand_path(ffmpeg) if ffmpeg else "ffmpeg"
33
34        rate = options.get("framerate", "auto")
35        if rate != "auto":
36            self.calculate_framerate = lambda _: (None, rate)
37
38        if options.get("ffmpeg-demuxer") == "image2":
39            self._process = self._image2
40        else:
41            self._process = self._concat
42
43        if options.get("libx264-prevent-odd", True):
44            # get last video-codec argument
45            vcodec = None
46            for index, arg in enumerate(self.args):
47                arg, _, stream = arg.partition(":")
48                if arg == "-vcodec" or arg in ("-c", "-codec") and (
49                        not stream or stream.partition(":")[0] in ("v", "V")):
50                    vcodec = self.args[index + 1]
51            # use filter when using libx264/5
52            self.prevent_odd = (
53                vcodec in ("libx264", "libx265") or
54                not vcodec and self.extension.lower() in ("mp4", "mkv"))
55        else:
56            self.prevent_odd = False
57
58        job.register_hooks(
59            {"prepare": self.prepare, "file": self.convert}, options)
60
61    def prepare(self, pathfmt):
62        self._frames = None
63
64        if pathfmt.extension != "zip":
65            return
66
67        if "frames" in pathfmt.kwdict:
68            self._frames = pathfmt.kwdict["frames"]
69        elif "pixiv_ugoira_frame_data" in pathfmt.kwdict:
70            self._frames = pathfmt.kwdict["pixiv_ugoira_frame_data"]["data"]
71        else:
72            return
73
74        if self.delete:
75            pathfmt.set_extension(self.extension)
76
77    def convert(self, pathfmt):
78        if not self._frames:
79            return
80
81        with tempfile.TemporaryDirectory() as tempdir:
82            # extract frames
83            try:
84                with zipfile.ZipFile(pathfmt.temppath) as zfile:
85                    zfile.extractall(tempdir)
86            except FileNotFoundError:
87                pathfmt.realpath = pathfmt.temppath
88                return
89
90            # process frames and collect command-line arguments
91            args = self._process(tempdir)
92            if self.args:
93                args += self.args
94            self.log.debug("ffmpeg args: %s", args)
95
96            # invoke ffmpeg
97            pathfmt.set_extension(self.extension)
98            try:
99                if self.twopass:
100                    if "-f" not in self.args:
101                        args += ("-f", self.extension)
102                    args += ("-passlogfile", tempdir + "/ffmpeg2pass", "-pass")
103                    self._exec(args + ["1", "-y", os.devnull])
104                    self._exec(args + ["2", pathfmt.realpath])
105                else:
106                    args.append(pathfmt.realpath)
107                    self._exec(args)
108            except OSError as exc:
109                print()
110                self.log.error("Unable to invoke FFmpeg (%s: %s)",
111                               exc.__class__.__name__, exc)
112                pathfmt.realpath = pathfmt.temppath
113            else:
114                if self.delete:
115                    pathfmt.delete = True
116                else:
117                    pathfmt.set_extension("zip")
118
119    def _concat(self, path):
120        ffconcat = path + "/ffconcat.txt"
121
122        content = ["ffconcat version 1.0"]
123        append = content.append
124        for frame in self._frames:
125            append("file '{}'\nduration {}".format(
126                frame["file"], frame["delay"] / 1000))
127        if self.repeat:
128            append("file '{}'".format(frame["file"]))
129        append("")
130
131        with open(ffconcat, "w") as file:
132            file.write("\n".join(content))
133
134        rate_in, rate_out = self.calculate_framerate(self._frames)
135        args = [self.ffmpeg, "-f", "concat"]
136        if rate_in:
137            args += ("-r", str(rate_in))
138        args += ("-i", ffconcat)
139        if rate_out:
140            args += ("-r", str(rate_out))
141        return args
142
143    def _image2(self, path):
144        path += "/"
145
146        # adjust frame mtime values
147        ts = 0
148        for frame in self._frames:
149            os.utime(path + frame["file"], ns=(ts, ts))
150            ts += frame["delay"] * 1000000
151
152        return [
153            self.ffmpeg,
154            "-f", "image2",
155            "-ts_from_file", "2",
156            "-pattern_type", "sequence",
157            "-i", "{}%06d.{}".format(
158                path.replace("%", "%%"), frame["file"].rpartition(".")[2]),
159        ]
160
161    def _exec(self, args):
162        out = None if self.output else subprocess.DEVNULL
163        return subprocess.Popen(args, stdout=out, stderr=out).wait()
164
165    @staticmethod
166    def calculate_framerate(framelist):
167        counter = collections.Counter(frame["delay"] for frame in framelist)
168        fps = "1000/{}".format(min(counter))
169        return (fps, None) if len(counter) == 1 else (None, fps)
170
171
172__postprocessor__ = UgoiraPP
173