lolistealer/src/work_async.rs

use super::*;
use std::{
    path::Path,
};
use tokio::{
    fs::{
	OpenOptions,
    },
    prelude::*,
    stream::StreamExt,
};

mod tasklist;
#[macro_use]
mod progress;

/// Decode a loli from path
pub async fn decode(from: impl AsRef<Path>, to: impl AsRef<Path>, tags: impl AsRef<[tags::Tag]>, progress: &mut progress::CommandSender) -> Result<loli::Loli, error::Error>
{
    prog_send!(progress.println("Mapping child"));
    let base = loli::BasedLoli::map(from)?;
    prog_send!(progress.println("Calculating bounds"));
    let bounds = base.calculate_bounds()?; // If server is returning error code, this will fail.
    prog_send!(progress.println(format!("Finding bounds ({:?}) {} -> {} bytes", bounds, base.as_ref().len(), base.decoded_size())));

    //Find extension
    
    let mut decoded = bounds.create_child(to.as_ref().with_extension(bounds.image().ext()))?;

    let tags=tags.as_ref();
    if tags.len() > 0 {
	let res = tags::search(decoded.tags().iter(), tags)?;
	if res.len() > 0 {
	    prog_send!(progress.println(format!("Matched tags {}", res.into_iter().map(|x| format!("{}", x)).join(", "))));
	}
    }
    prog_send!(progress.println("Decoding..."));    
    let sz = bounds.decode(&mut decoded)?;
    prog_send!(link progress.println(format!("Decode complete ({} bytes)", sz)));

    Ok(decoded)
}

/// Download a loli async
pub async fn perform(url: impl AsRef<str>, path: impl AsRef<Path>, progress: &mut progress::CommandSender) -> Result<(), error::Error>
{
    let url = url.as_ref();
    let path = path.as_ref();


    let resp = reqwest::get(url).await?;
    let len = resp.content_length();

    //prog_send!(link progress.push_task(&task));
    if let Some(len) = len {
	prog_send!(progress.bump_max(len));
    } else {
	prog_send!(progress.bump_max(1));
    }
    
    let mut file = OpenOptions::new()
	.create(true)
	.truncate(true)
	.write(true)
	.read(true)
	.open(path).await?;
    
    prog_send!(progress.println(format!("Starting download of {:?} bytes to {:?}", len, path)));
    let mut bytes = resp.bytes_stream();
    while let Some(buffer) = bytes.next().await {
	let slice = buffer?;
	let slice = slice.as_ref();
	file.write(slice).await?;

	if let Some(_) = len {
	    prog_send!(progress.bump(slice.len() as u64));
	}
    }
    file.flush().await?;
    if len.is_none() {
	prog_send!(progress.bump(1));
    }
    prog_send!(progress.println(format!("done for {}", url)));
    //prog_send!(link progress.pop_task(task));

    Ok(())
}


pub async fn work(conf: config::Config) -> Result<(), Box<dyn std::error::Error>>
{
    let rating = conf.rating;
    let mut children = Vec::with_capacity(conf.output.len());

    let (mut prog_writer, prog) = match conf.verbose {
	config::Verbosity::Full => {
	    let prog = progress::AsyncProgressCounter::<termprogress::progress::Bar>::new("Initialising...", 1);
	    let prog_writer = prog.writer();
	    let prog = prog.host();
	    (prog_writer, prog)
	},
	config::Verbosity::Silent => {
	    let prog = progress::AsyncProgressCounter::<progress::Silent>::new("Initialising...", 1);
	    let prog_writer = prog.writer();
	    let prog = prog.host();
	    (prog_writer, prog)
	},
    };
    
    
    for path in conf.output.into_iter()
    {
	let url = url::parse(&rating);
	let mut prog = prog_writer.clone_with(format!("-> {:?}", path));
	let tags = conf.tags.clone();
	children.push(tokio::task::spawn(async move {
	    prog.println(format!("Starting download ({})...", url)).await.expect("fatal");
	    
	    let task = format!("{:?}", path); //TODO: Real task name
	    prog_send!(link unwind prog.push_task(&task));
	    
	    let temp = tempfile::TempFile::new();
	    let return_value = loop {
		match perform(&url, &temp, &mut prog).await {
		    Err(e) => prog_send!(link prog.println(format!("Failed downloading {} -> {:?}: {}", url, temp, e))),
		    Ok(_) => {
			let path = match path {
			    config::OutputType::File(file) => file,
			    config::OutputType::Directory(dir) => unimplemented!(), //TODO: implement get hash to file
			};
			let loli = match decode(&temp, &path, &tags, &mut prog).await {
			    Ok(v) => v,
			    Err(e) => {
				prog_send!(link prog.println(format!("Failed decoding: {}", e)));
				break  Some(e);
			    },
			};

			prog_send!(link prog.println(format!("{:?} Complete", loli)));
			break None;
		    },
		};
		break Some(error::Error::Unknown);
	    };

	    prog_send!(link prog.pop_task(task));
	    return_value
	}));
    }

    prog_send!(link prog_writer.println("Children working..."));
    let mut done =0;
    let total = children.len();
    let mut failures = Vec::with_capacity(children.len());
    for child in children.into_iter()
    {
	match child.await {
	    Ok(None) => done+=1,
	    Ok(Some(err)) => failures.push(err),
	    Err(err) => {
		prog_send!(try link unwind prog_writer.println(format!("Child panic: {}", err)));
		failures.push(error::Error::ChildPanic);
	    },
	}
    }
    prog_send!(link prog_writer.set_title(""));
    prog_send!(try link prog_writer.kill());
    prog.await.expect("mpsc fatal");
    println!("Completed {} / {} lolis ({} failed).", done, total, total-done);
    if failures.len() > 0 {
	println!("Reasons for failure(s):");
	for failure in failures.into_iter() {
	    eprintln!("\t{}", failure);
	}
    }
    
    Ok(())
}