what2watch/src/imdb_utils.rs

use julid::Julid;
use sqlx::{Sqlite, SqlitePool};

use crate::{
    import_utils::{insert_credit, insert_star, insert_watch},
    Credit, ShowKind, Star, Watch,
};

pub type IdMap = std::collections::BTreeMap<String, Julid>;

const OMEGA_ID: Julid = Julid::omega();

#[derive(Debug, sqlx::FromRow, Clone)]
pub struct ImportImdbMovie {
    pub title: String,
    pub year: Option<String>,
    pub length: Option<String>,
    pub id: String,
    pub kind: Option<String>,
}

impl From<ImportImdbMovie> for Watch {
    fn from(value: ImportImdbMovie) -> Self {
        Watch {
            id: OMEGA_ID, // this is ignored by the inserter
            title: value.title,
            release_date: value.year,
            length: value.length.and_then(|v| v.parse::<i64>().ok()),
            kind: ShowKind::Movie,
            metadata_url: Some(format!("https://imdb.com/title/{}/", &value.id)),
            added_by: OMEGA_ID,
        }
    }
}

impl From<&ImportImdbMovie> for Watch {
    fn from(value: &ImportImdbMovie) -> Self {
        Watch {
            id: OMEGA_ID,
            title: value.title.to_string(),
            release_date: value.year.clone(),
            length: value.length.as_ref().and_then(|v| v.parse::<i64>().ok()),
            kind: ShowKind::Movie,
            metadata_url: Some(format!("https://imdb.com/title/{}/", value.id)),
            added_by: OMEGA_ID,
        }
    }
}

#[derive(Debug, sqlx::FromRow, Clone)]
pub struct ImdbStar {
    pub nconst: String,
    #[sqlx(rename = "primaryName")]
    pub name: String,
    #[sqlx(rename = "birthYear")]
    pub born: Option<String>,
    #[sqlx(rename = "deathYear")]
    pub died: Option<String>,
}

impl From<&ImdbStar> for Star {
    fn from(value: &ImdbStar) -> Self {
        let id = &value.nconst;
        let metadata_url = Some(format!("https://imdb.com/name/{id}/"));
        Self {
            name: value.name.clone(),
            metadata_url,
            born: value.born.clone(),
            died: value.died.clone(),
            ..Default::default()
        }
    }
}

pub async fn import_imdb_data(w2w_db: &SqlitePool, imdb: &SqlitePool, ids: &mut IdMap, num: u32) {
    const IMDB_QUERY: &str = "select * from watches order by year, title asc limit ?";

    let iwatches: Vec<ImportImdbMovie> = sqlx::query_as(IMDB_QUERY)
        .bind(num)
        .fetch_all(imdb)
        .await
        .unwrap();

    for batch in iwatches.chunks(5_000) {
        let mut tx = w2w_db.begin().await.unwrap();
        for iwatch in batch {
            let aid = iwatch.id.clone();
            let kind = show_kind(iwatch.kind.as_ref().unwrap());
            let mut watch: Watch = iwatch.into();
            watch.kind = kind;
            let watch_id: Julid = insert_watch(watch, &mut tx).await;
            add_imdb_stars(&mut tx, imdb, &aid, watch_id, ids).await;
            ids.insert(aid, watch_id);
        }
        tx.commit().await.unwrap();
    }
}

async fn add_imdb_stars(
    w2w_db: &mut sqlx::Transaction<'_, Sqlite>,
    imdb: &SqlitePool,
    iwatch: &str,
    watch: Julid,
    ids: &mut IdMap,
) {
    let principals_query = "select nconst, category from principals where tconst = ?";
    let principals = sqlx::query_as::<Sqlite, (String, String)>(principals_query)
        .bind(iwatch)
        .fetch_all(imdb)
        .await
        .unwrap();

    for row in principals {
        let (name_id, cat) = row;
        let star = if let Some(id) = ids.get(&name_id) {
            *id
        } else {
            let name_query =
                "select nconst, primaryName, birthYear, deathYear from names where nconst = ?";
            let istar: Option<ImdbStar> = sqlx::query_as(name_query)
                .bind(&name_id)
                .fetch_optional(imdb)
                .await
                .unwrap();
            if let Some(star) = istar {
                let star = (&star).into();
                let star_id = insert_star(&star, w2w_db).await;
                ids.insert(name_id, star_id);
                star_id
            } else {
                continue;
            }
        };

        let credit = Credit {
            star,
            watch,
            credit: Some(cat.to_string()),
        };
        insert_credit(&credit, w2w_db).await;
    }
}

fn show_kind(kind: &str) -> ShowKind {
    /*
    tvSeries
    tvMiniSeries
    tvMovie
    tvShort
    tvSpecial
    */
    match &kind[0..4] {
        "tvSe" => ShowKind::Series,
        "tvSh" => ShowKind::Short,
        "tvMi" => ShowKind::LimitedSeries,
        "tvSp" => ShowKind::Other,
        "tvMo" | "movi" => ShowKind::Movie,
        _ => ShowKind::Unknown,
    }
}