use super::Config;
use crate::{Input, InputFile, InputFileType, Result, remove_xml_tags};
use either::Either;
use log::{error, info, warn};
use quick_xml::{Reader, Writer};
use rayon::prelude::*;
use std::{
borrow::Cow,
fs::{self, File},
io::{BufRead, BufReader, BufWriter, Write},
iter,
path::{Component, Path, PathBuf},
sync::{Mutex, Once},
};
use walkdir::WalkDir;
use zip::{ZipArchive, ZipWriter, write::SimpleFileOptions};
impl Config {
pub fn run(&self) -> Result<()> {
let is_found_any = Once::new();
let it = Mutex::new(self.subdirs_src_iter());
let src_dests: Vec<(InputFile, Dest)> = (0..self.jobs)
.into_par_iter()
.map(|_| job_src_dests(self, &is_found_any, &it))
.collect::<std::result::Result<Vec<_>, String>>()
.map(|vecs| vecs.into_iter().flatten().collect())?;
if !src_dests.is_empty() {
force_overwrites(src_dests);
}
if !is_found_any.is_completed() {
if let Input::Dir(d) = &self.input {
warn!("not found any fb2 in directory '{}'", d.display());
}
}
Ok(())
}
fn subdirs_src_iter(&self) -> impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)> {
match &self.input {
Input::Dir(d) if self.recursive != 0 => {
let it = WalkDir::new(d)
.max_depth(self.recursive as usize)
.into_iter()
.filter_map(|e| e.ok())
.filter(|e| e.file_type().is_dir())
.filter(|e| !e.path().starts_with(&*self.output.dir))
.map(move |e| {
let subdirs: Vec<PathBuf> = e
.path()
.strip_prefix(d)
.unwrap_or(Path::new(""))
.components()
.filter_map(|x| match x {
Component::Normal(x) => Some(PathBuf::from(x)),
_ => None,
})
.collect();
(Some(subdirs), Input::Dir(e.into_path().into()))
})
.flat_map(|(subdirs, d)| iter::repeat(subdirs).zip(d.iter()));
Either::Left(it)
}
_ => Either::Right(iter::repeat(None).zip(self.input.iter())),
}
}
}
fn job_src_dests(
cfg: &Config,
is_found_any: &Once,
it: &Mutex<impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)>>,
) -> std::result::Result<Vec<(InputFile, Dest)>, String> {
let mut src_dests: Vec<(InputFile, Dest)> = Vec::new();
let mut zip_owner: Option<ZipArchive<File>> = None;
loop {
let (src, dest) = {
let mut it = it.lock().ok();
match it.as_mut().and_then(|it| it.next()) {
Some((subdirs, src)) => {
is_found_any.call_once(|| {
info!("Cleaning input files...");
});
let dest = Dest::new(cfg, subdirs, &src);
(src, dest)
}
None => break,
}
};
info!("Cleaning '{}'...", src.path.display());
if !cfg.force && dest.path.exists() {
warn!(
"output is already exists '{}'. Skipping",
dest.path.display()
);
continue;
}
match try_reader_writer(&mut zip_owner, &src, &dest)
.and_then(|(mut r, mut w)| remove_xml_tags(&mut r, &mut w, &cfg.tags))
{
Err(e) if cfg.exit_on_err => return Err(e.to_string()),
Err(e) => {
error!("{}. Skipping", e);
continue;
}
Ok(()) => {
info!("Success cleaned and saved to '{}'", dest.path.display())
}
}
if cfg.force {
src_dests.push((src, dest));
}
}
Ok(src_dests)
}
fn force_overwrites(src_dests: Vec<(InputFile, Dest)>) {
info!("\nOverwriting input files...");
for (src, dest) in &src_dests {
info!("Overwriting '{}'...", src.path.display());
match dest.force_overwrite(&src) {
Ok(()) => info!("Success overwrited from '{}'", dest.path.display()),
Err(e) => error!("fail overwrite: {}", e),
}
}
let mut dirs: Vec<PathBuf> = Vec::new();
for (_, dest) in src_dests {
for d in dest.created_dirs.into_iter().flat_map(|ds| ds.into_iter()) {
dirs.push(d);
}
}
dirs.par_sort_by(|a, b| b.as_os_str().len().cmp(&a.as_os_str().len()));
for d in dirs {
if let Err(e) = fs::remove_dir(&d) {
error!("fail remove temp directory '{}': {}", d.display(), e);
}
}
}
fn try_reader_writer<'a>(
zip_owner: &'a mut Option<ZipArchive<File>>,
src: &'a InputFile,
dest: &'a Dest,
) -> Result<(Reader<Box<dyn BufRead + 'a>>, Writer<Box<dyn Write + 'a>>)> {
let src_file = File::open(&src.path)?;
let dest_file = File::create(&dest.path)?;
let reader = match src.ty {
InputFileType::Fb2 => {
Reader::from_reader(Box::new(BufReader::new(src_file)) as Box<dyn BufRead>)
}
InputFileType::Fb2Zip => {
*zip_owner = Some(ZipArchive::new(src_file)?);
let zip = zip_owner.as_mut().unwrap();
let fb2_index = (0..zip.len())
.find(|&i| {
zip.by_index(i)
.ok()
.map(|f| {
let bytes = f.name().as_bytes();
let len = bytes.len();
len > 3 && bytes[len - 4..len].eq_ignore_ascii_case(b".fb2")
})
.unwrap_or(false)
})
.ok_or_else(|| format!("fb2 not found in archive '{}'", src.path.display()))?;
let fb2_file = zip.by_index(fb2_index)?;
Reader::from_reader(Box::new(BufReader::new(fb2_file)) as Box<dyn BufRead>)
}
};
let writer = match dest.ty {
InputFileType::Fb2 => Writer::new(Box::new(BufWriter::new(dest_file)) as Box<dyn Write>),
InputFileType::Fb2Zip => {
let mut zip_writer = ZipWriter::new(BufWriter::new(dest_file));
zip_writer.start_file(dest.zip_start_file(), SimpleFileOptions::default())?;
Writer::new(Box::new(zip_writer) as Box<dyn Write>)
}
};
Ok((reader, writer))
}
struct Dest {
created_dirs: Option<Vec<PathBuf>>,
ty: InputFileType,
stem: PathBuf,
path: PathBuf,
}
impl Dest {
fn new(cfg: &Config, subdirs: Option<Vec<PathBuf>>, src: &InputFile) -> Dest {
let ty = if cfg.zip {
InputFileType::Fb2Zip
} else if cfg.unzip {
InputFileType::Fb2
} else {
src.ty
};
let stem = match src.path.file_stem() {
Some(os) => {
let mut p = PathBuf::from(os);
if src.ty.is_fb2_zip() {
p.set_extension("");
}
p
}
_ => PathBuf::from("cleaned"),
};
let mut path = cfg.output.dir.clone().into_path_buf();
let mut created_dirs: Option<Vec<PathBuf>> = None;
if let Some(xs) = &subdirs {
for x in xs {
path.push(x);
if let Ok(()) = fs::create_dir(&path) {
created_dirs
.get_or_insert_with(|| Vec::with_capacity(xs.len()))
.push(path.clone())
}
}
}
path.push(&stem);
path.add_extension(ty.as_extension());
Dest {
created_dirs,
ty,
stem,
path,
}
}
fn force_overwrite(&self, src: &InputFile) -> Result<()> {
let force_path = self.force_path(src);
if let Err(_) = fs::rename(&self.path, &force_path) {
fs::copy(&self.path, &force_path)?;
if let Err(e) = fs::remove_file(&self.path) {
error!("fail remove temp file '{}': {}", self.path.display(), e);
}
}
if force_path != &*src.path {
if let Err(e) = fs::remove_file(&src.path) {
error!("fail remove input file '{}': {}", src.path.display(), e);
}
}
Ok(())
}
fn force_path<'a>(&self, src: &'a InputFile) -> Cow<'a, Path> {
if self.ty == src.ty {
return Cow::Borrowed(&*src.path);
}
let mut p = match src.path.parent() {
Some(p) => PathBuf::from(p),
None => PathBuf::from("."),
};
p.push(&self.stem);
p.add_extension(self.ty.as_extension());
Cow::Owned(p)
}
fn zip_start_file(&self) -> String {
let mut s = self.stem.to_string_lossy().into_owned();
s.push_str(".fb2");
s
}
}