From eeba604d11e61bec77e09d53d64804dcb53c55e7 Mon Sep 17 00:00:00 2001
From: louis <louis@poweris.moe>
Date: Sun, 30 May 2021 19:36:17 -0400
Subject: [PATCH] yt_common: chapters and audio mixins

---
 yt_common/yt_common/__init__.py |  2 +-
 yt_common/yt_common/chapters.py | 63 +++++++++++++++++++++++++++++++++
 yt_common/yt_common/config.py   | 53 +++++++++++++++++++++++++--
 yt_common/yt_common/source.py   |  5 +--
 4 files changed, 117 insertions(+), 6 deletions(-)
 create mode 100644 yt_common/yt_common/chapters.py

diff --git a/yt_common/yt_common/__init__.py b/yt_common/yt_common/__init__.py
index 6b683c2..57a5874 100644
--- a/yt_common/yt_common/__init__.py
+++ b/yt_common/yt_common/__init__.py
@@ -1 +1 @@
-from . import antialiasing, automation, config, data, deband, logging, source  # noqa: F401
+from . import antialiasing, automation, chapters, config, data, deband, logging, scale, source  # noqa: F401
diff --git a/yt_common/yt_common/chapters.py b/yt_common/yt_common/chapters.py
new file mode 100644
index 0000000..3d6a8e4
--- /dev/null
+++ b/yt_common/yt_common/chapters.py
@@ -0,0 +1,63 @@
+from lxml import etree
+from random import getrandbits
+
+from typing import Dict, List, NamedTuple, Set
+
+LANGMAP: Dict[str, str] = {
+    "eng": "en",
+    "und": "und",
+    "jpn": "ja",
+}
+
+
+class Chapter(NamedTuple):
+    title: str
+    frame: int
+    lang: str = "eng"
+
+
+class RandMan:
+    used: Set[int]
+
+    def __init__(self) -> None:
+        self.used = set()
+
+    def get_rand(self, bits: int = 64) -> str:
+        r = getrandbits(bits)
+        while r in self.used:
+            r = getrandbits(bits)
+        self.used.add(r)
+        return str(r)
+
+
+def timecode_to_timestamp(stamp: float) -> str:
+    m = int(stamp // 60)
+    stamp %= 60
+    h = int(m // 60)
+    m %= 60
+    return f"{h:02d}:{m:02d}:{stamp:06.3f}000000"
+
+
+def make_chapters(chapters: List[Chapter], timecodes: List[float], outfile: str) -> None:
+    chapters.sort(key=lambda c: c.frame)
+    rand = RandMan()
+
+    root = etree.Element("Chapters")
+    ed = etree.SubElement(root, "EditionEntry")
+    etree.SubElement(ed, "EditionUID").text = rand.get_rand()
+
+    for i, c in enumerate(chapters):
+        start = timecode_to_timestamp(timecodes[c.frame])
+        end = timecode_to_timestamp(timecodes[chapters[i+1].frame]) if i < len(chapters) - 1 else None
+        atom = etree.SubElement(ed, "ChapterAtom")
+        etree.SubElement(atom, "ChapterTimeStart").text = start
+        if end is not None:
+            etree.SubElement(atom, "ChapterTimeEnd").text = end
+        disp = etree.SubElement(atom, "ChapterDisplay")
+        etree.SubElement(disp, "ChapterString").text = c.title
+        etree.SubElement(disp, "ChapLanguageIETF").text = LANGMAP[c.lang]
+        etree.SubElement(disp, "ChapterLanguage").text = c.lang
+        etree.SubElement(atom, "ChapterUID").text = rand.get_rand()
+
+    with open(outfile, "wb") as o:
+        o.write(etree.tostring(root, encoding="utf-8", xml_declaration=True, pretty_print=True))
diff --git a/yt_common/yt_common/config.py b/yt_common/yt_common/config.py
index 0369f3d..d741eec 100644
--- a/yt_common/yt_common/config.py
+++ b/yt_common/yt_common/config.py
@@ -1,4 +1,8 @@
-from typing import Union
+import os
+import subprocess
+
+from abc import ABC, abstractmethod
+from typing import List, Union
 
 
 class Config():
@@ -16,5 +20,48 @@ class Config():
         self.datapath = datapath
 
     def format_filename(self, filename: str) -> str:
-        return filename.format(epnum=self.desc, title=self.title,
-                               title_long=self.title_long, resolution=self.resolution)
+        fname = filename.format(epnum=self.desc, title=self.title,
+                                title_long=self.title_long, resolution=self.resolution)
+        return os.path.join(f"../{self.desc}/", fname)
+
+    def encode_audio(self, afile: str) -> str:
+        return afile  # default: passthrough
+
+
+class AudioEncoder(ABC):
+    @abstractmethod
+    def encode_audio(self, afile: str) -> str:
+        pass
+
+
+class FFAudio(AudioEncoder):
+    def encode_audio(self, afile: str) -> str:
+        ffmpeg_args = [
+            "ffmpeg",
+            "-hide_banner", "-loglevel", "panic",
+            "-i", afile,
+            "-y",
+            "-map", "0:a",
+        ] + self.codec_args() + ["_ffaudio_encode.mka"]
+        print("+ " + " ".join(ffmpeg_args))
+        subprocess.call(ffmpeg_args)
+        return "_ffaudio_encode.mka"
+
+    @abstractmethod
+    def codec_args(self) -> List[str]:
+        pass
+
+
+class OpusMixin(FFAudio):
+    def codec_args(self) -> List[str]:
+        return ["-c:a", "libopus", "-b:a", "192k", "-sample_fmt", "s16"]
+
+
+class FdkAacMixin(FFAudio):
+    def codec_args(self) -> List[str]:
+        return ["-c:a", "libfdk_aac", "-b:a", "256k", "-sample_fmt", "s16"]
+
+
+class FlacMixin(FFAudio):
+    def codec_args(self) -> List[str]:
+        return ["-c:a", "flac"]
diff --git a/yt_common/yt_common/source.py b/yt_common/yt_common/source.py
index 6c1adcc..4cdc974 100644
--- a/yt_common/yt_common/source.py
+++ b/yt_common/yt_common/source.py
@@ -78,8 +78,9 @@ class FileSource(ABC):
 class SimpleSource(FileSource):
     src: List[FileTrim]
 
-    def __init__(self, src: Union[FileTrim, List[FileTrim]]) -> None:
-        self.src = src if isinstance(src, list) else [src]
+    def __init__(self, src: Union[str, List[str], FileTrim, List[FileTrim]]) -> None:
+        srcl = src if isinstance(src, list) else [src]
+        self.src = [FileTrim(s, (None, None)) if isinstance(s, str) else s for s in srcl]
 
     def get_audio(self) -> List[FileTrim]:
         return self.src