trakt import

1 files changed, 128 insertions, 63 deletions

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(())
+    }
 }