use std::collections::HashMap;
use std::collections::HashSet;
use csv_async::AsyncReaderBuilder;
use futures::future::ready;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use serde::de::Unexpected;
use serde::Deserialize;
use serde::Deserializer;
use crate::Error;
use crate::Episode;
use crate::Movie;
use crate::Show;
use crate::TITLES_URL;
use crate::get_episodes_filtered;
use crate::get_movies_filtered;
use crate::get_shows_filtered;
use crate::start_stream;
use crate::start_stream_lines;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Title {
#[serde(rename = "tconst", deserialize_with = "crate::util::parse_imdb_id")]
pub imdb_id: u64,
pub title_type: TitleType,
pub primary_title: String,
pub original_title: String,
pub is_adult: u8,
#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
pub start_year: Option<u16>,
#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
pub end_year: Option<u16>,
#[serde(deserialize_with = "crate::util::parse_janky_tsv_option")]
pub runtime_minutes: Option<u16>,
#[serde(with = "serde_with::StringWithSeparator::<serde_with::CommaSeparator>")]
pub genres: Vec<String>
}
#[derive(Debug, Eq, PartialEq)]
pub enum TitleType {
Short,
Movie,
Episode,
TVMiniSeries,
TVSeries,
TVShort,
TVSpecial,
Video,
VideoGame,
Audiobook,
RadioSeries
}
impl<'de> Deserialize<'de> for TitleType {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s: &str = Deserialize::deserialize(deserializer)?;
Ok(match s {
"short" => Self::Short,
"episode" | "tvEpisode" => Self::Episode,
"tvMiniSeries" => Self::TVMiniSeries,
"movie" | "tvMovie" => Self::Movie,
"tvSeries" => Self::TVSeries,
"tvShort" => Self::TVShort,
"tvSpecial" => Self::TVSpecial,
"video" => Self::Video,
"videoGame" => Self::VideoGame,
"audiobook" => Self::Audiobook,
"radioSeries" => Self::RadioSeries,
s => return Err(serde::de::Error::invalid_value(Unexpected::Str(&s), &"'short', 'movie', 'tvEpisode', 'tvMiniSeries', 'tvMovie', 'tvSeries', 'tvShort', 'tvSpecial', 'video', 'videoGame'"))
})
}
}
pub async fn get_titles_filtered_split(movie_ids: &[u64], show_ids: &[u64]) -> Result<(Vec<Movie>, Vec<Show>), Error> {
if(show_ids.len() == 0) {
return Ok((get_movies_filtered(movie_ids).await?, vec![]));
} else if(movie_ids.len() == 0) {
return Ok((vec![], get_shows_filtered(show_ids).await?));
}
let mut episode_links = get_episodes_filtered(show_ids).await?;
let mut all_ids = Vec::from(movie_ids);
all_ids.extend(show_ids);
all_ids.extend(episode_links.iter().map(|(id, _)| id));
let mut stream = stream_titles_filtered(&all_ids).await?;
let mut episodes_by_show = HashMap::new();
let mut movies = Vec::with_capacity(movie_ids.len());
let mut shows = Vec::with_capacity(show_ids.len());
while let Some(result) = stream.next().await {
let title = result?;
if let Some(link) = episode_links.remove(&title.imdb_id) {
let series_id = link.series_imdb_id;
if let Some(result) = Episode::from_title_and_link(title, link) {
let episode = result?;
let episodes = episodes_by_show.entry(series_id).or_insert_with(|| Vec::new());
episodes.push(episode);
}
} else if let Some(show) = Show::from_title(&title) {
shows.push(show?);
} else if let Some(movie) = Movie::from_title(&title) {
movies.push(movie?);
}
}
for show in shows.iter_mut() {
show.episodes = match episodes_by_show.remove(&show.imdb_id) {
Some(v) => v,
None => Vec::new()
};
}
Ok((movies, shows))
}
pub async fn get_titles_filtered(ids: &[u64]) -> Result<Vec<Title>, Error> {
let ids = Vec::from(ids);
let mut stream = stream_titles_filtered(&ids).await?;
let mut titles = Vec::with_capacity(ids.len());
while let Some(result) = stream.next().await {
titles.push(result?);
}
Ok(titles)
}
pub async fn stream_titles_filtered(ids: &[u64]) -> Result<impl Stream<Item = Result<Title, Error>> + '_, Error> {
let mut ids: HashSet<String> = ids.iter().map(|id| format!("tt{:07}", id)).collect();
let reader = start_stream_lines(TITLES_URL).await?;
let stream = AsyncReaderBuilder::new()
.delimiter(b'\t')
.has_headers(false)
.create_deserializer(reader
.filter(move |result| match result {
Ok(line) => {
if(ids.remove(line.split('\t').nth(0).unwrap())) {
ready(true)
} else {
ready(false)
}
}
Err(_) => ready(true)
})
.map_ok(|l| l + "\n")
.into_async_read()
)
.into_deserialize::<Title>();
Ok(stream
.err_into()
)
}
pub async fn stream_titles() -> Result<impl Stream<Item = Result<Title, Error>>, Error> {
let reader = start_stream(TITLES_URL).await?;
let stream = AsyncReaderBuilder::new()
.delimiter(b'\t')
.create_deserializer(reader)
.into_deserialize::<Title>();
Ok(stream
.err_into()
)
}