From 744780f4bd10405e1e4109324d9dfc15911e049f Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Thu, 7 Oct 2021 22:28:32 +0700
Subject: [PATCH] add optional EmbedThumbnail postprocessor

---
 tubearchivist/home/config.json                |  3 +-
 tubearchivist/home/src/download.py            | 35 +++++++++++++++----
 .../home/templates/home/settings.html         |  9 +++++
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/tubearchivist/home/config.json b/tubearchivist/home/config.json
index 3d3c50bb..ce0867b1 100644
--- a/tubearchivist/home/config.json
+++ b/tubearchivist/home/config.json
@@ -18,7 +18,8 @@
         "limit_speed": false,
         "sleep_interval": 3,
         "format": false,
-        "add_metadata": false
+        "add_metadata": false,
+        "add_thumbnail": false
     },
     "application": {
         "cache_dir": "/cache",
diff --git a/tubearchivist/home/src/download.py b/tubearchivist/home/src/download.py
index 6b217f5d..f34c1718 100644
--- a/tubearchivist/home/src/download.py
+++ b/tubearchivist/home/src/download.py
@@ -470,8 +470,8 @@ class VideoDownloader:
         }
         RedisArchivist().set_message("progress:download", mess_dict)
 
-    def dl_single_vid(self, youtube_id):
-        """download single video"""
+    def build_obs(self):
+        """build obs dictionary for yt-dlp"""
         obs = {
             "default_search": "ytsearch",
             "merge_output_format": "mp4",
@@ -485,6 +485,7 @@ class VideoDownloader:
             "quiet": True,
             "continuedl": True,
             "retries": 3,
+            "writethumbnail": False,
         }
         if self.config["downloads"]["format"]:
             obs["format"] = self.config["downloads"]["format"]
@@ -505,15 +506,29 @@ class VideoDownloader:
                 }
             )
 
+        if self.config["downloads"]["add_thumbnail"]:
+            postprocessors.append(
+                {
+                    "key": "EmbedThumbnail",
+                    "already_have_thumbnail": True,
+                }
+            )
+            obs["writethumbnail"] = True
+
         obs["postprocessors"] = postprocessors
 
+        return obs
+
+    def dl_single_vid(self, youtube_id):
+        """download single video"""
+        dl_cache = self.config["application"]["cache_dir"] + "/download/"
+        obs = self.build_obs()
+
         # check if already in cache to continue from there
-        cache_dir = self.config["application"]["cache_dir"]
-        cached = os.listdir(cache_dir + "/download/")
-        all_cached = ignore_filelist(cached)
+        all_cached = ignore_filelist(os.listdir(dl_cache))
         for file_name in all_cached:
             if youtube_id in file_name:
-                obs["outtmpl"] = cache_dir + "/download/" + file_name
+                obs["outtmpl"] = os.path.join(dl_cache, file_name)
         with youtube_dl.YoutubeDL(obs) as ydl:
             try:
                 ydl.download([youtube_id])
@@ -522,6 +537,14 @@ class VideoDownloader:
                 sleep(10)
                 ydl.download([youtube_id])
 
+        if obs["writethumbnail"]:
+            # webp files don't get cleaned up automatically
+            all_cached = ignore_filelist(os.listdir(dl_cache))
+            to_clean = [i for i in all_cached if not i.endswith(".mp4")]
+            for file_name in to_clean:
+                file_path = os.path.join(dl_cache, file_name)
+                os.remove(file_path)
+
     def move_to_archive(self, vid_dict):
         """move downloaded video from cache to archive"""
         videos = self.config["application"]["videos"]
diff --git a/tubearchivist/home/templates/home/settings.html b/tubearchivist/home/templates/home/settings.html
index c695b9d2..99458d3f 100644
--- a/tubearchivist/home/templates/home/settings.html
+++ b/tubearchivist/home/templates/home/settings.html
@@ -88,6 +88,15 @@
                 <option value="1">embed metadata</option>
             </select>
         </div>
+        <div class="settings-item">
+            <p>Current thumbnail embed setting: <span class="settings-current">{{ config.downloads.add_thumbnail }}</span></p>
+            <i>Embed thumbnail into the mediafile.</i><br>
+            <select name="downloads.add_thumbnail" id="downloads.add_thumbnail"">
+                <option value="" disabled selected> -- change thumbnail embed -- </option>
+                <option value="0">don't embed thumbnail</option>
+                <option value="1">embed thumbnail</option>
+            </select>
+        </div>
     </div>
     <button type="submit">Update Settings</button>
 </form>