1use super::Config;
2use crate::{Input, InputFile, InputFileType, Msg, Result, remove_xml_tags};
3use either::Either;
4use log::{error, info, warn};
5use quick_xml::{Reader, Writer};
6use rayon::prelude::*;
7use std::{
8 borrow::Cow,
9 fs::{self, File},
10 io::{BufRead, BufReader, BufWriter, Write},
11 iter,
12 path::{Component, Path, PathBuf},
13 sync::{Mutex, Once},
14};
15use walkdir::WalkDir;
16use zip::{ZipArchive, ZipWriter, write::SimpleFileOptions};
17
18impl Config {
19 pub fn run(&self) -> Result<()> {
21 let is_found_any = Once::new();
22 let it = Mutex::new(self.subdirs_src_iter());
23 let jobs = self.jobs.min(Self::default_jobs());
24
25 let src_dests: Vec<(InputFile, Dest)> = (0..jobs)
26 .into_par_iter()
27 .map(|_| job_src_dests(self, &is_found_any, &it))
28 .collect::<std::result::Result<Vec<_>, String>>()
29 .map(|vecs| vecs.into_iter().flatten().collect())?;
30
31 if !src_dests.is_empty() {
32 force_overwrites(src_dests);
33 }
34
35 if !is_found_any.is_completed() {
36 if let Input::Dir(d) = &self.input {
37 return Err(
38 format!("{} '{}'", Msg::NotFoundAnyBookInDirectory, d.display()).into(),
39 );
40 }
41 }
42
43 Ok(())
44 }
45
46 fn subdirs_src_iter(&self) -> impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)> {
47 match &self.input {
48 Input::Dir(d) if self.recursive != 0 => {
49 let it = WalkDir::new(d)
50 .max_depth(self.recursive as usize)
51 .into_iter()
52 .filter_map(|e| e.ok())
53 .filter(|e| e.file_type().is_dir())
54 .filter(|e| !e.path().starts_with(&*self.output.dir))
55 .map(move |e| {
56 let subdirs: Vec<PathBuf> = e
57 .path()
58 .strip_prefix(d)
59 .unwrap_or(Path::new(""))
60 .components()
61 .filter_map(|x| match x {
62 Component::Normal(x) => Some(PathBuf::from(x)),
63 _ => None,
64 })
65 .collect();
66
67 (Some(subdirs), Input::Dir(e.into_path().into()))
68 })
69 .flat_map(|(subdirs, d)| iter::repeat(subdirs).zip(d.iter()));
70
71 Either::Left(it)
72 }
73 _ => Either::Right(iter::repeat(None).zip(self.input.iter())),
74 }
75 }
76}
77
78fn job_src_dests(
79 cfg: &Config,
80 is_found_any: &Once,
81 it: &Mutex<impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)>>,
82) -> std::result::Result<Vec<(InputFile, Dest)>, String> {
83 let mut src_dests: Vec<(InputFile, Dest)> = Vec::new();
84 let mut zip_owner: Option<ZipArchive<File>> = None;
85
86 loop {
87 let (src, dest) = {
88 let mut it = it.lock().ok();
89 match it.as_mut().and_then(|it| it.next()) {
90 Some((subdirs, src)) => {
91 is_found_any.call_once(|| {
92 info!("{}...", Msg::CleaningBooks);
93 });
94 let dest = Dest::new(cfg, subdirs, &src);
95 (src, dest)
96 }
97 None => break,
98 }
99 };
100 info!("{} '{}'...", Msg::Cleaning, src.path.display());
101
102 if !cfg.overwrite && dest.path.exists() {
103 warn!(
104 "{} '{}'. {}",
105 Msg::FileIsAlreadyExists,
106 dest.path.display(),
107 Msg::Skipping,
108 );
109 continue;
110 }
111
112 match try_reader_writer(&mut zip_owner, &src, &dest)
113 .and_then(|(mut r, mut w)| remove_xml_tags(&mut r, &mut w, &cfg.tags))
114 {
115 Err(e) if cfg.exit_on_err => return Err(e.to_string()),
116 Err(e) => {
117 error!("{}. {}", e, Msg::Skipping);
118 continue;
119 }
120 Ok(()) => {
121 info!(
122 "{} '{}'",
123 Msg::SuccessCleanedAndSavedTo,
124 dest.path.display()
125 )
126 }
127 }
128
129 if cfg.overwrite {
130 src_dests.push((src, dest));
131 }
132 }
133 Ok(src_dests)
134}
135
136fn force_overwrites(src_dests: Vec<(InputFile, Dest)>) {
137 info!("\n{}...", Msg::OverwritingBooks);
138
139 for (src, dest) in &src_dests {
140 info!("{} '{}'...", Msg::Overwriting, src.path.display());
141 match dest.force_overwrite(&src) {
142 Ok(()) => info!("{} '{}'", Msg::SuccessOverwritedFrom, dest.path.display()),
143 Err(e) => error!("{}: {}", Msg::Overwriting, e),
144 }
145 }
146
147 let mut dirs: Vec<PathBuf> = Vec::new();
148 for (_, dest) in src_dests {
149 for d in dest.created_dirs.into_iter().flat_map(|ds| ds.into_iter()) {
150 dirs.push(d);
151 }
152 }
153 dirs.par_sort_by(|a, b| b.as_os_str().len().cmp(&a.as_os_str().len()));
154
155 for d in dirs {
156 if let Err(e) = fs::remove_dir(&d) {
157 error!("{} '{}': {}", Msg::RemovingTempDirectory, d.display(), e);
158 }
159 }
160}
161
162fn try_reader_writer<'a>(
163 zip_owner: &'a mut Option<ZipArchive<File>>,
164 src: &'a InputFile,
165 dest: &'a Dest,
166) -> Result<(Reader<Box<dyn BufRead + 'a>>, Writer<Box<dyn Write + 'a>>)> {
167 let src_file = File::open(&src.path)?;
168 let dest_file = File::create(&dest.path)?;
169
170 let reader = match src.ty {
171 InputFileType::Fb2 => {
172 Reader::from_reader(Box::new(BufReader::new(src_file)) as Box<dyn BufRead>)
173 }
174 InputFileType::Fb2Zip => {
175 *zip_owner = Some(ZipArchive::new(src_file)?);
176 let zip = zip_owner.as_mut().unwrap();
177 let fb2_index = (0..zip.len())
178 .find(|&i| {
179 zip.by_index(i)
180 .ok()
181 .map(|f| {
182 let bytes = f.name().as_bytes();
183 let len = bytes.len();
184 len > 3 && bytes[len - 4..len].eq_ignore_ascii_case(b".fb2")
185 })
186 .unwrap_or(false)
187 })
188 .ok_or_else(|| {
189 format!("{} '{}'", Msg::NotFoundAnyFb2InArchive, src.path.display())
190 })?;
191
192 let fb2_file = zip.by_index(fb2_index)?;
193 Reader::from_reader(Box::new(BufReader::new(fb2_file)) as Box<dyn BufRead>)
194 }
195 };
196
197 let writer = match dest.ty {
198 InputFileType::Fb2 => Writer::new(Box::new(BufWriter::new(dest_file)) as Box<dyn Write>),
199 InputFileType::Fb2Zip => {
200 let mut zip_writer = ZipWriter::new(BufWriter::new(dest_file));
201 zip_writer.start_file(dest.zip_start_file(), SimpleFileOptions::default())?;
202 Writer::new(Box::new(zip_writer) as Box<dyn Write>)
203 }
204 };
205
206 Ok((reader, writer))
207}
208
209struct Dest {
210 created_dirs: Option<Vec<PathBuf>>,
211 ty: InputFileType,
212 stem: PathBuf,
213 path: PathBuf,
214}
215
216impl Dest {
217 fn new(cfg: &Config, subdirs: Option<Vec<PathBuf>>, src: &InputFile) -> Dest {
218 let ty = if cfg.zip {
219 InputFileType::Fb2Zip
220 } else if cfg.unzip {
221 InputFileType::Fb2
222 } else {
223 src.ty
224 };
225
226 let stem = match src.path.file_stem() {
227 Some(os) => {
228 let mut p = PathBuf::from(os);
229 if src.ty.is_fb2_zip() {
230 p.set_extension("");
231 }
232 p
233 }
234 _ => PathBuf::from("cleaned"),
235 };
236
237 let mut path = cfg.output.dir.clone().into_path_buf();
238 let mut created_dirs: Option<Vec<PathBuf>> = None;
239
240 if let Some(xs) = &subdirs {
241 for x in xs {
242 path.push(x);
243 if let Ok(()) = fs::create_dir(&path) {
244 created_dirs
245 .get_or_insert_with(|| Vec::with_capacity(xs.len()))
246 .push(path.clone())
247 }
248 }
249 }
250 path.push(&stem);
251 path.add_extension(ty.as_extension());
252
253 Dest {
254 created_dirs,
255 ty,
256 stem,
257 path,
258 }
259 }
260
261 fn force_overwrite(&self, src: &InputFile) -> Result<()> {
262 let force_path = self.force_path(src);
263
264 if let Err(_) = fs::rename(&self.path, &force_path) {
265 fs::copy(&self.path, &force_path)?;
266 if let Err(e) = fs::remove_file(&self.path) {
267 error!("{} '{}': {}", Msg::RemovingTempFile, self.path.display(), e);
268 }
269 }
270
271 if force_path != &*src.path {
272 if let Err(e) = fs::remove_file(&src.path) {
273 error!("{} '{}': {}", Msg::RemovingInputFile, src.path.display(), e);
274 }
275 }
276 Ok(())
277 }
278
279 fn force_path<'a>(&self, src: &'a InputFile) -> Cow<'a, Path> {
280 if self.ty == src.ty {
281 return Cow::Borrowed(&*src.path);
282 }
283
284 let mut p = match src.path.parent() {
285 Some(p) => PathBuf::from(p),
286 None => PathBuf::from("."),
287 };
288 p.push(&self.stem);
289 p.add_extension(self.ty.as_extension());
290 Cow::Owned(p)
291 }
292
293 fn zip_start_file(&self) -> String {
294 let mut s = self.stem.to_string_lossy().into_owned();
295 s.push_str(".fb2");
296 s
297 }
298}