From 0e6146c6e8c0ffbfacfe65a02e6b90b7a0cd8c3c Mon Sep 17 00:00:00 2001
From: Aroy-Art <Aroy-Art@pm.me>
Date: Mon, 24 Feb 2025 22:33:59 +0100
Subject: [PATCH] Add: celery task for generating video thumbnails

---
 backend/apps/files/tasks.py | 79 +++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/backend/apps/files/tasks.py b/backend/apps/files/tasks.py
index d0e23d8..5228237 100644
--- a/backend/apps/files/tasks.py
+++ b/backend/apps/files/tasks.py
@@ -1,5 +1,7 @@
+import os
 import io
 import hashlib
+import subprocess
 
 from django.db import transaction
 
@@ -10,6 +12,8 @@ import blurhash
 
 from .models import PostFileModel
 
+from utils.hash import compute_file_hash_blake3, compute_blur_hash
+
 
 @shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5)
 def generate_blur_hash_PostFile(file_id):
@@ -52,3 +56,78 @@ def generate_md5_hash_PostFile(file_id):
         print(error_message)
         raise Retry(exc=e)  # Retry on exception
         return error_message  # This ensures the error message is stored in the results
+
+
+@shared_task(name="generate_video_thumbnail")
+def generate_video_thumbnail(
+    file_id: int, size: int = 0, timestamp=None, movie_strip: bool = False
+):
+    """
+    Generate video thumbnails using ffmpegthumbnailer and update the PostFileModel instance.
+
+    Args:
+        file_id (int): ID of the PostFileModel instance
+        size (int): Desired thumbnail width or height defulats to video size
+        timestamp (float): Timestamp(s) in seconds where the thumbnail should be extracted
+        movie_strip (bool): Create a movie strip overlay
+
+    Returns:
+        str: Success message or error message
+    """
+    try:
+        with transaction.atomic():
+            # Retrieve the PostFileModel instance with a lock
+            pstfile = PostFileModel.objects.select_for_update().get(id=file_id)
+
+            if not pstfile.file:
+                return "Error: Video file not found for the given file_id."
+
+            video_path = pstfile.file.path
+
+            # Create output directory if it doesn't exist
+            output_dir = os.path.join("/tmp/thumbgen/", str(file_id))
+            os.makedirs(output_dir, exist_ok=True)
+
+            thumbnail_filename = f"thumbnail_{pstfile.hash_blake3}.png"
+
+            thumbnail_file_path = os.path.join(output_dir, thumbnail_filename)
+
+            cmd = [
+                "ffmpegthumbnailer",
+                "-i",
+                video_path,
+                "-o",
+                thumbnail_file_path,
+                "-s",
+                str(size),
+                "-m",
+            ]
+
+            if movie_strip:
+                cmd.extend(["-f"])
+
+            # Generate thumbnail at specified timestamps
+            if timestamp is not None:
+                cmd.extend(["-t", f"{timestamp}"])
+
+            subprocess.run(cmd, check=True)
+
+            # Update the PostFileModel's thumbnail field with the new file
+            with open(thumbnail_file_path, "rb") as file:
+                pstfile.thumbnail.save(thumbnail_filename, file)
+
+            pstfile.thumbnail_hash_blake3 = compute_file_hash_blake3(
+                thumbnail_file_path
+            )
+
+            print(compute_blur_hash(thumbnail_file_path))
+            pstfile.thumbnail_blur_hash = compute_blur_hash(thumbnail_file_path)
+
+            pstfile.save()
+
+        return f"Thumbnail generated and saved to {thumbnail_file_path}"
+
+    except subprocess.CalledProcessError as e:
+        return f"Error generating thumbnail: {str(e)}"
+    except Exception as e:
+        return f"Unexpected error: {str(e)}"