mangadexrs/src/commands/download.rs

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::<i32>().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::<i32>().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<i32>, mutex: Arc<Mutex<DownloadData>>, 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::<Vec<_>>().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<Mutex<DownloadData>>) -> 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 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<RwLock<(structs::ChapterData, i32)>>, usize, PathBuf)>,
    pub is_done: bool
}