use crate::utils; use crate::structs; use std::env; use std::sync::Arc; use std::process::exit; use std::path::{Path, PathBuf}; use std::collections::VecDeque; use url::Url; use tokio::sync::Mutex; 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: u64 = 5000; const NO_ITEM_WAIT_TIME: u64 = 1000; const GET_MANGA_FAIL_WAIT_TIME: u64 = 30000; pub async fn run(mut args: env::Args) { let manga_id = match args.next() { Some(manga_id) => manga_id, None => { eprintln!("Missing manga id"); exit(1); } }; let mut chapter_numbers: Vec<_> = args.map(|i| i.trim().to_string()).collect(); chapter_numbers.sort(); chapter_numbers.dedup(); let mut chapters: Vec = Vec::new(); let client = reqwest::Client::new(); match utils::get_manga(client.clone(), &manga_id).await.unwrap() { structs::MangaOption::Manga(mut manga_info) => { if chapter_numbers.is_empty() { chapters = manga_info.chapters; } else { for chapter_number in chapter_numbers { let tmp = manga_info.chapters.iter().enumerate().find(|(_, chapter)| chapter_number == chapter.chapter_number); if tmp.is_some() { let (i, _) = tmp.unwrap(); chapters.push(manga_info.chapters.remove(i)); } else { eprintln!("Chapter {} does not exist", &chapter_number); exit(1); } } } }, structs::MangaOption::Redirect(_) => panic!("Nested redirect"), structs::MangaOption::DoesNotExist => { eprintln!("ID: {}\nError: does not exist", &manga_id); exit(1); } }; let mutex = Arc::new(Mutex::new(DownloadData { data: VecDeque::new(), is_done: false })); let handles: Vec> = summon_handles(client.clone(), Arc::clone(&mutex)).await; for chapter in chapters { loop { match utils::get_pages(client.clone(), &chapter, &manga_id).await { Ok(chapter_pages) => { if chapter_pages.is_empty() { sleep(Duration::from_millis(GET_MANGA_FAIL_WAIT_TIME)).await; continue; } let mut to_extend: Vec<(String, PathBuf, String)> = Vec::new(); for url in chapter_pages { let mut file_name = PathBuf::from(&chapter.chapter_number); file_name.push(Path::new(Url::parse(&url).unwrap().path()).file_name().unwrap()); if !file_name.exists() { to_extend.push((url, file_name, chapter.domain.clone())); } } if !to_extend.is_empty() { Arc::clone(&mutex).lock().await.data.extend(to_extend); } break; }, Err(_) => sleep(Duration::from_millis(GET_MANGA_FAIL_WAIT_TIME)).await }; } } { mutex.lock().await.is_done = true; } for handle in handles { handle.await.unwrap(); } } 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 (url, file_name, referer) = match data { Some(data) => data, None => { if is_done { break; } sleep(Duration::from_millis(NO_ITEM_WAIT_TIME)).await; continue; } }; eprintln!("[DW{}] Downloading {} to {}", worker_id, &url, file_name.display()); loop { match utils::download_file(cloned_client.clone(), &url, &file_name, &referer).await { Ok(result) if result => break, Ok(_) => (), Err(err) => eprintln!("[DW{}] Error while downloading {}: {}", worker_id, file_name.display(), err) }; sleep(Duration::from_millis(NON_IMAGE_WAIT_TIME)).await; } eprintln!("[DW{}] Downloaded {} to {}", worker_id, &url, file_name.display()); } eprintln!("[DW{}] Down!", worker_id); })); } handles } struct DownloadData { pub data: VecDeque<(String, PathBuf, String)>, pub is_done: bool }