use super::Config;
use crate::{Input, InputFile, InputFileType, Msg, Result, display, 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 jobs = self.jobs.min(Self::default_jobs());
let src_dests: Vec<(InputFile, Dest)> = (0..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 {
return Err(format!("{} '{}'", Msg::NotFoundAnyBookInDirectory, display(d)).into());
}
}
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_entry(|entry| {
if !entry.file_type().is_dir() {
return false;
}
let entry_bytes = entry.path().as_os_str().as_encoded_bytes();
let out_bytes = self.output.dir.as_os_str().as_encoded_bytes();
if entry_bytes.len() == out_bytes.len() && entry_bytes.ends_with(out_bytes)
{
return false;
}
true
})
.filter_map(|e| e.ok())
.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!("{}...", Msg::CleaningBooks);
});
let dest = Dest::new(cfg, subdirs, &src);
(src, dest)
}
None => break,
}
};
info!("{} '{}'...", Msg::Cleaning, display(&src.path));
if !cfg.overwrite && dest.path.exists() {
warn!(
"{} '{}'. {}",
Msg::FileAlreadyExists,
display(&dest.path),
Msg::Skipping,
);
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!("{}. {}", e, Msg::Skipping);
continue;
}
Ok(()) => {
info!("{} '{}'", Msg::SuccessfullyCleanedTo, display(&dest.path))
}
}
if cfg.overwrite {
src_dests.push((src, dest));
}
}
Ok(src_dests)
}
fn force_overwrites(src_dests: Vec<(InputFile, Dest)>) {
info!("\n{}...", Msg::OverwritingBooks);
for (src, dest) in &src_dests {
info!("{} '{}'...", Msg::Overwriting, display(&src.path));
match dest.force_overwrite(&src) {
Ok(()) => info!(
"{} '{}'",
Msg::SuccessfullyOverwritedFrom,
display(&dest.path)
),
Err(e) => error!("{}: {}", Msg::Overwriting, 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!("{} '{}': {}", Msg::RemovingTempDirectory, display(&d), 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| {
f.name().is_ok_and(|name| {
let bytes = 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!("{} '{}'", Msg::NotFoundAnyFb2InArchive, display(&src.path))
})?;
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!("{} '{}': {}", Msg::RemovingTempFile, display(&self.path), e);
}
}
if force_path != &*src.path {
if let Err(e) = fs::remove_file(&src.path) {
error!("{} '{}': {}", Msg::RemovingInputFile, display(&src.path), 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
}
}