From 764bef64576e3ab3096b78d9cf7910bf8a261e71 Mon Sep 17 00:00:00 2001
From: JMARyA <jmarya@hydrar.de>
Date: Mon, 16 Dec 2024 23:45:15 +0100
Subject: [PATCH] add thumbnail feature

---
 migrations/0004_thumbnails.sql |  7 ++++
 src/library/mod.rs             | 61 ++++++++++++++++++++++++++++++++++
 src/main.rs                    |  1 +
 src/meta.rs                    | 38 +++++++++++++++++++++
 src/pages/assets.rs            | 17 ++++------
 5 files changed, 113 insertions(+), 11 deletions(-)
 create mode 100644 migrations/0004_thumbnails.sql
 create mode 100644 src/meta.rs

diff --git a/migrations/0004_thumbnails.sql b/migrations/0004_thumbnails.sql
new file mode 100644
index 0000000..ed933fd
--- /dev/null
+++ b/migrations/0004_thumbnails.sql
@@ -0,0 +1,7 @@
+
+CREATE TABLE video_thumbnail (
+    "id" UUID NOT NULL PRIMARY KEY,
+    "time" FLOAT NOT NULL,
+    "precomputed" BYTEA NOT NULL,
+    "custom" BYTEA
+);
diff --git a/src/library/mod.rs b/src/library/mod.rs
index cdd7b3a..d0afb8e 100644
--- a/src/library/mod.rs
+++ b/src/library/mod.rs
@@ -2,10 +2,13 @@ use serde_json::json;
 use std::path::Path;
 use std::path::PathBuf;
 use std::str::FromStr;
+use tokio::io::AsyncReadExt;
 use walkdir::WalkDir;
 
 use func::is_video_file;
 pub use video::Video;
+
+use crate::meta;
 mod func;
 pub mod user;
 mod video;
@@ -52,6 +55,64 @@ impl Library {
             .unwrap()
     }
 
+    pub async fn get_thumbnail(&self, video: &Video) -> Option<Vec<u8>> {
+        // get custom thumbnail
+        let res: Option<(Option<Vec<u8>>,)> =
+            sqlx::query_as("SELECT custom FROM video_thumbnail WHERE id = $1")
+                .bind(video.id)
+                .fetch_optional(&self.conn)
+                .await
+                .unwrap();
+
+        if let Some(res) = res {
+            if let Some(data) = res.0 {
+                log::info!("Returned Custom Thumbnail");
+                return Some(data);
+            }
+        } else {
+            let half_time = video.duration / 2.5;
+            let precomputed = meta::extract_video_thumbnail(&video.path, half_time);
+
+            sqlx::query(
+                "INSERT INTO video_thumbnail (\"id\", time, precomputed) VALUES ($1, $2, $3)",
+            )
+            .bind(video.id)
+            .bind(half_time)
+            .bind(precomputed)
+            .execute(&self.conn)
+            .await
+            .unwrap();
+            log::info!("Computed Thumbnail");
+        }
+
+        // get sidecar thumbnail
+        let path = std::path::Path::new(&video.path);
+        let parent = path.parent().unwrap();
+        let thumbnail_path = path.file_stem().unwrap().to_str().unwrap();
+        let thumbnail_path = parent.join(thumbnail_path);
+        let thumbnail_path = thumbnail_path.to_str().unwrap();
+
+        for ext in ["jpg", "jpeg", "png", "avif"] {
+            if let Ok(mut file) = tokio::fs::File::open(format!("{thumbnail_path}.{ext}")).await {
+                let mut content = Vec::with_capacity(1024);
+                file.read_to_end(&mut content).await.unwrap();
+                log::info!("Returned Sidecar Thumbnail");
+                return Some(content);
+            }
+        }
+
+        // get precomputed thumbnail
+        let res: (Vec<u8>,) =
+            sqlx::query_as("SELECT precomputed FROM video_thumbnail WHERE id = $1")
+                .bind(video.id)
+                .fetch_one(&self.conn)
+                .await
+                .unwrap();
+
+        log::info!("Returned Precomputed Thumbnail");
+        return Some(res.0);
+    }
+
     // YT
 
     pub async fn get_channel_name_yt(&self, id: &str) -> String {
diff --git a/src/main.rs b/src/main.rs
index b2bd6aa..1a21083 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,6 +5,7 @@ use rocket::{http::Method, routes};
 use tokio::sync::OnceCell;
 
 mod library;
+mod meta;
 mod pages;
 mod yt_meta;
 
diff --git a/src/meta.rs b/src/meta.rs
new file mode 100644
index 0000000..55febad
--- /dev/null
+++ b/src/meta.rs
@@ -0,0 +1,38 @@
+use std::io::{self, Read};
+use std::process::{Command, Stdio};
+
+pub fn extract_video_thumbnail(path: &str, time: f64) -> Vec<u8> {
+    let output = Command::new("ffmpeg")
+        .args([
+            "-i",
+            path, // Input video file path
+            "-ss",
+            &format!("{}", time), // Seek to the specific time
+            "-vframes",
+            "1", // Extract a single frame
+            "-f",
+            "image2", // Specify the image format
+            "-",      // Output to stdout
+        ])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .output();
+
+    match output {
+        Ok(result) => {
+            if !result.status.success() {
+                log::error!(
+                    "thumbnail_extraction:: ffmpeg error: {}",
+                    String::from_utf8_lossy(&result.stderr)
+                );
+                return Vec::new(); // Return an empty vector on failure
+            }
+            // Return the raw bytes of the image from stdout
+            result.stdout
+        }
+        Err(err) => {
+            log::error!("thumbnail_extraction:: Failed to execute ffmpeg: {}", err);
+            Vec::new()
+        }
+    }
+}
diff --git a/src/pages/assets.rs b/src/pages/assets.rs
index 010e31d..72f6b4d 100644
--- a/src/pages/assets.rs
+++ b/src/pages/assets.rs
@@ -36,23 +36,18 @@ pub async fn video_file(
 }
 
 #[get("/video/thumbnail?<v>")]
-pub async fn video_thumbnail(v: &str, library: &State<Library>) -> Option<NamedFile> {
+pub async fn video_thumbnail(
+    v: &str,
+    library: &State<Library>,
+) -> Option<(Status, (ContentType, Vec<u8>))> {
     let video = if let Some(video) = library.get_video_by_id(v).await {
         video
     } else {
         library.get_video_by_youtube_id(v).await.unwrap()
     };
 
-    let path = std::path::Path::new(&video.path);
-    let parent = path.parent().unwrap();
-    let thumbnail_path = path.file_stem().unwrap().to_str().unwrap();
-    let thumbnail_path = parent.join(thumbnail_path);
-    let thumbnail_path = thumbnail_path.to_str().unwrap();
-
-    for ext in ["jpg", "jpeg", "png", "avif"] {
-        if let Ok(file) = NamedFile::open(format!("{thumbnail_path}.{ext}")).await {
-            return Some(file);
-        }
+    if let Some(data) = library.get_thumbnail(&video).await {
+        return Some((Status::Ok, (ContentType::PNG, data)));
     }
 
     None