fb2_clean/config/

run.rs

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 {
    /// Run for current [`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| {
                        // skip non-directory
                        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();

                        // skip output directory
                        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
    }
}