oxidize-pdf 5.0.0

//! Worker pool for parallel batch processing

use crate::batch::{BatchJob, BatchProgress, JobResult};
use crate::error::PdfError;
use crate::operations::existing_document::{
    extract_pdf_pages, merge_pdfs, split_pdf, ExistingDocumentMergeInput, ExistingDocumentPolicy,
};
use crate::operations::OperationError;
use crate::operations::PageRange;
use crate::parser::PdfReader;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

/// Options for worker pool
#[derive(Debug, Clone)]
pub struct WorkerOptions {
    /// Number of worker threads
    pub num_workers: usize,
    /// Memory limit per worker (bytes)
    pub memory_limit: usize,
    /// Timeout for individual jobs
    pub job_timeout: Option<Duration>,
}

/// Message sent to workers
enum WorkerMessage {
    Job(usize, BatchJob),
    Shutdown,
}

/// Worker pool for parallel processing
pub struct WorkerPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<WorkerMessage>,
}

impl WorkerPool {
    /// Create a new worker pool
    pub fn new(options: WorkerOptions) -> Self {
        let (sender, receiver) = mpsc::channel();
        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(options.num_workers);

        for id in 0..options.num_workers {
            workers.push(Worker::new(
                id,
                Arc::clone(&receiver),
                options.memory_limit,
                options.job_timeout,
            ));
        }

        Self { workers, sender }
    }

    /// Process a batch of jobs
    pub fn process_jobs(
        self,
        jobs: Vec<BatchJob>,
        progress: Arc<BatchProgress>,
        cancelled: Arc<AtomicBool>,
        stop_on_error: bool,
    ) -> Vec<JobResult> {
        let num_jobs = jobs.len();
        let (result_sender, result_receiver) = mpsc::channel();

        // Spawn result collector thread
        let results = vec![None; num_jobs];
        let results_handle = {
            let mut results = results;
            thread::spawn(move || {
                for (idx, result) in result_receiver {
                    results[idx] = Some(result);
                }
                results
            })
        };

        // Send jobs to workers
        for (idx, job) in jobs.into_iter().enumerate() {
            if cancelled.load(Ordering::SeqCst) {
                let _ = result_sender.send((
                    idx,
                    JobResult::Cancelled {
                        job_name: job.display_name(),
                    },
                ));
                continue;
            }

            let job_name = job.display_name();
            let progress_clone = Arc::clone(&progress);
            let result_sender_clone = result_sender.clone();
            let cancelled_clone = Arc::clone(&cancelled);

            // Wrap job with progress tracking
            let wrapped_job = match job {
                BatchJob::Custom { name, operation } => BatchJob::Custom {
                    name,
                    operation: Box::new(move || {
                        progress_clone.start_job();
                        let start = Instant::now();

                        let result = if cancelled_clone.load(Ordering::SeqCst) {
                            Err(PdfError::OperationCancelled)
                        } else {
                            operation()
                        };

                        let duration = start.elapsed();

                        match result {
                            Ok(()) => {
                                progress_clone.complete_job();
                                let _ = result_sender_clone.send((
                                    idx,
                                    JobResult::Success {
                                        job_name: job_name.clone(),
                                        duration,
                                        output_files: vec![],
                                    },
                                ));
                            }
                            Err(ref e) => {
                                progress_clone.fail_job();
                                let _ = result_sender_clone.send((
                                    idx,
                                    JobResult::Failed {
                                        job_name: job_name.clone(),
                                        duration,
                                        error: e.to_string(),
                                    },
                                ));
                            }
                        }

                        result
                    }),
                },
                _ => {
                    // Handle other job types
                    let progress_clone2 = Arc::clone(&progress);
                    let result_sender_clone2 = result_sender.clone();

                    BatchJob::Custom {
                        name: job_name.clone(),
                        operation: Box::new(move || {
                            progress_clone2.start_job();
                            let start = Instant::now();

                            let result = execute_job(job);
                            let duration = start.elapsed();

                            match &result {
                                Ok(output_files) => {
                                    progress_clone2.complete_job();
                                    let _ = result_sender_clone2.send((
                                        idx,
                                        JobResult::Success {
                                            job_name: job_name.clone(),
                                            duration,
                                            output_files: output_files.clone(),
                                        },
                                    ));
                                }
                                Err(e) => {
                                    progress_clone2.fail_job();
                                    let _ = result_sender_clone2.send((
                                        idx,
                                        JobResult::Failed {
                                            job_name: job_name.clone(),
                                            duration,
                                            error: e.to_string(),
                                        },
                                    ));

                                    if stop_on_error {
                                        cancelled_clone.store(true, Ordering::SeqCst);
                                    }
                                }
                            }

                            result.map(|_| ())
                        }),
                    }
                }
            };

            if self
                .sender
                .send(WorkerMessage::Job(idx, wrapped_job))
                .is_err()
            {
                break;
            }
        }

        // Drop the original sender to close the channel
        drop(result_sender);
        drop(self.sender);

        // Wait for workers to finish
        for worker in self.workers {
            worker.join();
        }

        // Collect results
        let results = results_handle.join().unwrap_or_else(|_| {
            tracing::debug!("Result collection thread panicked");
            Vec::new()
        });
        results.into_iter().flatten().collect()
    }

    /// Shutdown the worker pool
    pub fn shutdown(self) {
        for _ in &self.workers {
            let _ = self.sender.send(WorkerMessage::Shutdown);
        }

        for worker in self.workers {
            worker.join();
        }
    }
}

/// Worker thread
struct Worker {
    #[allow(dead_code)]
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    /// Create a new worker
    fn new(
        id: usize,
        receiver: Arc<Mutex<mpsc::Receiver<WorkerMessage>>>,
        _memory_limit: usize,
        job_timeout: Option<Duration>,
    ) -> Self {
        let thread = thread::spawn(move || {
            loop {
                let message = {
                    let receiver = match receiver.lock() {
                        Ok(r) => r,
                        Err(_) => {
                            tracing::debug!("Worker {} receiver lock poisoned", id);
                            break;
                        }
                    };
                    receiver.recv()
                };

                match message {
                    Ok(WorkerMessage::Job(_idx, job)) => {
                        // Execute job with optional timeout
                        if let Some(_timeout) = job_timeout {
                            // In a real implementation, we'd use a timeout mechanism
                            // For now, just execute normally
                            if let BatchJob::Custom { operation, .. } = job {
                                let _ = operation();
                            }
                        } else if let BatchJob::Custom { operation, .. } = job {
                            let _ = operation();
                        }
                    }
                    Ok(WorkerMessage::Shutdown) => break,
                    Err(_) => break,
                }
            }
        });

        Self {
            id,
            thread: Some(thread),
        }
    }

    /// Wait for the worker to finish
    fn join(mut self) {
        if let Some(thread) = self.thread.take() {
            let _ = thread.join();
        }
    }
}

/// Execute a non-custom job
fn execute_job(job: BatchJob) -> std::result::Result<Vec<PathBuf>, PdfError> {
    match job {
        BatchJob::Split {
            input,
            output_pattern,
            pages_per_file,
        } => {
            if pages_per_file == 0 {
                return Err(PdfError::InvalidStructure(
                    "pages_per_file must be greater than zero".to_string(),
                ));
            }
            let page_count = PdfReader::open(&input)?.page_count()? as usize;
            if page_count == 0 {
                return Err(PdfError::InvalidStructure(
                    "cannot split a PDF with no pages".to_string(),
                ));
            }
            let mut ranges = Vec::new();
            let mut outputs = Vec::new();
            for (index, start) in (0..page_count).step_by(pages_per_file).enumerate() {
                let end = (start + pages_per_file - 1).min(page_count - 1);
                ranges.push(PageRange::Range(start, end));
                outputs.push(PathBuf::from(
                    output_pattern
                        .replace("%d", &(index + 1).to_string())
                        .replace("{}", &format!("{}-{}", start + 1, end + 1))
                        .replace("{n}", &(index + 1).to_string())
                        .replace("{start}", &(start + 1).to_string())
                        .replace("{end}", &(end + 1).to_string()),
                ));
            }
            split_pdf(
                &input,
                &ranges,
                &outputs,
                ExistingDocumentPolicy::preserve_base(),
            )
            .map_err(operation_error_to_pdf_error)?;
            Ok(outputs)
        }

        BatchJob::Merge { inputs, output } => {
            let merge_inputs: Vec<_> = inputs
                .into_iter()
                .map(ExistingDocumentMergeInput::new)
                .collect();
            merge_pdfs(
                &merge_inputs,
                &output,
                ExistingDocumentPolicy::preserve_base(),
            )
            .map_err(operation_error_to_pdf_error)?;
            Ok(vec![output])
        }

        BatchJob::Rotate {
            input,
            output,
            rotation: _,
            pages: _,
        } => {
            // Rotate not implemented in current API, just copy
            std::fs::copy(&input, &output)?;
            Ok(vec![output])
        }

        BatchJob::Extract {
            input,
            output,
            pages,
        } => {
            extract_pdf_pages(
                &input,
                &output,
                &pages,
                ExistingDocumentPolicy::preserve_base(),
            )
            .map_err(operation_error_to_pdf_error)?;
            Ok(vec![output])
        }

        BatchJob::Compress {
            input,
            output,
            quality: _,
        } => {
            // Compression not implemented yet, just copy
            std::fs::copy(&input, &output)?;
            Ok(vec![output])
        }

        BatchJob::Custom { .. } => {
            unreachable!("Custom jobs should be handled separately")
        }
    }
}

fn operation_error_to_pdf_error(error: OperationError) -> PdfError {
    match error {
        OperationError::Io(error) => PdfError::Io(error),
        OperationError::PdfError(error) => error,
        OperationError::ParseError(message) => PdfError::ParseError(message),
        OperationError::InvalidPath { reason }
        | OperationError::InvalidPageRange(reason)
        | OperationError::ResourceConflict(reason)
        | OperationError::ProcessingError(reason) => PdfError::InvalidOperation(reason),
        OperationError::PageIndexOutOfBounds(page, _) => u32::try_from(page)
            .map(PdfError::InvalidPageNumber)
            .unwrap_or_else(|_| {
                PdfError::InvalidOperation(format!("page index {page} is too large"))
            }),
        OperationError::NoPagesToProcess => {
            PdfError::InvalidOperation("no pages to process".to_string())
        }
        OperationError::InvalidRotation(rotation) => {
            PdfError::InvalidOperation(format!("invalid rotation angle {rotation}"))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Document, Page};

    #[test]
    fn batch_extraction_uses_the_fail_closed_preserving_engine() {
        let directory = tempfile::tempdir().unwrap();
        let input = directory.path().join("source.pdf");
        let output = directory.path().join("extracted.pdf");
        let mut document = Document::new();
        document.add_page(Page::a4());
        document.add_page(Page::a4());
        document.save(&input).unwrap();
        let source = std::fs::read(&input).unwrap();

        execute_job(BatchJob::Extract {
            input,
            output: output.clone(),
            pages: vec![0],
        })
        .unwrap();

        assert!(std::fs::read(output).unwrap().starts_with(&source));
    }

    #[test]
    fn operation_errors_keep_their_public_category() {
        let io = operation_error_to_pdf_error(OperationError::Io(std::io::Error::new(
            std::io::ErrorKind::PermissionDenied,
            "denied",
        )));
        let parse = operation_error_to_pdf_error(OperationError::ParseError("broken".into()));
        let permission = operation_error_to_pdf_error(OperationError::PdfError(
            PdfError::PermissionDenied("locked".into()),
        ));

        assert!(matches!(io, PdfError::Io(_)));
        assert!(matches!(parse, PdfError::ParseError(_)));
        assert!(matches!(permission, PdfError::PermissionDenied(_)));
    }

    #[test]
    fn split_rejects_a_zero_page_pdf() {
        let directory = tempfile::tempdir().unwrap();
        let input = directory.path().join("empty.pdf");
        let bytes = b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [] /Count 0 >>\nendobj\nxref\n0 3\n0000000000 65535 f \n0000000009 00000 n \n0000000058 00000 n \ntrailer\n<< /Size 3 /Root 1 0 R >>\nstartxref\n110\n%%EOF\n";
        std::fs::write(&input, bytes).unwrap();

        let error = execute_job(BatchJob::Split {
            input,
            output_pattern: directory.path().join("page_%d.pdf").display().to_string(),
            pages_per_file: 1,
        })
        .unwrap_err();

        assert!(error.to_string().contains("no pages"));
    }

    #[test]
    fn test_worker_pool_creation() {
        let options = WorkerOptions {
            num_workers: 2,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        assert_eq!(pool.workers.len(), 2);

        pool.shutdown();
    }

    #[test]
    fn test_worker_pool_empty_jobs() {
        let options = WorkerOptions {
            num_workers: 2,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let results = pool.process_jobs(vec![], progress, cancelled, false);
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_worker_pool_custom_jobs() {
        let options = WorkerOptions {
            num_workers: 2,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let jobs = vec![
            BatchJob::Custom {
                name: "Test Job 1".to_string(),
                operation: Box::new(|| Ok(())),
            },
            BatchJob::Custom {
                name: "Test Job 2".to_string(),
                operation: Box::new(|| Ok(())),
            },
        ];

        progress.add_job();
        progress.add_job();

        let results = pool.process_jobs(jobs, progress, cancelled, false);

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r.is_success()));
    }

    #[test]
    fn test_worker_pool_with_failures() {
        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let jobs = vec![
            BatchJob::Custom {
                name: "Success Job".to_string(),
                operation: Box::new(|| Ok(())),
            },
            BatchJob::Custom {
                name: "Failing Job".to_string(),
                operation: Box::new(|| Err(PdfError::InvalidStructure("Test error".to_string()))),
            },
        ];

        progress.add_job();
        progress.add_job();

        let results = pool.process_jobs(jobs, progress, cancelled, false);

        assert_eq!(results.len(), 2);
        assert!(results[0].is_success());
        assert!(results[1].is_failed());
    }

    #[test]
    fn test_worker_pool_shutdown_with_active_jobs() {
        // Test graceful shutdown while jobs are running
        let options = WorkerOptions {
            num_workers: 2,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        // Jobs that take time to complete
        let jobs = vec![BatchJob::Custom {
            name: "Long Running Job".to_string(),
            operation: Box::new(|| {
                std::thread::sleep(std::time::Duration::from_millis(50));
                Ok(())
            }),
        }];

        progress.add_job();

        // Process jobs and shutdown - should complete gracefully
        let results = pool.process_jobs(jobs, progress, cancelled, false);
        assert_eq!(results.len(), 1);
        assert!(results[0].is_success());
    }

    #[test]
    fn test_worker_pool_job_panic_recovery() {
        // Test that worker pool recovers from panicking jobs
        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let jobs = vec![
            BatchJob::Custom {
                name: "Panicking Job".to_string(),
                operation: Box::new(|| {
                    // Convert panic to error for testing
                    Err(PdfError::InvalidStructure("Simulated panic".to_string()))
                }),
            },
            BatchJob::Custom {
                name: "Normal Job".to_string(),
                operation: Box::new(|| Ok(())),
            },
        ];

        progress.add_job();
        progress.add_job();

        let results = pool.process_jobs(jobs, progress, cancelled, false);

        assert_eq!(results.len(), 2);
        assert!(results[0].is_failed());
        assert!(results[1].is_success());
    }

    #[test]
    fn test_worker_pool_memory_pressure() {
        // Test behavior under memory constraints
        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024, // Very low limit
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let jobs = vec![BatchJob::Custom {
            name: "Memory Test Job".to_string(),
            operation: Box::new(|| {
                // Simulate work that could use memory
                let _data = vec![0u8; 512]; // Small allocation
                Ok(())
            }),
        }];

        progress.add_job();

        let results = pool.process_jobs(jobs, progress, cancelled, false);
        assert_eq!(results.len(), 1);
        // Should succeed with small allocation
        assert!(results[0].is_success());
    }

    #[test]
    fn test_worker_pool_cancellation_during_processing() {
        // Test cancellation while jobs are being processed
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::sync::Arc;

        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let cancelled_clone = Arc::clone(&cancelled);
        let jobs = vec![
            BatchJob::Custom {
                name: "Job Before Cancel".to_string(),
                operation: Box::new(|| Ok(())),
            },
            BatchJob::Custom {
                name: "Job After Cancel".to_string(),
                operation: Box::new(move || {
                    // This should be cancelled
                    if cancelled_clone.load(Ordering::SeqCst) {
                        Err(PdfError::InvalidStructure("Cancelled".to_string()))
                    } else {
                        Ok(())
                    }
                }),
            },
        ];

        progress.add_job();
        progress.add_job();

        // Cancel after starting
        cancelled.store(true, Ordering::SeqCst);

        let results = pool.process_jobs(jobs, progress, cancelled, false);
        assert_eq!(results.len(), 2);
        // Some jobs might be cancelled
    }

    #[test]
    fn test_worker_pool_timeout_handling() {
        // Test job timeout handling
        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024 * 1024,
            job_timeout: Some(std::time::Duration::from_millis(10)), // Very short timeout
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(false));

        let jobs = vec![
            BatchJob::Custom {
                name: "Quick Job".to_string(),
                operation: Box::new(|| Ok(())), // Should complete quickly
            },
            BatchJob::Custom {
                name: "Slow Job".to_string(),
                operation: Box::new(|| {
                    // Simulate timeout scenario with error
                    std::thread::sleep(std::time::Duration::from_millis(5));
                    Ok(())
                }),
            },
        ];

        progress.add_job();
        progress.add_job();

        let results = pool.process_jobs(jobs, Arc::clone(&progress), cancelled, false);

        assert_eq!(results.len(), 2);
        assert_eq!(results.iter().filter(|r| r.is_success()).count(), 2);
        assert_eq!(results.iter().filter(|r| r.is_failed()).count(), 0);

        let info = progress.get_info();
        assert_eq!(info.completed_jobs, 2);
        assert_eq!(info.failed_jobs, 0);
    }

    #[test]
    fn test_worker_pool_cancellation() {
        let options = WorkerOptions {
            num_workers: 1,
            memory_limit: 1024 * 1024,
            job_timeout: None,
        };

        let pool = WorkerPool::new(options);
        let progress = Arc::new(BatchProgress::new());
        let cancelled = Arc::new(AtomicBool::new(true)); // Pre-cancelled

        let jobs = vec![BatchJob::Custom {
            name: "Should be cancelled".to_string(),
            operation: Box::new(|| Ok(())),
        }];

        progress.add_job();

        let results = pool.process_jobs(jobs, progress, cancelled, false);

        assert_eq!(results.len(), 1);
        assert!(results[0].is_cancelled());
    }
}