use crate::utils; use crate::structs; use std::sync::Arc; use std::fs::create_dir_all; use std::collections::VecDeque; use std::path::{Path, PathBuf}; use std::process::exit; use clap::ArgMatches; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinHandle; use tokio::time::{sleep, Duration}; extern crate tokio; extern crate reqwest; const DOWNLOAD_WORKERS: usize = 5; const NON_IMAGE_WAIT_TIME: Duration = Duration::from_millis(5000); const NO_ITEM_WAIT_TIME: Duration = Duration::from_millis(1000); const GET_DATA_FAIL_WAIT_TIME: Duration = Duration::from_millis(30000); pub async fn download(arg_m: &ArgMatches<'_>) { let print_only = arg_m.is_present("print"); let languages: Vec<_> = arg_m.values_of("language").unwrap_or_default().collect(); let mut chapter_ids: Vec<_> = arg_m.values_of("chapter_ids").unwrap_or_default().map(|i| i.parse::().unwrap()).collect(); chapter_ids.sort(); chapter_ids.dedup(); let mut manga_ids: Vec<_> = arg_m.values_of("manga_ids").unwrap_or_default().map(|i| i.parse::().unwrap()).collect(); manga_ids.sort(); manga_ids.dedup(); let client = reqwest::Client::new(); let mut return_fail = false; let mutex = Arc::new(Mutex::new(DownloadData { data: VecDeque::new(), is_done: false })); let handles = summon_handles(client.clone(), Arc::clone(&mutex)).await; handle_chapters(client.clone(), chapter_ids.clone(), Arc::clone(&mutex), &mut return_fail, print_only).await; for manga_id in manga_ids { let cloned_client = client.clone(); let cloned_mutex = Arc::clone(&mutex); loop { let mut manga = match utils::get_manga(cloned_client.clone(), manga_id).await { Ok(Some(i)) => i, Ok(None) => { eprintln!("Manga ID: {}\nError: does not exist", manga_id); return_fail = true; break; }, Err(err) => { eprintln!("Manga ID: {}\nError: {}", manga_id, err); sleep(GET_DATA_FAIL_WAIT_TIME).await; continue; } }; let mut mchapter_ids = Vec::new(); manga.data.chapters.reverse(); for chapter in manga.data.chapters { if chapter_ids.contains(&chapter.id) || (!languages.is_empty() && !languages.contains(&chapter.language.as_str())) { continue; } mchapter_ids.push(chapter.id); } if !mchapter_ids.is_empty() { handle_chapters(cloned_client.clone(), mchapter_ids, Arc::clone(&cloned_mutex), &mut return_fail, print_only).await; } break; } } mutex.lock().await.is_done = true; for handle in handles { handle.await.unwrap(); } if return_fail { exit(1); } } async fn handle_chapters(client: reqwest::Client, chapter_ids: Vec, mutex: Arc>, return_fail: &mut bool, print_only: bool) { let mut chapter_datas = Vec::new(); for chapter_id in chapter_ids { let cloned_client = client.clone(); loop { let chapter = match utils::get_chapter(cloned_client.clone(), chapter_id).await { Ok(Some(i)) => i, Ok(None) => { eprintln!("Chapter ID: {}\nError: does not exist", chapter_id); *return_fail = true; break; }, Err(err) => { eprintln!("Chapter ID: {}\nError: {}", chapter_id, err); sleep(GET_DATA_FAIL_WAIT_TIME).await; continue; } }; if print_only { println!("{}", chapter.data.pages.iter().map(|i| format!("{}{}/{}", &chapter.data.server, &chapter.data.hash, i)).collect::>().join(";")); } else { let manga_slug = utils::generate_slug(&chapter.data.manga_title); let local_dir: PathBuf = [&manga_slug, &chapter.data.chapter].iter().collect(); let chapter_pages = chapter.data.pages.clone(); let chapter_rwlock = Arc::new(RwLock::new((chapter.data, 0))); for (i, server_file) in chapter_pages.iter().enumerate() { let mut local_file = local_dir.clone(); local_file.push((i + 1).to_string()); if let Some(ext) = Path::new(&server_file).extension() { local_file.set_extension(ext); } if local_file.exists() { continue; } chapter_datas.push((Arc::clone(&chapter_rwlock), i, local_file)); } } break; } } if !chapter_datas.is_empty() { mutex.lock().await.data.extend(chapter_datas); } } async fn summon_handles(client: reqwest::Client, mutex: Arc>) -> Vec> { let mut handles = Vec::with_capacity(DOWNLOAD_WORKERS); for worker_id in 0..DOWNLOAD_WORKERS { let tcloned_mutex = Arc::clone(&mutex); let tcloned_client = client.clone(); handles.push(tokio::spawn(async move { eprintln!("[DW{}] Up!", worker_id); loop { let cloned_mutex = Arc::clone(&tcloned_mutex); let cloned_client = tcloned_client.clone(); let mut download_data = cloned_mutex.lock().await; let data = download_data.data.pop_front(); let is_done = download_data.is_done; drop(download_data); let (chapter_rwlock, page_index, local_file) = match data { Some(data) => data, None => { if is_done { break; } sleep(NO_ITEM_WAIT_TIME).await; continue; } }; let read_guard = chapter_rwlock.read().await; let chapter_id = read_guard.0.id; let mut server = read_guard.0.server.clone(); let mut server_fallback = read_guard.0.server_fallback.clone(); let mut hash = read_guard.0.hash.clone(); let mut server_file = read_guard.0.pages[page_index].clone(); let mut local_update_count = read_guard.1; drop(read_guard); if let Some(parent) = local_file.parent() { match create_dir_all(&parent) { Ok(()) => (), Err(err) => eprintln!("[DW{}] Error while creating directories {}: {}", worker_id, parent.display(), err) }; } eprintln!("[DW{}] Downloading {}", worker_id, local_file.display()); loop { match utils::download_page(cloned_client.clone(), &server, &server_fallback, &hash, &server_file, &local_file).await { Ok(true) => break, Ok(false) => eprintln!("[DW{}] Error while downloading {}: does not exist", worker_id, local_file.display()), Err(err) => eprintln!("[DW{}] Error while downloading {}: {}", worker_id, local_file.display(), err) }; sleep(NON_IMAGE_WAIT_TIME).await; let mut write_guard = chapter_rwlock.write().await; if local_update_count >= write_guard.1 { match utils::get_chapter(cloned_client.clone(), chapter_id).await { Ok(Some(chapter_data)) => { write_guard.0 = chapter_data.data; write_guard.1 = local_update_count + 1; }, Ok(None) => eprintln!("[DW{}] Error while fetching chapter data {}: does not exist", worker_id, chapter_id), Err(err) => eprintln!("[DW{}] Error while fetching chapter data {}: {}", worker_id, chapter_id, err) }; } server = write_guard.0.server.clone(); server_fallback = write_guard.0.server_fallback.clone(); hash = write_guard.0.hash.clone(); server_file = write_guard.0.pages[page_index].clone(); local_update_count = write_guard.1; } eprintln!("[DW{}] Downloaded {}", worker_id, local_file.display()); } eprintln!("[DW{}] Down!", worker_id); })); } handles } struct DownloadData { pub data: VecDeque<(Arc>, usize, PathBuf)>, pub is_done: bool }