Skip to main content

audiobook_forge/core/
processor.rs

1//! Single book processor
2
3use crate::audio::{
4    generate_chapters_from_files, inject_chapters_mp4box, inject_metadata_atomicparsley,
5    parse_cue_file, write_mp4box_chapters, AacEncoder, FFmpeg,
6};
7use crate::models::{BookFolder, ProcessingResult};
8use anyhow::{Context, Result};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use std::time::Instant;
12use tokio::sync::Semaphore;
13
14/// Processor for converting a single audiobook
15pub struct Processor {
16    ffmpeg: FFmpeg,
17    keep_temp: bool,
18    encoder: AacEncoder,
19    enable_parallel_encoding: bool,
20    max_concurrent_files: usize,
21    quality_preset: Option<String>,
22}
23
24impl Processor {
25    /// Create a new processor
26    pub fn new() -> Result<Self> {
27        Ok(Self {
28            ffmpeg: FFmpeg::new()?,
29            keep_temp: false,
30            encoder: crate::audio::get_encoder(),
31            enable_parallel_encoding: true,
32            max_concurrent_files: 8,
33            quality_preset: None,
34        })
35    }
36
37    /// Create processor with options
38    pub fn with_options(
39        keep_temp: bool,
40        encoder: AacEncoder,
41        enable_parallel_encoding: bool,
42        max_concurrent_files: usize,
43        quality_preset: Option<String>,
44    ) -> Result<Self> {
45        Ok(Self {
46            ffmpeg: FFmpeg::new()?,
47            keep_temp,
48            encoder,
49            enable_parallel_encoding,
50            max_concurrent_files: max_concurrent_files.clamp(1, 32),
51            quality_preset,
52        })
53    }
54
55    /// Process a single book folder
56    pub async fn process_book(
57        &self,
58        book_folder: &BookFolder,
59        output_dir: &Path,
60        chapter_source: &str,
61    ) -> Result<ProcessingResult> {
62        let start_time = Instant::now();
63        let result = ProcessingResult::new(book_folder.name.clone());
64
65        tracing::info!("=== Starting book processing: {} ===", book_folder.name);
66
67        // Create output directory if it doesn't exist
68        if !output_dir.exists() {
69            std::fs::create_dir_all(output_dir)
70                .context("Failed to create output directory")?;
71        }
72
73        // Determine output file path
74        let output_filename = book_folder.get_output_filename();
75        let output_path = output_dir.join(&output_filename);
76
77        // Create temp directory
78        let temp_dir = self.create_temp_dir(&book_folder.name)?;
79
80        // Check if we can use copy mode
81        let use_copy = book_folder.can_use_concat_copy();
82
83        tracing::info!(
84            "Processing {} - {} tracks, copy_mode={}",
85            book_folder.name,
86            book_folder.tracks.len(),
87            use_copy
88        );
89
90        // Get quality profile (auto-detected from source)
91        let mut quality = book_folder
92            .get_best_quality_profile(true)
93            .context("No tracks found")?
94            .clone();
95
96        // Apply quality preset override if specified
97        if let Some(ref preset) = self.quality_preset {
98            quality = quality.apply_preset(Some(preset.as_str()));
99            tracing::info!("Applying quality preset '{}': {}", preset, quality);
100        }
101
102        if book_folder.tracks.len() == 1 {
103            // Single file - just convert
104            self.ffmpeg
105                .convert_single_file(
106                    &book_folder.tracks[0].file_path,
107                    &output_path,
108                    &quality,
109                    use_copy,
110                    self.encoder,
111                )
112                .await
113                .context("Failed to convert audio file")?;
114        } else if use_copy {
115            // Copy mode - can concatenate directly
116            let concat_file = temp_dir.join("concat.txt");
117            let file_refs: Vec<&Path> = book_folder
118                .tracks
119                .iter()
120                .map(|t| t.file_path.as_path())
121                .collect();
122            FFmpeg::create_concat_file(&file_refs, &concat_file)?;
123
124            self.ffmpeg
125                .concat_audio_files(
126                    &concat_file,
127                    &output_path,
128                    &quality,
129                    use_copy,
130                    self.encoder,
131                )
132                .await
133                .context("Failed to concatenate audio files")?;
134        } else if self.enable_parallel_encoding && book_folder.tracks.len() > 1 {
135            // Transcode mode - encode files in parallel with throttling
136            let effective_limit = self.max_concurrent_files.min(book_folder.tracks.len());
137
138            tracing::info!(
139                "Using parallel encoding: {} files with max {} concurrent",
140                book_folder.tracks.len(),
141                effective_limit
142            );
143
144            // Create semaphore to limit concurrent file encodings
145            let semaphore = Arc::new(Semaphore::new(effective_limit));
146
147            // Step 1: Encode all files to AAC/M4A in parallel (with throttling)
148            let mut encoded_files = Vec::new();
149            let mut tasks = Vec::new();
150
151            for (i, track) in book_folder.tracks.iter().enumerate() {
152                let temp_output = temp_dir.join(format!("encoded_{:04}.m4a", i));
153                encoded_files.push(temp_output.clone());
154
155                tracing::info!(
156                    "[{}/{}] Encoding: {} ({:.1} min)",
157                    i + 1,
158                    book_folder.tracks.len(),
159                    track.file_path.file_name().unwrap().to_string_lossy(),
160                    track.quality.duration / 60.0
161                );
162
163                let ffmpeg = self.ffmpeg.clone();
164                let input = track.file_path.clone();
165                let output = temp_output;
166                let quality = quality.clone();
167                let encoder = self.encoder;
168                let sem = Arc::clone(&semaphore);
169
170                // Spawn parallel encoding task with semaphore
171                let task = tokio::spawn(async move {
172                    // Acquire permit before encoding (blocks if limit reached)
173                    let _permit = sem.acquire().await.unwrap();
174
175                    ffmpeg
176                        .convert_single_file(&input, &output, &quality, false, encoder)
177                        .await
178                    // Permit automatically released when _permit drops
179                });
180
181                tasks.push(task);
182            }
183
184            // Wait for all encoding tasks to complete
185            for (i, task) in tasks.into_iter().enumerate() {
186                match task.await {
187                    Ok(Ok(())) => continue,
188                    Ok(Err(e)) => {
189                        return Err(e).context(format!("Track {} encoding failed", i));
190                    }
191                    Err(e) => {
192                        return Err(anyhow::anyhow!("Task {} panicked: {}", i, e));
193                    }
194                }
195            }
196
197            tracing::info!("All {} files encoded, now concatenating...", encoded_files.len());
198
199            // Step 2: Concatenate the encoded files (fast, no re-encoding)
200            let concat_file = temp_dir.join("concat.txt");
201            let file_refs: Vec<&Path> = encoded_files.iter().map(|p| p.as_path()).collect();
202            FFmpeg::create_concat_file(&file_refs, &concat_file)?;
203
204            self.ffmpeg
205                .concat_audio_files(
206                    &concat_file,
207                    &output_path,
208                    &quality,
209                    true, // use copy mode for concatenation
210                    self.encoder,
211                )
212                .await
213                .context("Failed to concatenate encoded files")?;
214        } else {
215            // Serial mode - concatenate and encode in one FFmpeg call (traditional method)
216            tracing::info!("Using serial encoding (parallel encoding disabled in config)");
217
218            let concat_file = temp_dir.join("concat.txt");
219            let file_refs: Vec<&Path> = book_folder
220                .tracks
221                .iter()
222                .map(|t| t.file_path.as_path())
223                .collect();
224            FFmpeg::create_concat_file(&file_refs, &concat_file)?;
225
226            self.ffmpeg
227                .concat_audio_files(
228                    &concat_file,
229                    &output_path,
230                    &quality,
231                    false, // transcode mode
232                    self.encoder,
233                )
234                .await
235                .context("Failed to concatenate audio files")?;
236        }
237
238        tracing::info!("Audio processing complete: {}", output_path.display());
239
240        // Step 3: Generate and inject chapters
241        let chapters = self.generate_chapters(book_folder, chapter_source)?;
242
243        if !chapters.is_empty() {
244            tracing::info!("Injecting {} chapters using MP4Box", chapters.len());
245
246            let chapters_file = temp_dir.join("chapters.txt");
247            write_mp4box_chapters(&chapters, &chapters_file)
248                .context("Failed to write chapter file")?;
249
250            inject_chapters_mp4box(&output_path, &chapters_file)
251                .await
252                .context("Failed to inject chapters")?;
253
254            tracing::info!("✓ Chapter injection complete");
255        }
256
257        // Step 4: Inject metadata
258        let title = book_folder.get_album_title();
259        let artist = book_folder.get_album_artist();
260        let year = book_folder.get_year();
261        let genre = book_folder.get_genre();
262        let comment = book_folder.get_comment();
263        let composer = book_folder.get_composer();
264
265        tracing::info!("Injecting metadata using AtomicParsley");
266        tracing::debug!(
267            "Metadata: title={:?}, artist={:?}",
268            title,
269            artist
270        );
271
272        inject_metadata_atomicparsley(
273            &output_path,
274            title.as_deref(),
275            artist.as_deref(),
276            title.as_deref(), // Use title as album
277            artist.as_deref(), // Album artist
278            year,
279            genre.as_deref(),
280            composer.as_deref(),
281            comment.as_deref(),
282            book_folder.cover_file.as_deref(),
283        )
284        .await
285        .context("Failed to inject metadata")?;
286
287        tracing::info!("✓ Metadata injection complete");
288
289        // Clean up extracted cover file if it was auto-extracted
290        if let Some(cover_path) = &book_folder.cover_file {
291            if cover_path.file_name().and_then(|n| n.to_str()) == Some(".extracted_cover.jpg") {
292                if let Err(e) = std::fs::remove_file(cover_path) {
293                    tracing::debug!("Failed to remove extracted cover file: {}", e);
294                } else {
295                    tracing::debug!("Cleaned up extracted cover file");
296                }
297            }
298        }
299
300        // Clean up temp directory
301        if !self.keep_temp {
302            if let Err(e) = std::fs::remove_dir_all(&temp_dir) {
303                tracing::warn!("Failed to remove temp directory: {}", e);
304            }
305        }
306
307        // Calculate processing time
308        let processing_time = start_time.elapsed().as_secs_f64();
309
310        tracing::info!(
311            "=== Completed: {} in {:.1}s ===",
312            book_folder.name,
313            processing_time
314        );
315
316        // Return success result
317        Ok(result.success(output_path, processing_time, use_copy))
318    }
319
320    /// Generate chapters for the book
321    fn generate_chapters(
322        &self,
323        book_folder: &BookFolder,
324        chapter_source: &str,
325    ) -> Result<Vec<crate::audio::Chapter>> {
326        match chapter_source {
327            "cue" => {
328                // Use CUE file if available
329                if let Some(ref cue_file) = book_folder.cue_file {
330                    tracing::info!("Using CUE file for chapters: {}", cue_file.display());
331                    return parse_cue_file(cue_file);
332                }
333                Ok(Vec::new())
334            }
335            "files" | "auto" => {
336                // Generate chapters from files
337                if book_folder.tracks.len() > 1 {
338                    let files: Vec<&Path> = book_folder
339                        .tracks
340                        .iter()
341                        .map(|t| t.file_path.as_path())
342                        .collect();
343                    let durations: Vec<f64> = book_folder
344                        .tracks
345                        .iter()
346                        .map(|t| t.quality.duration)
347                        .collect();
348
349                    tracing::info!(
350                        "Generating {} chapters from files",
351                        book_folder.tracks.len()
352                    );
353                    Ok(generate_chapters_from_files(&files, &durations))
354                } else {
355                    // Single file - check for CUE
356                    if let Some(ref cue_file) = book_folder.cue_file {
357                        tracing::info!("Using CUE file for single-file book");
358                        parse_cue_file(cue_file)
359                    } else {
360                        Ok(Vec::new())
361                    }
362                }
363            }
364            "none" => Ok(Vec::new()),
365            _ => {
366                tracing::warn!("Unknown chapter source: {}, using auto", chapter_source);
367                self.generate_chapters(book_folder, "auto")
368            }
369        }
370    }
371
372    /// Create temporary directory for processing
373    fn create_temp_dir(&self, book_name: &str) -> Result<PathBuf> {
374        let temp_base = std::env::temp_dir();
375        let sanitized_name = sanitize_filename::sanitize(book_name);
376        let temp_dir = temp_base.join(format!("audiobook-forge-{}", sanitized_name));
377
378        if temp_dir.exists() {
379            std::fs::remove_dir_all(&temp_dir).ok();
380        }
381
382        std::fs::create_dir_all(&temp_dir).context("Failed to create temp directory")?;
383
384        Ok(temp_dir)
385    }
386}
387
388impl Default for Processor {
389    fn default() -> Self {
390        Self::new().expect("Failed to create processor")
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    #[test]
399    fn test_processor_creation() {
400        let processor = Processor::new();
401        assert!(processor.is_ok());
402    }
403
404    #[test]
405    fn test_processor_with_options() {
406        let processor = Processor::with_options(true, AacEncoder::AppleSilicon, true, 8, None).unwrap();
407        assert!(processor.keep_temp);
408        assert_eq!(processor.encoder, AacEncoder::AppleSilicon);
409        assert_eq!(processor.max_concurrent_files, 8);
410        assert_eq!(processor.quality_preset, None);
411
412        let processor_with_preset = Processor::with_options(false, AacEncoder::Native, true, 4, Some("high".to_string())).unwrap();
413        assert_eq!(processor_with_preset.quality_preset, Some("high".to_string()));
414    }
415
416    #[test]
417    fn test_create_temp_dir() {
418        let processor = Processor::new().unwrap();
419        let temp_dir = processor.create_temp_dir("Test Book").unwrap();
420
421        assert!(temp_dir.exists());
422        assert!(temp_dir.to_string_lossy().contains("audiobook-forge"));
423
424        // Clean up
425        std::fs::remove_dir_all(temp_dir).ok();
426    }
427}