101 lines
3.8 KiB
Rust
101 lines
3.8 KiB
Rust
use crate::utils;
|
|
use crate::structs;
|
|
|
|
use std::env;
|
|
use std::sync::Arc;
|
|
use std::process::exit;
|
|
use std::path::{Path, PathBuf};
|
|
use tokio::sync::Mutex;
|
|
use tokio::task::JoinHandle;
|
|
extern crate tokio;
|
|
extern crate reqwest;
|
|
|
|
const DOWNLOAD_WORKERS: usize = 5;
|
|
|
|
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.collect();
|
|
chapter_numbers.sort();
|
|
chapter_numbers.dedup();
|
|
let mut chapters: Vec<structs::Chapter> = 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.trim() == chapter.chapter_number.as_str());
|
|
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::None => {
|
|
eprintln!("ID: {}\nError: does not exist", &manga_id);
|
|
exit(1);
|
|
}
|
|
};
|
|
let mutex: Arc<Mutex<Vec<(String, PathBuf, String)>>> = Arc::new(Mutex::new(Vec::new()));
|
|
let mut handles: Option<Vec<JoinHandle<()>>> = None;
|
|
for chapter in chapters {
|
|
let cloned_mutex = Arc::clone(&mutex);
|
|
let chapter_pages = utils::get_pages(client.clone(), &chapter, &manga_id).await.unwrap();
|
|
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(reqwest::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() {
|
|
cloned_mutex.lock().await.extend(to_extend);
|
|
}
|
|
if handles.is_none() {
|
|
handles = Some(summon_handles(client.clone(), cloned_mutex).await);
|
|
}
|
|
}
|
|
for handle in handles.unwrap() {
|
|
handle.await.unwrap();
|
|
}
|
|
}
|
|
|
|
async fn summon_handles(client: reqwest::Client, mutex: Arc<Mutex<Vec<(String, PathBuf, String)>>>) -> Vec<JoinHandle<()>> {
|
|
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 vec = cloned_mutex.lock().await;
|
|
if vec.is_empty() {
|
|
break;
|
|
}
|
|
let (url, file_name, referer) = vec.remove(0);
|
|
drop(vec);
|
|
eprintln!("[DW{}] Downloading {} to {}", worker_id, &url, file_name.display());
|
|
utils::download_file(cloned_client, &url, &file_name, &referer).await.unwrap();
|
|
eprintln!("[DW{}] Downloaded {} to {}", worker_id, &url, file_name.display());
|
|
}
|
|
eprintln!("[DW{}] Down!", worker_id);
|
|
}));
|
|
}
|
|
handles
|
|
}
|