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, to: impl AsRef, tags: impl AsRef<[tags::Tag]>, progress: &mut progress::CommandSender) -> Result { 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, path: impl AsRef, 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> { 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::::new("Initialising...", 1); let prog_writer = prog.writer(); let prog = prog.host(); (prog_writer, prog) }, config::Verbosity::Silent => { let prog = progress::AsyncProgressCounter::::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(()) }