path: root/import/src/plugins/tmdb.rs

/*
    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) 2026 metamuffin <metamuffin.org>
*/
use crate::{
    USER_AGENT,
    plugins::{ImportPlugin, PluginContext, PluginInfo},
    source_rank::ObjectImportSourceExt,
};
use anyhow::{Context, Result, anyhow, bail};
use chrono::{Utc, format::Parsed};
use jellycache::{Cache, EscapeKey, HashKey};
use jellycommon::*;
use jellydb::RowNum;
use log::info;
use reqwest::{
    Client, ClientBuilder,
    header::{HeaderMap, HeaderName, HeaderValue},
};
use serde::{Deserialize, Serialize};
use std::{fmt::Display, sync::Arc};
use tokio::runtime::Handle;

pub struct Tmdb {
    client: Client,
    image_client: Client,
    key: String,
}

impl Tmdb {
    pub fn new(api_key: &str) -> Self {
        let client = ClientBuilder::new()
            .default_headers(HeaderMap::from_iter([
                (
                    HeaderName::from_static("accept"),
                    HeaderValue::from_static("application/json"),
                ),
                (
                    HeaderName::from_static("user-agent"),
                    HeaderValue::from_static(USER_AGENT),
                ),
            ]))
            .build()
            .unwrap();
        let image_client = ClientBuilder::new().build().unwrap();
        Self {
            client,
            image_client,
            key: api_key.to_owned(),
        }
    }
    pub fn search(
        &self,
        cache: &Cache,
        kind: TmdbKind,
        query: &str,
        rt: &Handle,
    ) -> Result<Arc<TmdbQuery>> {
        cache
            .cache_memory(
                &format!("ext/tmdb/search/{kind}-{}.json", HashKey(query)),
                move || {
                    rt.block_on(async {
                        info!("searching tmdb: {query:?}");
                        Ok(self
                            .client
                            .get(format!(
                                "https://api.themoviedb.org/3/search/{kind}?query={}?api_key={}",
                                query.replace(" ", "+"),
                                self.key
                            ))
                            .send()
                            .await?
                            .error_for_status()?
                            .json::<TmdbQuery>()
                            .await?)
                    })
                },
            )
            .context("tmdb search")
    }
    pub fn details(
        &self,
        cache: &Cache,
        kind: TmdbKind,
        id: u64,
        rt: &Handle,
    ) -> Result<Arc<TmdbDetails>> {
        cache
            .cache_memory(&format!("ext/tmdb/details/{kind}-{id}.json"), move || {
                rt.block_on(async {
                    info!("fetching details: {id:?}");
                    Ok(self
                        .client
                        .get(format!(
                            "https://api.themoviedb.org/3/{kind}/{id}?api_key={}",
                            self.key,
                        ))
                        .send()
                        .await?
                        .error_for_status()?
                        .json()
                        .await?)
                })
            })
            .context("tmdb details")
    }
    pub fn person_image(
        &self,
        cache: &Cache,
        id: u64,
        rt: &Handle,
    ) -> Result<Arc<TmdbPersonImage>> {
        cache
            .cache_memory(&format!("ext/tmdb/person/images/{id}.json"), move || {
                rt.block_on(async {
                    Ok(self
                        .client
                        .get(format!(
                            "https://api.themoviedb.org/3/person/{id}/images?api_key={}",
                            self.key,
                        ))
                        .send()
                        .await?
                        .error_for_status()?
                        .json()
                        .await?)
                })
            })
            .context("tmdb person images")
    }
    pub fn image(&self, cache: &Cache, path: &str, rt: &Handle) -> Result<String> {
        cache
            .store(
                format!("ext/tmdb/image/{}.image", EscapeKey(path)),
                move || {
                    rt.block_on(async {
                        info!("downloading image {path:?}");
                        Ok(self
                            .image_client
                            .get(format!("https://image.tmdb.org/t/p/original{path}"))
                            .send()
                            .await?
                            .error_for_status()?
                            .bytes()
                            .await?
                            .to_vec())
                    })
                },
            )
            .context("tmdb image download")
    }

    pub fn episode_details(
        &self,
        cache: &Cache,
        series_id: u64,
        season: u64,
        episode: u64,
        rt: &Handle,
    ) -> Result<Arc<TmdbEpisode>> {
        cache.cache_memory(&format!("ext/tmdb/episode-details/{series_id}-S{season}-E{episode}.json"), move || {
            rt.block_on(async {
                info!("tmdb episode details {series_id} S={season} E={episode}");
                Ok(self
                    .image_client
                    .get(format!("https://api.themoviedb.org/3/tv/{series_id}/season/{season}/episode/{episode}?api_key={}", self.key))
                    .send()
                    .await?
                    .error_for_status()?
                    .json()
                    .await?)
            })
        })
        .context("tmdb episode details")
    }
}

impl ImportPlugin for Tmdb {
    fn info(&self) -> PluginInfo {
        PluginInfo {
            name: "tmdb",
            tag: MSOURCE_TMDB,
            handle_process: true,
            ..Default::default()
        }
    }
    fn process(&self, ct: &PluginContext, node: RowNum) -> Result<()> {
        self.process_primary(ct, node)?;
        self.process_episode(ct, node)?;
        self.process_person(ct, node)?;
        Ok(())
    }
}
impl Tmdb {
    fn process_primary(&self, ct: &PluginContext, node: RowNum) -> Result<()> {
        let data = ct.ic.get_node(node)?.unwrap();
        let data = data.as_object();

        let (tmdb_kind, tmdb_id): (_, u64) = if let Some(id) = data
            .get(NO_IDENTIFIERS)
            .unwrap_or_default()
            .get(IDENT_TMDB_SERIES)
        {
            (TmdbKind::Tv, id.parse()?)
        } else if let Some(id) = data
            .get(NO_IDENTIFIERS)
            .unwrap_or_default()
            .get(IDENT_TMDB_MOVIE)
        {
            (TmdbKind::Movie, id.parse()?)
        } else {
            return Ok(());
        };

        let details = self.details(&ct.ic.cache, tmdb_kind, tmdb_id, ct.rt)?;
        let backdrop = details
            .backdrop_path
            .as_ref()
            .map(|path| self.image(&ct.ic.cache, &path, ct.rt))
            .transpose()
            .context("backdrop image")?;
        let poster = details
            .poster_path
            .as_ref()
            .map(|path| self.image(&ct.ic.cache, &path, ct.rt))
            .transpose()
            .context("poster image")?;

        let release_date = details
            .release_date
            .as_ref()
            .map(|s| parse_release_date(s))
            .transpose()?
            .flatten();

        ct.ic.update_node(node, |mut node| {
            if let Some(title) = &details.title {
                node = node.as_object().insert_s(ct.is, NO_TITLE, &title);
            }
            if let Some(tagline) = &details.tagline {
                node = node.as_object().insert_s(ct.is, NO_TAGLINE, &tagline);
            }
            node = node
                .as_object()
                .insert_s(ct.is, NO_DESCRIPTION, &details.overview);
            node = node.as_object().update(NO_RATINGS, |rat| {
                rat.insert_s(ct.is, RTYP_TMDB, details.vote_average)
            });
            if let Some(poster) = &poster {
                node = node
                    .as_object()
                    .update(NO_PICTURES, |rat| rat.insert_s(ct.is, PICT_COVER, &poster));
            }
            if let Some(backdrop) = &backdrop {
                node = node.as_object().update(NO_PICTURES, |rat| {
                    rat.insert_s(ct.is, PICT_BACKDROP, &backdrop)
                });
            }
            if let Some(releasedate) = release_date {
                node = node
                    .as_object()
                    .insert_s(ct.is, NO_RELEASEDATE, releasedate);
            }
            node
        })?;
        Ok(())
    }
    fn process_episode(&self, ct: &PluginContext, node: RowNum) -> Result<()> {
        let data = ct.ic.get_node(node)?.unwrap();
        let data = data.as_object();

        let (Some(episode), Some(season)) = (data.get(NO_INDEX), data.get(NO_SEASON_INDEX)) else {
            return Ok(());
        };
        let mut series_id = None;
        ct.ic.db.transaction(&mut |txn| {
            for parent in data.iter(NO_PARENT) {
                let parent_data = txn.get(parent)?.ok_or(anyhow!("parent missing"))?;
                if let Some(id) = parent_data
                    .as_object()
                    .get(NO_IDENTIFIERS)
                    .unwrap_or_default()
                    .get(IDENT_TMDB_SERIES)
                {
                    series_id = Some(id.parse::<u64>()?);
                    break;
                }
            }
            Ok(())
        })?;
        let Some(series_id) = series_id else {
            return Ok(());
        };
        let details = self.episode_details(&ct.ic.cache, series_id, season, episode, ct.rt)?;
        let cover = details
            .still_path
            .as_ref()
            .map(|path| self.image(&ct.ic.cache, &path, ct.rt))
            .transpose()
            .context("still image download")?;
        let release_date = parse_release_date(&details.air_date)?;
        ct.ic.update_node(node, |mut node| {
            node = node.as_object().insert_s(ct.is, NO_TITLE, &details.name);
            node = node
                .as_object()
                .insert_s(ct.is, NO_DESCRIPTION, &details.overview);
            if let Some(release_date) = release_date {
                node = node
                    .as_object()
                    .insert_s(ct.is, NO_RELEASEDATE, release_date)
            }
            node = node.as_object().update(NO_RATINGS, |rat| {
                rat.insert_s(ct.is, RTYP_TMDB, details.vote_average)
            });
            if let Some(cover) = &cover {
                node = node.as_object().update(NO_PICTURES, |picts| {
                    picts.insert_s(ct.is, PICT_COVER, &cover)
                });
            }
            node
        })
    }

    fn process_person(&self, ct: &PluginContext, node: RowNum) -> Result<()> {
        let data = ct.ic.get_node(node)?.unwrap();
        let data = data.as_object();

        let Some(id) = data
            .get(NO_IDENTIFIERS)
            .unwrap_or_default()
            .get(IDENT_TMDB_PERSON)
        else {
            return Ok(());
        };
        let id = id.parse()?;

        let images = self.person_image(&ct.ic.cache, id, ct.rt)?;
        let Some(prof) = images.profiles.first() else {
            return Ok(());
        };

        let image = self.image(&ct.ic.cache, &prof.file_path, ct.rt)?;

        ct.ic.update_node(node, |node| {
            node.as_object()
                .update(NO_PICTURES, |pict| pict.insert_s(ct.is, PICT_COVER, &image))
        })?;

        Ok(())
    }
}

pub fn parse_release_date(d: &str) -> Result<Option<i64>> {
    if d.is_empty() {
        return Ok(None);
    } else if d.len() < 10 {
        bail!(anyhow!("date string too short"))
    }
    let (year, month, day) = (&d[0..4], &d[5..7], &d[8..10]);
    let (year, month, day) = (
        year.parse().context("parsing year")?,
        month.parse().context("parsing month")?,
        day.parse().context("parsing day")?,
    );

    let mut p = Parsed::new();
    p.year = Some(year);
    p.month = Some(month);
    p.day = Some(day);
    p.hour_div_12 = Some(0);
    p.hour_mod_12 = Some(0);
    p.minute = Some(0);
    p.second = Some(0);
    Ok(Some(p.to_datetime_with_timezone(&Utc)?.timestamp_millis()))
}

impl Display for TmdbKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            TmdbKind::Tv => "tv",
            TmdbKind::Movie => "movie",
        })
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbEpisode {
    pub air_date: String,
    pub overview: String,
    pub name: String,
    pub id: u64,
    pub runtime: f64,
    pub still_path: Option<String>,
    pub vote_average: f64,
    pub vote_count: usize,
}

#[derive(Debug, Clone, Copy, Hash, Serialize, Deserialize)]
pub enum TmdbKind {
    Tv,
    Movie,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbPersonImage {
    pub id: u64,
    pub profiles: Vec<TmdbPersonImageProfile>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbPersonImageProfile {
    pub aspect_ratio: f64,
    pub height: u32,
    pub width: u32,
    pub file_path: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbQuery {
    pub page: usize,
    pub results: Vec<TmdbQueryResult>,
    pub total_pages: usize,
    pub total_results: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbQueryResult {
    pub adult: bool,
    pub backdrop_path: Option<String>,
    pub genre_ids: Vec<u64>,
    pub id: u64,
    pub original_language: Option<String>,
    pub original_title: Option<String>,
    pub overview: String,
    pub popularity: f64,
    pub poster_path: Option<String>,
    pub release_date: Option<String>,
    pub title: Option<String>,
    pub name: Option<String>,
    pub vote_average: f64,
    pub vote_count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbDetails {
    pub adult: bool,
    pub backdrop_path: Option<String>,
    pub genres: Vec<TmdbGenre>,
    pub id: u64,
    pub original_language: Option<String>,
    pub original_title: Option<String>,
    pub overview: String,
    pub popularity: f64,
    pub poster_path: Option<String>,
    pub release_date: Option<String>,
    pub title: Option<String>,
    pub name: Option<String>,
    pub vote_average: f64,
    pub vote_count: usize,
    pub budget: Option<usize>,
    pub homepage: Option<String>,
    pub imdb_id: Option<String>,
    pub production_companies: Vec<TmdbProductionCompany>,
    pub revenue: Option<usize>,
    pub tagline: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbGenre {
    pub id: u64,
    pub name: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TmdbProductionCompany {
    pub id: u64,
    pub name: String,
    pub logo_path: Option<String>,
}