1use std::env;
2use std::path::PathBuf;
3use std::sync::Arc;
4use std::time::Duration;
5
6use clap::Args;
7use color_eyre::eyre::{self, Result};
8use fluent_templates::Loader;
9use novel_api::{CiweimaoClient, CiyuanjiClient, Client, ContentInfo, SfacgClient, VolumeInfos};
10use scc::HashMap;
11use tokio::sync::Semaphore;
12use url::Url;
13
14use crate::cmd::{Convert, Format, Source};
15use crate::utils::{self, Chapter, Content, Novel, ProgressBar, Volume};
16use crate::{LANG_ID, LOCALES, renderer};
17
18#[must_use]
19#[derive(Args)]
20#[command(arg_required_else_help = true,
21 about = LOCALES.lookup(&LANG_ID, "download_command"))]
22pub struct Download {
23 #[arg(help = LOCALES.lookup(&LANG_ID, "novel_id"))]
24 pub novel_id: u32,
25
26 #[arg(short, long,
27 help = LOCALES.lookup(&LANG_ID, "source"))]
28 pub source: Source,
29
30 #[arg(short, long, value_enum,
31 help = LOCALES.lookup(&LANG_ID, "format"))]
32 pub format: Format,
33
34 #[arg(short, long, value_enum, value_delimiter = ',',
35 help = LOCALES.lookup(&LANG_ID, "converts"))]
36 pub converts: Vec<Convert>,
37
38 #[arg(long, default_value_t = false,
39 help = LOCALES.lookup(&LANG_ID, "ignore_images"))]
40 pub ignore_images: bool,
41
42 #[arg(long, default_value_t = false,
43 help = LOCALES.lookup(&LANG_ID, "force_update_novel_db"))]
44 pub force_update_novel_db: bool,
45
46 #[arg(long, default_value_t = false,
47 help = LOCALES.lookup(&LANG_ID, "order_novel"))]
48 pub order_novel: bool,
49
50 #[arg(long, default_value_t = false,
51 help = LOCALES.lookup(&LANG_ID, "ignore_keyring"))]
52 pub ignore_keyring: bool,
53
54 #[arg(short, long, default_value_t = 4, value_parser = clap::value_parser!(u8).range(1..=8),
55 help = LOCALES.lookup(&LANG_ID, "maximum_concurrency"))]
56 pub maximum_concurrency: u8,
57
58 #[arg(long, help = LOCALES.lookup(&LANG_ID, "sleep"))]
59 pub sleep: Option<u64>,
60
61 #[arg(long, num_args = 0..=1, default_missing_value = super::DEFAULT_PROXY,
62 help = LOCALES.lookup(&LANG_ID, "proxy"))]
63 pub proxy: Option<Url>,
64
65 #[arg(long, default_value_t = false,
66 help = LOCALES.lookup(&LANG_ID, "no_proxy"))]
67 pub no_proxy: bool,
68
69 #[arg(long, num_args = 0..=1, default_missing_value = super::default_cert_path(),
70 help = super::cert_help_msg())]
71 pub cert: Option<PathBuf>,
72
73 #[arg(long, default_value_t = false,
74 help = LOCALES.lookup(&LANG_ID, "skip_login"))]
75 pub skip_login: bool,
76}
77
78pub async fn execute(mut config: Download) -> Result<()> {
79 check_skip_login_flag(&config)?;
80
81 match config.source {
82 Source::Sfacg => {
83 let mut client = SfacgClient::new().await?;
84 super::set_options(&mut client, &config.proxy, &config.no_proxy, &config.cert);
85 do_execute(client, config).await?;
86 }
87 Source::Ciweimao => {
88 let mut client = CiweimaoClient::new().await?;
89 super::set_options(&mut client, &config.proxy, &config.no_proxy, &config.cert);
90 do_execute(client, config).await?;
91 }
92 Source::Ciyuanji => {
93 if config.maximum_concurrency > 1 {
94 tracing::warn!(
95 "ciyuanji does not support concurrent downloads, set `maximum_concurrency` to 1"
96 );
97 config.maximum_concurrency = 1;
98 }
99
100 if config.sleep.is_none() {
101 tracing::warn!(
102 "ciyuanji has a limit on the number of downloads per minute, set `sleep` to 1"
103 );
104 config.sleep = Some(1)
105 }
106
107 let mut client = CiyuanjiClient::new().await?;
108 super::set_options(&mut client, &config.proxy, &config.no_proxy, &config.cert);
109 do_execute(client, config).await?;
110 }
111 }
112
113 Ok(())
114}
115
116fn check_skip_login_flag(config: &Download) -> Result<()> {
117 if config.skip_login && (config.source == Source::Ciweimao || config.source == Source::Ciyuanji)
118 {
119 eyre::bail!(
120 "This source cannot skip login: `{}`",
121 config.source.as_ref()
122 );
123 }
124
125 Ok(())
126}
127
128async fn do_execute<T>(client: T, config: Download) -> Result<()>
129where
130 T: Client + Send + Sync + 'static,
131{
132 if !config.skip_login {
133 if config.source == Source::Ciyuanji {
134 utils::log_in_without_password(&client).await?;
135 } else {
136 utils::log_in(&client, &config.source, config.ignore_keyring).await?;
137 }
138
139 let user_info = client.user_info().await?;
140 println!(
141 "{}",
142 utils::locales_with_arg("login_msg", "✨", user_info.nickname)
143 );
144 }
145
146 let mut novel = download_novel(client, &config).await?;
147 println!("{}", utils::locales("download_complete_msg", "👌"));
148
149 utils::convert(&mut novel, &config.converts)?;
150
151 match config.format {
152 Format::Pandoc => renderer::generate_pandoc_markdown(novel, &config.converts)?,
153 Format::Mdbook => renderer::generate_mdbook(novel, &config.converts).await?,
154 };
155
156 Ok(())
157}
158
159async fn download_novel<T>(client: T, config: &Download) -> Result<Novel>
160where
161 T: Client + Send + Sync + 'static,
162{
163 let client = Arc::new(client);
164 super::handle_shutdown_signal(&client);
165
166 if config.force_update_novel_db {
167 unsafe {
168 env::set_var("FORCE_UPDATE_NOVEL_DB", "true");
169 }
170 }
171
172 let novel_info = utils::novel_info(&client, config.novel_id).await?;
173
174 let mut novel = Novel {
175 name: novel_info.name,
176 author_name: novel_info.author_name,
177 introduction: novel_info.introduction,
178 cover_image: None,
179 volumes: Vec::new(),
180 };
181
182 println!(
183 "{}",
184 utils::locales_with_arg("start_msg", "🚚", &novel.name)
185 );
186
187 if novel_info.cover_url.is_some() && !config.ignore_images {
188 match client.image(&novel_info.cover_url.unwrap()).await {
189 Ok(image) => novel.cover_image = Some(image),
190 Err(error) => {
191 tracing::error!("Cover image download failed: `{error}`");
192 }
193 };
194 }
195
196 let Some(mut volume_infos) = client.volume_infos(config.novel_id).await? else {
197 eyre::bail!("Unable to get chapter information");
198 };
199
200 if config.order_novel {
201 client.order_novel(config.novel_id, &volume_infos).await?;
202 volume_infos = client.volume_infos(config.novel_id).await?.unwrap();
203
204 println!("{}", utils::locales("order_msg", "💰"));
205 }
206
207 let mut handles = Vec::with_capacity(128);
208 let pb = ProgressBar::new(chapter_count(&volume_infos))?;
209 let semaphore = Arc::new(Semaphore::new(config.maximum_concurrency as usize));
210 let chapter_map = Arc::new(HashMap::with_capacity(128));
211
212 let ignore_image = config.ignore_images;
213 let sleep = config.sleep;
214
215 for volume_info in volume_infos {
216 novel.volumes.push(Volume {
217 title: volume_info.title,
218 chapters: Vec::with_capacity(32),
219 });
220
221 let volume = novel.volumes.last_mut().unwrap();
222
223 for chapter_info in volume_info
224 .chapter_infos
225 .iter()
226 .filter(|x| !x.can_download())
227 {
228 tracing::info!(
229 "`{}-{}` can not be downloaded",
230 volume.title,
231 chapter_info.title
232 );
233 }
234
235 let can_download_chapter_infos = volume_info
236 .chapter_infos
237 .into_iter()
238 .filter(|x| x.can_download())
239 .collect::<Vec<_>>();
240
241 let chunk_size = match config.source {
242 Source::Sfacg => 1,
243 Source::Ciweimao => 10,
244 Source::Ciyuanji => 1,
245 };
246
247 for chapter_infos in can_download_chapter_infos
248 .chunks(chunk_size)
249 .map(|v| v.to_vec())
250 {
251 for chapter_info in &chapter_infos {
252 volume.chapters.push(Chapter {
253 id: chapter_info.id,
254 title: chapter_info.title.clone(),
255 contents: None,
256 });
257 }
258
259 let client = Arc::clone(&client);
260 let permit = semaphore.clone().acquire_owned().await.unwrap();
261 let mut pb = pb.clone();
262 let chapter_map = Arc::clone(&chapter_map);
263
264 handles.push(tokio::spawn(async move {
265 let msg = if chunk_size > 1 {
266 format!("等 {} 章", chunk_size)
267 } else {
268 String::new()
269 };
270 pb.inc(
271 format!("{}{}", chapter_infos[0].title, msg),
272 chapter_infos.len(),
273 )?;
274
275 let content_infos_multiple = client.content_infos_multiple(&chapter_infos).await?;
276 if let Some(sleep) = sleep {
277 tokio::time::sleep(Duration::from_secs(sleep)).await;
278 }
279 drop(permit);
280
281 for (index, content_infos) in content_infos_multiple.into_iter().enumerate() {
282 let mut contents = Vec::with_capacity(32);
283 for content_info in content_infos {
284 match content_info {
285 ContentInfo::Text(text) => contents.push(Content::Text(text)),
286 ContentInfo::Image(url) if !ignore_image => {
287 match client.image(&url).await {
288 Ok(image) => {
289 contents.push(Content::Image(image));
290 }
291 Err(error) => {
292 tracing::error!(
293 "Image download failed: `{error}`, url: `{url}`"
294 );
295 }
296 }
297 }
298 _ => (),
299 }
300 }
301
302 chapter_map
303 .insert_sync(chapter_infos[index].id, Some(contents))
304 .unwrap();
305 }
306
307 eyre::Ok(())
308 }));
309 }
310 }
311
312 for handle in handles {
313 handle.await??;
314 }
315
316 let chapter_map = Arc::into_inner(chapter_map).unwrap();
317 for volume in &mut novel.volumes {
318 for chapter in &mut volume.chapters {
319 if let Some((_, contents)) = chapter_map.remove_sync(&chapter.id) {
320 chapter.contents = contents;
321 }
322 }
323 }
324
325 pb.finish()?;
326
327 Ok(novel)
328}
329
330#[must_use]
331fn chapter_count(volume_infos: &VolumeInfos) -> u64 {
332 let mut count = 0;
333
334 for volume_info in volume_infos {
335 for chapter_info in &volume_info.chapter_infos {
336 if chapter_info.can_download() {
337 count += 1;
338 }
339 }
340 }
341
342 count
343}