hanimers/src/commands/download.rs

use crate::utils;

use std::fs::rename;
use std::path::Path;
use std::process::{exit, Command};
use std::collections::HashMap;
use clap::ArgMatches;
use reqwest::Client;
use reqwest::redirect::Policy;

const MAX_DOWNLOAD_ATTEMPTS: i32 = 5;

pub async fn download(arg_m: &ArgMatches<'_>) {
    let print_only = arg_m.is_present("print");
    let resolution = arg_m.value_of("resolution");
    let ids = arg_m.values_of("id").unwrap().collect::<Vec<_>>();
    let policy = Policy::custom(|attempt| {
        if attempt.previous().len() > 10 {
            attempt.error("too many redirects")
        } else if attempt.url().path() == "/404" {
            attempt.stop()
        } else {
            attempt.follow()
        }
    });
    let client = Client::builder().redirect(policy).build().unwrap();
    let mut return_fail = false;
    for id in ids {
        let hentai_info = utils::get_hentai(client.clone(), id).await;
        match hentai_info {
            Ok(hentai_info) => {
                match hentai_info {
                    Some(hentai_info) => {
                        let slug = hentai_info.state.data.video.hentai_video.slug;
                        let filename = format!("{}.mkv", &slug);
                        if !print_only && Path::new(&filename).exists() {
                            continue;
                        }
                        let mut download_url = None;
                        let mut perm_urls: HashMap<String, (String, i32)> = HashMap::new();
                        let mut temp_urls: HashMap<String, (String, i32)> = HashMap::new();
                        for server in hentai_info.state.data.video.videos_manifest.servers {
                            let mut to_hashmap = match server.is_permanent {
                                true => perm_urls.clone(),
                                false => temp_urls.clone()
                            };
                            for stream in server.streams {
                                if stream.url.is_empty() {
                                    continue;
                                }
                                if server.is_permanent && Some(stream.height.as_str()) == resolution {
                                    download_url = Some((stream.url, stream.filesize_mbs));
                                    break;
                                }
                                if !to_hashmap.contains_key(&stream.height) {
                                    to_hashmap.insert(stream.height, (stream.url, stream.filesize_mbs));
                                };
                            }
                            if download_url.is_some() {
                                break;
                            }
                            match server.is_permanent {
                                true => perm_urls.extend(to_hashmap),
                                false => temp_urls.extend(to_hashmap)
                            };
                        }
                        if download_url.is_none() {
                            if resolution.is_some() {
                                download_url = temp_urls.get(resolution.unwrap()).map(|i| (i.0.to_string(), i.1));
                            }
                            if download_url.is_none() {
                                download_url = magic_thing(perm_urls).or_else(|| { magic_thing(temp_urls) });
                                if download_url.is_none() {
                                    eprintln!("Failed to get {}: cannot get download url", id);
                                    return_fail = true;
                                    continue;
                                }
                            }
                        }
                        let (download_url, filesize_mbs) = download_url.unwrap();
                        if print_only {
                            println!("{}", download_url);
                        } else {
                            let mut fail_dl = true;
                            let tmp_filename = format!("{}.tmp", &slug);
                            for i in 0..MAX_DOWNLOAD_ATTEMPTS {
                                eprintln!("Downloading {} ({}MB, attempt {})", &filename, filesize_mbs, i);
                                match Command::new("ffmpeg").args(&["-v", "warning", "-stats", "-nostdin", "-y", "-i", &download_url, "-c", "copy", "-f", "matroska", &tmp_filename]).spawn() {
                                    Ok(mut child) => {
                                        match child.wait() {
                                            Ok(exit_status) => {
                                                if exit_status.success() {
                                                    fail_dl = false;
                                                    match rename(&tmp_filename, &filename) {
                                                        Ok(_) => (),
                                                        Err(err) => eprintln!("Failed to rename {} to {} due to {}", &tmp_filename, &filename, err)
                                                    };
                                                    break;
                                                }
                                                eprintln!("ffmpeg exited with {:?}", exit_status.code());
                                            },
                                            Err(err) => eprintln!("Failed to wait on ffmpeg process due to {}", err)
                                        };
                                    },
                                    Err(err) => eprintln!("Failed to spawn ffmpeg process due to {}", err)
                                };
                            }
                            if fail_dl {
                                eprintln!("Failed to download {}", &filename);
                                return_fail = true;
                            }
                        }
                    },
                    None => {
                        eprintln!("Failed to get {}: does not exist", id);
                        return_fail = true;
                    }
                };
            },
            Err(err) => {
                eprintln!("Failed to get {}: {}", id, err);
                return_fail = true;
            }
        };
    }
    if return_fail {
        exit(1);
    }
}

fn magic_thing(map: HashMap<String, (String, i32)>) -> Option<(String, i32)> {
    let mut keys = Vec::new();
    for i in map.keys() {
        match i.parse::<i32>() {
            Ok(i) => keys.push(i),
            Err(_) => {
                keys.clear();
                break;
            }
        };
    }
    keys.sort();
    match keys.pop() {
        Some(key) => map.get(&key.to_string()).map(|i| i.clone()),
        None => {
            let mut keys: Vec<_> = map.keys().collect();
            keys.sort();
            match keys.pop() {
                Some(key) => map.get(key.as_str()).map(|i| i.clone()),
                None => None
            }
        }
    }
}