thumbnail generation

7 files changed, 164 insertions, 26 deletions

diff --git a/client/src/lib.rs b/client/src/lib.rs
index f88ecc5..e751958 100644
--- a/client/src/lib.rs
+++ b/client/src/lib.rs
@@ -63,6 +63,9 @@ pub struct Session {
     session_token: String,
 }
 
+pub trait UnpinWrite: tokio::io::AsyncWrite + std::marker::Unpin {}
+impl<T: tokio::io::AsyncWrite + std::marker::Unpin> UnpinWrite for T {}
+
 impl Session {
     fn session_param(&self) -> String {
         format!("session={}", self.session_token)
@@ -79,24 +82,65 @@ impl Session {
             .await?)
     }
 
-    // TODO use AssetRole instead of str
     pub async fn node_asset(
         &self,
+        writer: impl UnpinWrite,
         id: &str,
         role: AssetRole,
         width: usize,
-        mut writer: impl tokio::io::AsyncWrite + std::marker::Unpin,
     ) -> Result<()> {
         debug!("downloading asset {role:?} for {id:?}");
-        let mut r = self
-            .client
-            .get(format!(
+        self.download_url(
+            writer,
+            format!(
                 "{}/n/{id}/asset?role={}&width={width}",
                 self.instance.base(),
                 role.as_str()
-            ))
-            .send()
-            .await?;
+            ),
+        )
+        .await
+    }
+
+    pub async fn node_thumbnail(
+        &self,
+        writer: impl UnpinWrite,
+        id: &str,
+        width: usize,
+        time: f64,
+    ) -> Result<()> {
+        debug!("downloading thumbnail for {id:?} at {time}s");
+        self.download_url(
+            writer,
+            format!(
+                "{}/n/{id}/thumbnail?t={time}&width={width}",
+                self.instance.base(),
+            ),
+        )
+        .await
+    }
+
+    pub async fn stream(
+        &self,
+        writer: impl UnpinWrite,
+        id: &str,
+        stream_spec: &StreamSpec,
+    ) -> Result<()> {
+        self.download_url(writer, self.stream_url(id, stream_spec))
+            .await
+    }
+
+    pub fn stream_url(&self, id: &str, stream_spec: &StreamSpec) -> String {
+        format!(
+            "{}/n/{}/stream?{}&{}",
+            self.instance.base(),
+            id,
+            stream_spec.to_query(),
+            self.session_param()
+        )
+    }
+
+    pub async fn download_url(&self, mut writer: impl UnpinWrite, url: String) -> Result<()> {
+        let mut r = self.client.get(url).send().await?;
         while let Some(chunk) = r.chunk().await? {
             writer.write_all(&chunk).await?;
         }
@@ -113,14 +157,4 @@ impl Session {
         info!("done");
         Ok(())
     }
-
-    pub fn stream(&self, id: &str, stream_spec: &StreamSpec) -> String {
-        format!(
-            "{}/n/{}/stream?{}&{}",
-            self.instance.base(),
-            id,
-            stream_spec.to_query(),
-            self.session_param()
-        )
-    }
 }
diff --git a/import/src/lib.rs b/import/src/lib.rs
index 4776bd2..692cf7f 100644
--- a/import/src/lib.rs
+++ b/import/src/lib.rs
@@ -572,7 +572,7 @@ async fn cache_federation_asset(
         move |out| async move {
             let session = session;
             session
-                .node_asset(identifier.as_str(), role, 1024, out)
+                .node_asset(out, identifier.as_str(), role, 1024)
                 .await
         },
     )
diff --git a/server/src/routes/mod.rs b/server/src/routes/mod.rs
index c3299d3..47fd6d2 100644
--- a/server/src/routes/mod.rs
+++ b/server/src/routes/mod.rs
@@ -27,7 +27,7 @@ use ui::{
         r_admin_remove_invite,
         user::{r_admin_remove_user, r_admin_user, r_admin_user_permission, r_admin_users},
     },
-    assets::r_item_assets,
+    assets::{r_item_assets, r_node_thumbnail},
     browser::r_all_items_filter,
     error::{r_api_catch, r_catch},
     home::{r_home, r_home_unpriv},
@@ -94,6 +94,7 @@ pub fn build_rocket(database: Database, federation: Federation) -> Rocket<Build>
                 r_assets_js_map,
                 r_stream,
                 r_node_userdata,
+                r_node_thumbnail,
                 r_player,
                 r_player_progress,
                 r_player_watched,
diff --git a/server/src/routes/stream.rs b/server/src/routes/stream.rs
index 8d16ded..e8b14b5 100644
--- a/server/src/routes/stream.rs
+++ b/server/src/routes/stream.rs
@@ -96,7 +96,7 @@ pub async fn r_stream(
                 })
                 .await?;
 
-            let uri = session.stream(
+            let uri = session.stream_url(
                 node.public.id.as_ref().unwrap(),
                 &StreamSpec {
                     tracks: vec![remote_index],
diff --git a/server/src/routes/ui/assets.rs b/server/src/routes/ui/assets.rs
index f00c416..7df7cf0 100644
--- a/server/src/routes/ui/assets.rs
+++ b/server/src/routes/ui/assets.rs
@@ -5,12 +5,15 @@
 */
 use crate::{
     database::Database,
-    routes::ui::{account::session::Session, error::MyError, CacheControlFile},
+    routes::ui::{account::session::Session, error::MyResult, CacheControlFile},
 };
 use anyhow::{anyhow, Context};
-use jellybase::{permission::NodePermissionExt, AssetLocationExt};
-use jellycommon::AssetLocation;
+use jellybase::{
+    cache::async_cache_file, federation::Federation, permission::NodePermissionExt,
+    AssetLocationExt,
+};
 pub use jellycommon::AssetRole;
+use jellycommon::{AssetLocation, LocalTrack, SourceTrackKind, TrackSource};
 use log::info;
 use rocket::{get, http::ContentType, State};
 use std::{path::PathBuf, str::FromStr};
@@ -23,7 +26,7 @@ pub async fn r_item_assets(
     id: &str,
     role: AssetRole,
     width: Option<usize>,
-) -> Result<(ContentType, CacheControlFile), MyError> {
+) -> MyResult<(ContentType, CacheControlFile)> {
     let node = db
         .node
         .get(&id.to_string())?
@@ -45,8 +48,72 @@ pub async fn r_item_assets(
     let asset = asset.unwrap_or(AssetLocation::Assets(
         PathBuf::from_str("fallback.avif").unwrap(),
     ));
+    Ok(asset_with_res(asset, width).await?)
+}
+
+#[get("/n/<id>/thumbnail?<t>&<width>")]
+pub async fn r_node_thumbnail(
+    session: Session,
+    db: &State<Database>,
+    fed: &State<Federation>,
+    id: &str,
+    t: f64,
+    width: Option<usize>,
+) -> MyResult<(ContentType, CacheControlFile)> {
+    let node = db
+        .node
+        .get(&id.to_string())?
+        .only_if_permitted(&session.user.permissions)
+        .ok_or(anyhow!("node does not exist"))?;
+
+    let media = node.public.media.ok_or(anyhow!("no media"))?;
+    let (thumb_track_index, thumb_track) = media
+        .tracks
+        .iter()
+        .enumerate()
+        .find(|(_i, t)| matches!(t.kind, SourceTrackKind::Video { .. }))
+        .ok_or(anyhow!("no video track to create a thumbnail of"))?;
+    let source = node.private.source.ok_or(anyhow!("no source"))?;
+    let thumb_track_source = &source[thumb_track_index];
+
+    if t < 0. || t > media.duration {
+        Err(anyhow!("thumbnail instant not within media duration"))?
+    }
+
+    let step = 8.;
+    let t = (t / step).floor() * step;
+
+    let asset = match thumb_track_source {
+        TrackSource::Local(LocalTrack { path, .. }) => {
+            // the track selected might be different from thumb_track
+            jellytranscoder::thumbnail::create_thumbnail(path, t).await?
+        }
+        TrackSource::Remote(_) => {
+            let session = fed
+                .get_session(
+                    thumb_track
+                        .federated
+                        .last()
+                        .ok_or(anyhow!("federation broken"))?,
+                )
+                .await?;
+
+            async_cache_file(&["fed-thumb", id, &format!("{t}")], |out| {
+                session.node_thumbnail(out, id, 2048, t)
+            })
+            .await?
+        }
+    };
+
+    Ok(asset_with_res(asset, width).await?)
+}
+
+async fn asset_with_res(
+    asset: AssetLocation,
+    width: Option<usize>,
+) -> MyResult<(ContentType, CacheControlFile)> {
     // fit the resolution into a finite set so the maximum cache is finite too.
-    let width = 2usize.pow(width.unwrap_or(2048).clamp(128, 8196).ilog2());
+    let width = 2usize.pow(width.unwrap_or(2048).clamp(128, 2048).ilog2());
     let path = jellytranscoder::image::transcode(asset, 50., 5, width)
         .await
         .context("transcoding asset")?;
diff --git a/transcoder/src/lib.rs b/transcoder/src/lib.rs
index d540e57..63d2cb8 100644
--- a/transcoder/src/lib.rs
+++ b/transcoder/src/lib.rs
@@ -10,6 +10,7 @@ use tokio::sync::Semaphore;
 pub mod image;
 pub mod snippet;
 pub mod subtitles;
+pub mod thumbnail;
 
 static LOCAL_IMAGE_TRANSCODING_TASKS: Semaphore = Semaphore::const_new(8);
 static LOCAL_VIDEO_TRANSCODING_TASKS: Semaphore = Semaphore::const_new(2);
diff --git a/transcoder/src/thumbnail.rs b/transcoder/src/thumbnail.rs
new file mode 100644
index 0000000..5baf888
--- /dev/null
+++ b/transcoder/src/thumbnail.rs
@@ -0,0 +1,35 @@
+use crate::LOCAL_IMAGE_TRANSCODING_TASKS;
+use jellybase::cache::async_cache_file;
+use jellycommon::AssetLocation;
+use log::info;
+use std::{path::Path, process::Stdio};
+use tokio::{io::copy, process::Command};
+
+pub async fn create_thumbnail(path: &Path, time: f64) -> anyhow::Result<AssetLocation> {
+    Ok(async_cache_file(
+        &["thumb", path.to_str().unwrap(), &format!("{time}")],
+        move |mut output| async move {
+            let _permit = LOCAL_IMAGE_TRANSCODING_TASKS.acquire().await?;
+            info!("creating thumbnail of {path:?} at {time}s",);
+
+            let mut proc = Command::new("ffmpeg")
+                .stdout(Stdio::piped())
+                .args(&["-ss", &format!("{time}")])
+                .args(&["-f", "matroska", "-i", path.to_str().unwrap()])
+                .args(&["-frames:v", "1"])
+                .args(&["-c:v", "qoi"])
+                .args(&["-f", "image2"])
+                .args(&["-update", "1"])
+                .arg("pipe:1")
+                .spawn()?;
+
+            let mut stdout = proc.stdout.take().unwrap();
+            copy(&mut stdout, &mut output).await?;
+
+            proc.wait().await.unwrap().exit_ok()?;
+            info!("done");
+            Ok(())
+        },
+    )
+    .await?)
+}