From e0ea9a59ed8f9187557c9e551b8994b84af955dc Mon Sep 17 00:00:00 2001
From: Jesse Bannon <jbann1994@gmail.com>
Date: Sun, 23 Jul 2023 23:35:11 -0700
Subject: [PATCH] [FEATURE] Dedicated embed_thumbnail plugin

---
 src/ytdl_sub/config/preset_class_mappings.py |  2 +
 src/ytdl_sub/plugins/embed_thumbnail.py      | 81 ++++++++++++++++++++
 src/ytdl_sub/plugins/music_tags.py           |  5 ++
 tests/e2e/plugins/test_audio_extract.py      |  2 +
 4 files changed, 90 insertions(+)
 create mode 100644 src/ytdl_sub/plugins/embed_thumbnail.py

diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py
index dac9e996..a77e4988 100644
--- a/src/ytdl_sub/config/preset_class_mappings.py
+++ b/src/ytdl_sub/config/preset_class_mappings.py
@@ -8,6 +8,7 @@ from ytdl_sub.downloaders.url.url import UrlDownloader
 from ytdl_sub.plugins.audio_extract import AudioExtractPlugin
 from ytdl_sub.plugins.chapters import ChaptersPlugin
 from ytdl_sub.plugins.date_range import DateRangePlugin
+from ytdl_sub.plugins.embed_thumbnail import EmbedThumbnailPlugin
 from ytdl_sub.plugins.file_convert import FileConvertPlugin
 from ytdl_sub.plugins.internal.view import ViewPlugin
 from ytdl_sub.plugins.match_filters import MatchFiltersPlugin
@@ -107,6 +108,7 @@ class PluginMapping:
         "_view": ViewPlugin,
         "audio_extract": AudioExtractPlugin,
         "date_range": DateRangePlugin,
+        "embed_thumbnail": EmbedThumbnailPlugin,
         "file_convert": FileConvertPlugin,
         "match_filters": MatchFiltersPlugin,
         "music_tags": MusicTagsPlugin,
diff --git a/src/ytdl_sub/plugins/embed_thumbnail.py b/src/ytdl_sub/plugins/embed_thumbnail.py
new file mode 100644
index 00000000..2c448d78
--- /dev/null
+++ b/src/ytdl_sub/plugins/embed_thumbnail.py
@@ -0,0 +1,81 @@
+from typing import Optional
+
+from ytdl_sub.config.preset_options import OptionsValidator
+from ytdl_sub.entries.entry import Entry
+from ytdl_sub.plugins.plugin import Plugin
+from ytdl_sub.utils.ffmpeg import FFMPEG
+from ytdl_sub.utils.file_handler import FileHandler
+from ytdl_sub.utils.file_handler import FileMetadata
+from ytdl_sub.utils.logger import Logger
+from ytdl_sub.utils.thumbnail import convert_download_thumbnail
+from ytdl_sub.validators.validators import BoolValidator
+
+logger = Logger.get("embed_thumbnail")
+
+
+class EmbedThumbnailOptions(BoolValidator, OptionsValidator):
+    """
+    Whether to embed thumbnails to the audio/video file or not.
+
+    Usage:
+
+    .. code-block:: yaml
+
+       presets:
+         my_example_preset:
+           embed_thumbnail: True
+    """
+
+
+class EmbedThumbnailPlugin(Plugin[EmbedThumbnailOptions]):
+    plugin_options_type = EmbedThumbnailOptions
+
+    @property
+    def _embed_thumbnail(self) -> bool:
+        return self.plugin_options.value
+
+    def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]:
+        """
+        Maybe embed the thumbnail
+        """
+        if not self._embed_thumbnail:
+            return None
+
+        if entry.ext == "webm":
+            logger.warning("webm does not support embedded thumbnails, skipping")
+            return None
+
+        if not self.is_dry_run:
+            # convert the entry thumbnail so it is embedded as jpg
+            convert_download_thumbnail(entry=entry)
+
+            file_path = entry.get_download_file_path()
+            thumbnail_path = entry.get_download_thumbnail_path()
+            tmp_file_path = FFMPEG.tmp_file_path(file_path)
+            try:
+                FFMPEG.run(
+                    [
+                        "-i",
+                        file_path,
+                        "-i",
+                        thumbnail_path,
+                        "-map",
+                        "0",
+                        "-map",
+                        "1",
+                        "-dn",  # ignore data streams
+                        "-c",
+                        "-copy",
+                        "-c:v:1",
+                        entry.thumbnail_ext,
+                        "-disposition:v:1",
+                        "attached_pic",
+                        "-bitexact",  # for reproducibility
+                        tmp_file_path,
+                    ]
+                )
+                FileHandler.move(tmp_file_path, file_path)
+            finally:
+                FileHandler.delete(tmp_file_path)
+
+        return FileMetadata("Embedded thumbnail")
diff --git a/src/ytdl_sub/plugins/music_tags.py b/src/ytdl_sub/plugins/music_tags.py
index 97280bd6..722128ba 100644
--- a/src/ytdl_sub/plugins/music_tags.py
+++ b/src/ytdl_sub/plugins/music_tags.py
@@ -149,6 +149,11 @@ class MusicTagsPlugin(Plugin[MusicTagsOptions]):
                     setattr(audio_file, tag_name, tag_value[0])
 
             if self.plugin_options.embed_thumbnail:
+                logger.warning(
+                    "music_tags.embed_thumbnail is now deprecated. Use the dedicated "
+                    "embed_thumbnail plugin instead"
+                )
+
                 # convert the entry thumbnail so it is embedded as jpg
                 convert_download_thumbnail(entry=entry)
 
diff --git a/tests/e2e/plugins/test_audio_extract.py b/tests/e2e/plugins/test_audio_extract.py
index 1a3563c7..48aeb9d0 100644
--- a/tests/e2e/plugins/test_audio_extract.py
+++ b/tests/e2e/plugins/test_audio_extract.py
@@ -11,6 +11,8 @@ def single_song_preset_dict(output_directory):
         "preset": "single",
         # test multi-tags
         "music_tags": {"embed_thumbnail": True, "tags": {"genres": ["multi_tag_1", "multi_tag_2"]}},
+        # test the new embed_thumbnail plugin
+        "embed_thumbnail": True,
         # download the worst format so it is fast
         "ytdl_options": {
             "format": "worst[ext=mp4]",