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