mangafetchi/src/commands/download.rs

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<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 == 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<JoinHandle<()>> = 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<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 (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
}