1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
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,
	//_title_type_string: String, // tsv seems to have a bug where it doesn't remove the TitleType value after we deserialize it.  It works when we use JSON.
	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'"))
		})
	}
}

// TODO:  Rejigger this whole shebang so that we don't need to duplicate most
//    of the code from both get_movies_filtered() and get_shows_filtered(); in
//    all likelihood, this means we pass a closure down to the individual
//    stream_*_filtered() functions.
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) => {
					// This is split out into multiple lines currently to make it easier to add everybody's favorite - printf debugging!
					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
		//.map_err(|e| Error::from(e))
		/*
		.filter(move |result| match result {
			//Ok(line) => ready(ids.contains(line.split('\t').nth(0).unwrap())),
			Ok(title) => ready(ids.contains(&title.imdb_id)),
			Err(_) => ready(true)
		})
		*/
		//.map(|r| r.flatten())
		.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
		//.map_err(|e| Error::from(e))
		//.map_ok(move |s| tsv::de::from_str(&s).map_err(|e| Error::from(e)))
		//.map(|r| r.flatten())
		.err_into()
	)
}