trakt import

3 files changed, 181 insertions, 398 deletions

diff --git a/import/src/lib.rs b/import/src/lib.rs
index ab74ecb..d6eb54f 100644
--- a/import/src/lib.rs
+++ b/import/src/lib.rs
@@ -40,9 +40,14 @@ use std::{
 };
 use tmdb::{parse_release_date, tmdb_image};
 use tokio::{io::AsyncWriteExt, sync::Semaphore, task::spawn_blocking};
+use trakt::Trakt;
 
 static IMPORT_SEM: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(1));
 
+struct Apis {
+    trakt: Option<Trakt>,
+}
+
 pub async fn import(db: &DataAcid, fed: &Federation) -> anyhow::Result<()> {
     let permit = IMPORT_SEM.try_acquire()?;
 
@@ -56,8 +61,13 @@ pub async fn import(db: &DataAcid, fed: &Federation) -> anyhow::Result<()> {
         drop(table);
         txn.commit()?;
     }
+
+    let ap = Apis {
+        trakt: SECRETS.api.trakt.as_ref().map(|key| Trakt::new(key)),
+    };
+
     info!("loading sources...");
-    import_path(CONF.library_path.clone(), vec![], db, fed)
+    import_path(CONF.library_path.clone(), vec![], db, fed, &ap)
         .await
         .context("indexing")?;
     info!("removing old nodes...");
@@ -155,11 +165,12 @@ fn compare_index_path(x: &[usize], y: &[usize]) -> Ordering {
 }
 
 #[async_recursion]
-pub async fn import_path(
+async fn import_path(
     path: PathBuf,
     index_path: Vec<usize>,
     db: &DataAcid,
     fed: &Federation,
+    ap: &Apis,
 ) -> anyhow::Result<()> {
     if path.is_dir() {
         let mut children_paths = path
@@ -192,6 +203,7 @@ pub async fn import_path(
                     },
                     db,
                     fed,
+                    ap,
                 )
             })
             .collect();
@@ -209,7 +221,7 @@ pub async fn import_path(
         };
 
         for s in opts.sources {
-            process_source(opts.id.clone(), s, &path, &index_path, db, fed).await?;
+            process_source(opts.id.clone(), s, &path, &index_path, db, fed, ap).await?;
         }
     }
     Ok(())
@@ -225,6 +237,7 @@ async fn process_source(
     index_path: &[usize],
     db: &DataAcid,
     fed: &Federation,
+    ap: &Apis,
 ) -> anyhow::Result<()> {
     let insert_node = move |id: &str, n: Node| -> anyhow::Result<()> {
         let txn = db.inner.begin_write()?;
@@ -240,6 +253,42 @@ async fn process_source(
     };
     match s {
         ImportSource::Override(n) => insert_node(&id, n)?,
+        ImportSource::Trakt { id: tid, kind } => {
+            let trakt_object = ap
+                .trakt
+                .as_ref()
+                .ok_or(anyhow!("trakt api key is required"))?
+                .lookup(kind, tid, true)
+                .await?;
+
+            let mut node = Node::default();
+            node.public.title = Some(trakt_object.title.to_owned());
+            if let Some(overview) = &trakt_object.overview {
+                node.public.description = Some(overview.to_owned())
+            }
+            if let Some(tagline) = &trakt_object.tagline {
+                node.public.tagline = Some(tagline.to_owned())
+            }
+            if let Some(rating) = &trakt_object.rating {
+                node.public.ratings.insert(Rating::Trakt, *rating);
+            }
+            insert_node(&id, node)?;
+
+            if let Some(tid) = trakt_object.ids.tmdb {
+                let mut index_path = index_path.to_vec();
+                index_path.push(1);
+                process_source(
+                    id,
+                    ImportSource::Tmdb { id: tid },
+                    path,
+                    &index_path,
+                    db,
+                    fed,
+                    ap,
+                )
+                .await?;
+            }
+        }
         ImportSource::Tmdb { id: tid } => {
             info!("tmdb lookup {id}");
             let key = SECRETS
@@ -324,6 +373,7 @@ async fn process_source(
                         index_path,
                         db,
                         fed,
+                        ap,
                     )
                     .await
                     .context(anyhow!("recursive media import: {:?}", f.path()))?;
@@ -624,16 +674,3 @@ async fn cache_federation_asset(
     )
     .await
 }
-
-// fn make_ident(s: &str) -> String {
-//     let mut out = String::new();
-//     for s in s.chars() {
-//         match s {
-//             'a'..='z' | '0'..='9' => out.push(s),
-//             'A'..='Z' => out.push(s.to_ascii_lowercase()),
-//             '-' | ' ' | '_' | ':' => out.push('-'),
-//             _ => (),
-//         }
-//     }
-//     out
-// }
diff --git a/import/src/mod.rs b/import/src/mod.rs
deleted file mode 100644
index 0c43cde..0000000
--- a/import/src/mod.rs
+++ /dev/null
@@ -1,319 +0,0 @@
-/*
-    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
-    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
-    Copyright (C) 2023 metamuffin <metamuffin.org>
-*/
-pub mod infojson;
-pub mod tmdb;
-
-use crate::{make_ident, ok_or_warn, Action};
-use anyhow::Context;
-use infojson::YVideo;
-use jellycommon::{
-    AssetLocation, LocalTrack, MediaInfo, Node, NodeKind, NodePrivate, NodePublic, Rating,
-    TrackSource,
-};
-use jellymatroska::read::EbmlReader;
-use jellyremuxer::import::import_metadata;
-use log::{debug, info, warn};
-use std::{
-    collections::BTreeMap,
-    fs::{remove_file, File},
-    io::{stdin, BufReader, Write},
-};
-use tmdb::{tmdb_details, tmdb_image};
-
-pub(crate) fn import(action: Action, dry: bool) -> anyhow::Result<()> {
-    match action {
-        Action::New {
-            path,
-            tmdb_id,
-            tmdb_search,
-            input,
-            series,
-            ident_prefix,
-            ignore_attachments,
-            copy,
-            video,
-            ignore_metadata,
-            r#move,
-            title,
-            skip_existing,
-        } => {
-            if std::env::current_dir().unwrap().file_name().unwrap() != "library" {
-                warn!("new command can only be used in the library directory; what you are doing right now probably wont work.")
-            }
-
-            if skip_existing {
-                if let Some(input) = &input {
-                    let guessed_path = path.join(input.file_stem().unwrap_or(input.as_os_str()));
-                    if guessed_path.exists() {
-                        info!("guessed output ({guessed_path:?}) exists, skipping import");
-                        return Ok(());
-                    } else {
-                        debug!("guessed output ({guessed_path:?}) missing");
-                    }
-                }
-            }
-
-            let tmdb_kind = if series { "tv" } else { "movie" };
-            let tmdb_id = if let Some(id) = tmdb_id {
-                Some(id.parse().unwrap())
-            } else if let Some(title) = tmdb_search {
-                let tmdb_key = std::env::var("TMDB_API_KEY").context("tmdb api key required")?;
-                let results = tmdb::tmdb_search(tmdb_kind, &title, &tmdb_key)?;
-                info!("results:");
-                for (i, r) in results.results.iter().enumerate() {
-                    info!(
-                        "\t[{i}] {}: {} ({})",
-                        r.id,
-                        r.title.as_ref().or(r.name.as_ref()).unwrap(),
-                        r.overview.chars().take(100).collect::<String>()
-                    );
-                }
-                let res_index = if results.results.len() > 1 {
-                    stdin()
-                        .lines()
-                        .next()
-                        .unwrap()
-                        .unwrap()
-                        .parse::<usize>()
-                        .unwrap()
-                } else {
-                    0
-                };
-                Some(results.results[res_index].id)
-            } else {
-                None
-            };
-
-            let tmdb_details = tmdb_id
-                .map(|id| {
-                    let tmdb_key =
-                        std::env::var("TMDB_API_KEY").context("tmdb api key required")?;
-                    let td = tmdb_details(tmdb_kind, id, &tmdb_key)
-                        .context("fetching details")
-                        .unwrap();
-                    Ok::<_, anyhow::Error>(td)
-                })
-                .transpose()?;
-
-            let mut kind = NodeKind::Series;
-            let mut file_meta = None;
-            let mut infojson = None;
-
-            if let Some(input_path) = &input {
-                file_meta = Some({
-                    let input = BufReader::new(File::open(&input_path).unwrap());
-                    let mut input = EbmlReader::new(input);
-                    import_metadata(&mut input)?
-                });
-                if ignore_attachments {
-                    let file_meta = file_meta.as_mut().unwrap();
-                    file_meta.cover = None;
-                    file_meta.infojson = None;
-                }
-                if ignore_metadata {
-                    let file_meta = file_meta.as_mut().unwrap();
-                    file_meta.description = None;
-                    file_meta.tagline = None;
-                    file_meta.title = None;
-                }
-
-                if let Some(ij) = &file_meta.as_ref().unwrap().infojson {
-                    infojson =
-                        Some(serde_json::from_str::<YVideo>(ij).context("parsing info.json")?);
-                }
-
-                kind = if video {
-                    NodeKind::Video
-                } else {
-                    NodeKind::Movie
-                };
-            }
-
-            let title = title
-                .or(tmdb_details
-                    .as_ref()
-                    .map(|d| d.title.clone().or(d.name.clone()))
-                    .flatten())
-                .or(file_meta.as_ref().map(|m| m.title.clone()).flatten())
-                .expect("no title detected");
-
-            let ident = format!(
-                "{}{}",
-                ident_prefix.unwrap_or(String::new()),
-                make_ident(
-                    &infojson
-                        .as_ref()
-                        .map(|i| i.id.clone())
-                        .unwrap_or(title.clone())
-                ),
-            );
-            let path = path.join(&ident);
-            let source_path = input.as_ref().map(|_| path.join(format!("source.mkv")));
-
-            let (mut poster, mut backdrop) = (None, None);
-            if !dry {
-                std::fs::create_dir_all(&path)?;
-
-                poster = file_meta
-                    .as_ref()
-                    .map(|m| {
-                        m.cover
-                            .as_ref()
-                            .map(|(mime, data)| {
-                                let pu = path.join(format!(
-                                    "cover.{}",
-                                    match mime.as_str() {
-                                        "image/webp" => "webp",
-                                        "image/jpeg" => "jpeg",
-                                        "image/png" => "png",
-                                        _ => {
-                                            warn!("unknown mime, just using webp");
-                                            "webp"
-                                        }
-                                    }
-                                ));
-                                if !pu.exists() {
-                                    let mut f = File::create(&pu)?;
-                                    f.write_all(&data)?;
-                                }
-                                Ok::<_, anyhow::Error>(pu)
-                            })
-                            .transpose()
-                    })
-                    .transpose()?
-                    .flatten()
-                    .or(tmdb_details
-                        .as_ref()
-                        .map(|d| {
-                            d.poster_path
-                                .as_ref()
-                                .map(|p| {
-                                    let pu = path.join("poster.jpeg");
-                                    let mut f = File::create(&pu)?;
-                                    tmdb_image(&p, &mut f)?;
-                                    Ok::<_, anyhow::Error>(pu)
-                                })
-                                .transpose()
-                        })
-                        .transpose()?
-                        .flatten());
-
-                backdrop = tmdb_details
-                    .as_ref()
-                    .map(|d| {
-                        d.backdrop_path
-                            .as_ref()
-                            .map(|p| {
-                                let pu = path.join("backdrop.jpeg");
-                                let mut f = File::create(&pu)?;
-                                tmdb_image(&p, &mut f)?;
-                                Ok::<_, anyhow::Error>(pu)
-                            })
-                            .transpose()
-                    })
-                    .transpose()?
-                    .flatten();
-            }
-
-            let mut ratings = BTreeMap::new();
-
-            ratings.extend(
-                infojson
-                    .as_ref()
-                    .map(|i| (Rating::YoutubeViews, i.view_count as f64)),
-            );
-            ratings.extend(
-                infojson
-                    .as_ref()
-                    .map(|i| i.like_count.map(|l| (Rating::YoutubeLikes, l as f64)))
-                    .flatten(),
-            );
-            ratings.extend(
-                tmdb_details
-                    .as_ref()
-                    .map(|d| (Rating::Tmdb, d.vote_average)),
-            );
-
-            let node = Node {
-                private: NodePrivate {
-                    id: Some(ident.clone()),
-                    backdrop: backdrop.clone().map(AssetLocation::Library),
-                    poster: poster.clone().map(AssetLocation::Library),
-                    source: file_meta.as_ref().map(|m| MediaSource::Local {
-                        tracks: m
-                            .track_sources
-                            .clone()
-                            .into_iter()
-                            .map(|t| LocalTrack {
-                                path: source_path.clone().unwrap(),
-                                ..t
-                            })
-                            .collect(),
-                    }),
-                },
-                public: NodePublic {
-                    federated: None,
-                    ratings,
-                    description: file_meta
-                        .as_ref()
-                        .map(|m| m.description.clone())
-                        .flatten()
-                        .or(tmdb_details.as_ref().map(|d| d.overview.to_owned())),
-                    tagline: file_meta.as_ref().map(|m| m.tagline.clone()).flatten().or(
-                        tmdb_details
-                            .as_ref()
-                            .map(|d| d.tagline.to_owned())
-                            .flatten(),
-                    ),
-                    title: Some(title),
-                    index: None,
-                    kind: Some(kind),
-                    children: Vec::new(),
-                    media: file_meta.as_ref().map(|m| MediaInfo {
-                        chapters: m.chapters.clone(),
-                        duration: m.duration,
-                        tracks: m.tracks.clone(),
-                    }),
-                    release_date: tmdb_details
-                        .as_ref()
-                        .map(|d| tmdb::parse_release_date(&d.release_date.clone()?).ok())
-                        .flatten()
-                        .or(infojson
-                            .as_ref()
-                            .and_then(|j| ok_or_warn(infojson::parse_upload_date(&j.upload_date)))),
-                    ..Default::default()
-                },
-            };
-
-            if dry {
-                println!("{}", serde_json::to_string_pretty(&node)?);
-            } else {
-                if let Some(source_path) = source_path {
-                    let input = input.clone().unwrap();
-                    if r#move {
-                        std::fs::rename(&input, &source_path)?;
-                    } else if copy {
-                        std::fs::copy(&input, &source_path)?;
-                    } else {
-                        if source_path.is_symlink() {
-                            remove_file(&source_path)?;
-                        }
-                        std::os::unix::fs::symlink(&input, &source_path)?;
-                    }
-                }
-                let f = File::create(path.join(if series {
-                    "directory.json"
-                } else {
-                    "item.jelly"
-                }))?;
-                serde_json::to_writer_pretty(f, &node)?;
-            }
-
-            Ok(())
-        }
-        _ => unreachable!(),
-    }
-}
diff --git a/import/src/trakt.rs b/import/src/trakt.rs
index e142eb6..9674351 100644
--- a/import/src/trakt.rs
+++ b/import/src/trakt.rs
@@ -1,8 +1,12 @@
+use bincode::{Decode, Encode};
+use jellybase::cache::async_cache_memory;
+use jellycommon::TraktKind;
 use reqwest::{
     header::{HeaderMap, HeaderName, HeaderValue},
     Client, ClientBuilder,
 };
 use serde::{Deserialize, Serialize};
+use std::{fmt::Display, sync::Arc};
 
 pub struct Trakt {
     client: Client,
@@ -32,25 +36,50 @@ impl Trakt {
 
     pub async fn search(
         &self,
-        types: &[TraktKind],
+        kinds: &[TraktKind],
         query: &str,
         extended: bool,
-    ) -> anyhow::Result<Vec<TraktSearchResult>> {
-        let res = self
-            .client
-            .get(format!(
-                "https://api.trakt.tv/search/{}?query={}{}",
-                types
-                    .iter()
-                    .map(|t| serde_json::to_string(t).unwrap())
-                    .collect::<Vec<_>>()
-                    .join(","),
-                urlencoding::encode(query),
-                optext(extended)
-            ))
-            .send()
-            .await?;
-        Ok(res.json().await?)
+    ) -> anyhow::Result<Arc<Vec<TraktSearchResult>>> {
+        async_cache_memory(
+            &["api-trakt-lookup", query, if extended { "a" } else { "b" }],
+            || async move {
+                let url = format!(
+                    "https://api.trakt.tv/search/{}?query={}{}",
+                    kinds
+                        .iter()
+                        .map(|t| t.singular())
+                        .collect::<Vec<_>>()
+                        .join(","),
+                    urlencoding::encode(query),
+                    optext(extended)
+                );
+                let res = self.client.get(url).send().await?.error_for_status()?;
+                Ok(res.json().await?)
+            },
+        )
+        .await
+    }
+
+    pub async fn lookup(
+        &self,
+        kind: TraktKind,
+        id: u64,
+        extended: bool,
+    ) -> anyhow::Result<Arc<TraktMediaObject>> {
+        async_cache_memory(
+            &["api-trakt-lookup", &format!("{id} {extended}")],
+            || async move {
+                let url = format!(
+                    "https://api.trakt.tv/{}/{}{}",
+                    kind.plural(),
+                    id,
+                    optext2(extended)
+                );
+                let res = self.client.get(url).send().await?.error_for_status()?;
+                Ok(res.json().await?)
+            },
+        )
+        .await
     }
 }
 
@@ -61,49 +90,23 @@ fn optext(extended: bool) -> &'static str {
         ""
     }
 }
-
-#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
-#[serde(rename_all = "snake_case")]
-pub enum TraktKind {
-    Movie,
-    Show,
-    Season,
-    Episode,
-    Person,
-    User,
-}
-
-impl TraktKind {
-    pub fn singular(self) -> &'static str {
-        match self {
-            TraktKind::Movie => "movie",
-            TraktKind::Show => "show",
-            TraktKind::Season => "season",
-            TraktKind::Episode => "episode",
-            TraktKind::Person => "person",
-            TraktKind::User => "user",
-        }
-    }
-    pub fn plural(self) -> &'static str {
-        match self {
-            TraktKind::Movie => "movies",
-            TraktKind::Show => "shows",
-            TraktKind::Season => "seasons",
-            TraktKind::Episode => "episodes",
-            TraktKind::Person => "people",
-            TraktKind::User => "user", // //! not used in API
-        }
+fn optext2(extended: bool) -> &'static str {
+    if extended {
+        "?extended=full"
+    } else {
+        ""
     }
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Encode, Decode)]
 pub struct TraktSearchResult {
-    r#type: TraktKind,
-    score: f64,
-    inner: TraktKindObject,
+    pub r#type: TraktKind,
+    pub score: f64,
+    #[serde(flatten)]
+    pub inner: TraktKindObject,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Encode, Decode)]
 #[serde(rename_all = "snake_case")]
 pub enum TraktKindObject {
     Movie(TraktMediaObject),
@@ -114,18 +117,80 @@ pub enum TraktKindObject {
     User(TraktMediaObject),
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+impl TraktKindObject {
+    pub fn inner(&self) -> &TraktMediaObject {
+        match self {
+            TraktKindObject::Movie(x)
+            | TraktKindObject::Show(x)
+            | TraktKindObject::Season(x)
+            | TraktKindObject::Episode(x)
+            | TraktKindObject::Person(x)
+            | TraktKindObject::User(x) => x,
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize, Encode, Decode, Clone)]
 pub struct TraktMediaObject {
-    title: String,
-    year: Option<u32>,
-    ids: TraktMediaObjectIds,
+    pub title: String,
+    pub year: Option<u32>,
+    pub ids: TraktMediaObjectIds,
+
+    pub tagline: Option<String>,
+    pub overview: Option<String>,
+    pub released: Option<String>,
+    pub runtime: Option<usize>,
+    pub country: Option<String>,
+    pub trailer: Option<String>,
+    pub homepage: Option<String>,
+    pub status: Option<String>,
+    pub rating: Option<f64>,
+    pub votes: Option<usize>,
+    pub comment_count: Option<usize>,
+    pub language: Option<String>,
+    pub available_translations: Option<Vec<String>>,
+    pub genres: Option<Vec<String>>,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Encode, Decode, Clone)]
 pub struct TraktMediaObjectIds {
-    trakt: u64,
-    slug: String,
+    pub trakt: u64,
+    pub slug: Option<String>,
+    pub imdb: Option<String>,
+    pub tmdb: Option<u64>,
+    pub omdb: Option<u64>,
+    pub tvdb: Option<u64>,
+}
 
-    imdb: Option<String>,
-    tmdb: Option<u64>,
+impl Display for TraktSearchResult {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!(
+            "{}: {} ({}) \x1b[2m[{}]\x1b[0m",
+            self.r#type.to_string(),
+            self.inner.inner().title,
+            self.inner.inner().year.unwrap_or(0),
+            self.inner.inner().ids
+        ))
+    }
+}
+impl Display for TraktMediaObjectIds {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("trakt")?;
+        if self.slug.is_some() {
+            f.write_str(",slug")?;
+        }
+        if self.tmdb.is_some() {
+            f.write_str(",tmdb")?;
+        }
+        if self.imdb.is_some() {
+            f.write_str(",imdb")?;
+        }
+        if self.tvdb.is_some() {
+            f.write_str(",tvdb")?;
+        }
+        if self.omdb.is_some() {
+            f.write_str(",omdb")?;
+        }
+        Ok(())
+    }
 }