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 let Some(cover_url) = &novel_info.cover_url
188 && !config.ignore_images
189 {
190 match client.image(cover_url).await {
191 Ok(image) => novel.cover_image = Some(image),
192 Err(error) => {
193 tracing::error!("Cover image download failed: `{error}`");
194 }
195 };
196 }
197
198 let Some(mut volume_infos) = client.volume_infos(config.novel_id).await? else {
199 eyre::bail!("Unable to get chapter information");
200 };
201
202 if config.order_novel {
203 client.order_novel(config.novel_id, &volume_infos).await?;
204 volume_infos = client.volume_infos(config.novel_id).await?.unwrap();
205
206 println!("{}", utils::locales("order_msg", "💰"));
207 }
208
209 let mut handles = Vec::with_capacity(128);
210 let pb = ProgressBar::new(chapter_count(&volume_infos))?;
211 let semaphore = Arc::new(Semaphore::new(config.maximum_concurrency as usize));
212 let chapter_map = Arc::new(HashMap::with_capacity(128));
213
214 let ignore_image = config.ignore_images;
215 let sleep = config.sleep;
216
217 for volume_info in volume_infos {
218 novel.volumes.push(Volume {
219 title: volume_info.title,
220 chapters: Vec::with_capacity(32),
221 });
222
223 let volume = novel.volumes.last_mut().unwrap();
224
225 for chapter_info in volume_info
226 .chapter_infos
227 .iter()
228 .filter(|x| !x.can_download())
229 {
230 tracing::info!(
231 "`{}-{}` can not be downloaded",
232 volume.title,
233 chapter_info.title
234 );
235 }
236
237 let can_download_chapter_infos = volume_info
238 .chapter_infos
239 .into_iter()
240 .filter(|x| x.can_download())
241 .collect::<Vec<_>>();
242
243 let chunk_size = match config.source {
244 Source::Sfacg => 1,
245 Source::Ciweimao => 10,
246 Source::Ciyuanji => 1,
247 };
248
249 for chapter_infos in can_download_chapter_infos
250 .chunks(chunk_size)
251 .map(|v| v.to_vec())
252 {
253 for chapter_info in &chapter_infos {
254 volume.chapters.push(Chapter {
255 id: chapter_info.id,
256 title: chapter_info.title.clone(),
257 contents: None,
258 });
259 }
260
261 let client = Arc::clone(&client);
262 let permit = semaphore.clone().acquire_owned().await.unwrap();
263 let mut pb = pb.clone();
264 let chapter_map = Arc::clone(&chapter_map);
265
266 handles.push(tokio::spawn(async move {
267 let msg = if chunk_size > 1 {
268 format!("等 {} 章", chunk_size)
269 } else {
270 String::new()
271 };
272 pb.inc(
273 format!("{}{}", chapter_infos[0].title, msg),
274 chapter_infos.len(),
275 )?;
276
277 let content_infos_multiple = client.content_infos_multiple(&chapter_infos).await?;
278 if let Some(sleep) = sleep {
279 tokio::time::sleep(Duration::from_secs(sleep)).await;
280 }
281 drop(permit);
282
283 for (index, content_infos) in content_infos_multiple.into_iter().enumerate() {
284 let mut contents = Vec::with_capacity(32);
285 for content_info in content_infos {
286 match content_info {
287 ContentInfo::Text(text) => contents.push(Content::Text(text)),
288 ContentInfo::Image(url) if !ignore_image => {
289 match client.image(&url).await {
290 Ok(image) => {
291 contents.push(Content::Image(image));
292 }
293 Err(error) => {
294 tracing::error!(
295 "Image download failed: `{error}`, url: `{url}`"
296 );
297 }
298 }
299 }
300 _ => (),
301 }
302 }
303
304 chapter_map
305 .insert_sync(chapter_infos[index].id, Some(contents))
306 .unwrap();
307 }
308
309 eyre::Ok(())
310 }));
311 }
312 }
313
314 for handle in handles {
315 handle.await??;
316 }
317
318 let chapter_map = Arc::into_inner(chapter_map).unwrap();
319 for volume in &mut novel.volumes {
320 for chapter in &mut volume.chapters {
321 if let Some((_, contents)) = chapter_map.remove_sync(&chapter.id) {
322 chapter.contents = contents;
323 }
324 }
325 }
326
327 pb.finish()?;
328
329 Ok(novel)
330}
331
332#[must_use]
333fn chapter_count(volume_infos: &VolumeInfos) -> u64 {
334 let mut count = 0;
335
336 for volume_info in volume_infos {
337 for chapter_info in &volume_info.chapter_infos {
338 if chapter_info.can_download() {
339 count += 1;
340 }
341 }
342 }
343
344 count
345}