From d533c7acfe731ecee067d978cbb1f6722465ca5b Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Thu, 16 Mar 2023 18:13:37 +0700
Subject: [PATCH] refactor re_sync_thumbs task

---
 tubearchivist/home/src/download/thumbnails.py | 94 ++++++++++---------
 tubearchivist/home/tasks.py                   |  4 +-
 2 files changed, 53 insertions(+), 45 deletions(-)

diff --git a/tubearchivist/home/src/download/thumbnails.py b/tubearchivist/home/src/download/thumbnails.py
index b47fd434..ab4743e8 100644
--- a/tubearchivist/home/src/download/thumbnails.py
+++ b/tubearchivist/home/src/download/thumbnails.py
@@ -10,7 +10,6 @@ from io import BytesIO
 from time import sleep
 
 import requests
-from home.src.download import queue  # partial import
 from home.src.es.connect import ElasticWrap, IndexPaginate
 from home.src.ta.config import AppConfig
 from mutagen.mp4 import MP4, MP4Cover
@@ -328,56 +327,65 @@ class ThumbValidator:
 
 
 class ThumbFilesystem:
-    """filesystem tasks for thumbnails"""
+    """sync thumbnail files to media files"""
+
+    INDEX_NAME = "ta_video"
+
+    def __init__(self, task=False):
+        self.task = task
+
+    def embed(self):
+        """entry point"""
+        data = {
+            "query": {"match_all": {}},
+            "_source": ["media_url", "youtube_id"],
+        }
+        paginate = IndexPaginate(
+            index_name=self.INDEX_NAME,
+            data=data,
+            callback=EmbedCallback,
+            task=self.task,
+            total=self._get_total(),
+        )
+        _ = paginate.get_results()
+
+    def _get_total(self):
+        """get total documents in index"""
+        path = f"{self.INDEX_NAME}/_count"
+        response, _ = ElasticWrap(path).get()
+
+        return response.get("count")
+
+
+class EmbedCallback:
+    """callback class to embed thumbnails"""
 
     CONFIG = AppConfig().config
     CACHE_DIR = CONFIG["application"]["cache_dir"]
     MEDIA_DIR = CONFIG["application"]["videos"]
-    VIDEO_DIR = os.path.join(CACHE_DIR, "videos")
+    FORMAT = MP4Cover.FORMAT_JPEG
 
-    def sync(self):
-        """embed thumbnails to mediafiles"""
-        video_list = self.get_thumb_list()
-        self._embed_thumbs(video_list)
+    def __init__(self, source, index_name):
+        self.source = source
+        self.index_name = index_name
 
-    def get_thumb_list(self):
-        """get list of mediafiles and matching thumbnails"""
-        pending = queue.PendingList()
-        pending.get_download()
-        pending.get_indexed()
-
-        video_list = []
-        for video in pending.all_videos:
-            video_id = video["youtube_id"]
-            media_url = os.path.join(self.MEDIA_DIR, video["media_url"])
+    def run(self):
+        """run embed"""
+        for video in self.source:
+            video_id = video["_source"]["youtube_id"]
+            media_url = os.path.join(
+                self.MEDIA_DIR, video["_source"]["media_url"]
+            )
             thumb_path = os.path.join(
                 self.CACHE_DIR, ThumbManager(video_id).vid_thumb_path()
             )
-            video_list.append(
-                {
-                    "media_url": media_url,
-                    "thumb_path": thumb_path,
-                }
-            )
+            if os.path.exists(thumb_path):
+                self.embed(media_url, thumb_path)
 
-        return video_list
+    def embed(self, media_url, thumb_path):
+        """embed thumb in single media file"""
+        video = MP4(media_url)
+        with open(thumb_path, "rb") as f:
+            video["covr"] = [MP4Cover(f.read(), imageformat=self.FORMAT)]
 
-    @staticmethod
-    def _embed_thumbs(video_list):
-        """rewrite the thumbnail into media file"""
-
-        counter = 1
-        for video in video_list:
-            # loop through all videos
-            media_url = video["media_url"]
-            thumb_path = video["thumb_path"]
-
-            mutagen_vid = MP4(media_url)
-            with open(thumb_path, "rb") as f:
-                mutagen_vid["covr"] = [
-                    MP4Cover(f.read(), imageformat=MP4Cover.FORMAT_JPEG)
-                ]
-            mutagen_vid.save()
-            if counter % 50 == 0:
-                print(f"thumbnail write progress {counter}/{len(video_list)}")
-            counter = counter + 1
+        video.save()
diff --git a/tubearchivist/home/tasks.py b/tubearchivist/home/tasks.py
index 20961af1..8c99181d 100644
--- a/tubearchivist/home/tasks.py
+++ b/tubearchivist/home/tasks.py
@@ -297,7 +297,7 @@ def thumbnail_check(self):
     ThumbValidator(task=self).validate()
 
 
-@shared_task(bind=True, name="resync_thumbs")
+@shared_task(bind=True, name="resync_thumbs", base=BaseTask)
 def re_sync_thumbs(self):
     """sync thumbnails to mediafiles"""
     manager = TaskManager()
@@ -306,7 +306,7 @@ def re_sync_thumbs(self):
         return
 
     manager.init(self)
-    ThumbFilesystem().sync()
+    ThumbFilesystem(task=self).embed()
 
 
 @shared_task(name="subscribe_to")