rig-core 0.44.0

use epub::doc::EpubDoc;

use std::fs::File;
use std::io::BufReader;
use std::marker::PhantomData;
use std::path::PathBuf;

use super::RawTextProcessor;
use super::errors::EpubLoaderError;
use super::text_processors::TextProcessor;

loadable_trait!(
    Loadable,
    EpubLoaderError,
    EpubDoc<BufReader<File>>,
    load,
    load_with_path
);

impl Loadable for PathBuf {
    fn load(self) -> Result<EpubDoc<BufReader<File>>, EpubLoaderError> {
        EpubDoc::new(self).map_err(EpubLoaderError::EpubError)
    }

    fn load_with_path(self) -> Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError> {
        let contents = EpubDoc::new(&self).map_err(EpubLoaderError::EpubError);
        Ok((self, contents?))
    }
}

/// Iterator pipeline for synchronous EPUB loading and chapter processing.
/// `P` selects text processing and defaults to [`RawTextProcessor`].
/// Loading and processing errors are yielded unless filtered; unavailable chapter
/// text is skipped.
pub struct EpubFileLoader<'a, T, P = RawTextProcessor> {
    iterator: Box<dyn Iterator<Item = T> + 'a>,
    _processor: PhantomData<P>,
}

type EpubLoaded = Result<(PathBuf, EpubDoc<BufReader<File>>), EpubLoaderError>;

impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P> {
    /// Loads each EPUB during iteration, yielding document-loading errors.
    ///
    /// ```no_run
    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    ///
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load().into_iter();
    /// for result in content {
    ///     match result {
    ///         Ok(doc) => println!("{doc:?}"),
    ///         Err(e) => eprintln!("Error reading epub: {e}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn load(self) -> EpubFileLoader<'a, Result<EpubDoc<BufReader<File>>, EpubLoaderError>, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(Loadable::load)),
            _processor: PhantomData,
        }
    }

    /// Loads each EPUB paired with its source path, yielding loading errors.
    ///
    /// ```no_run
    /// use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    ///
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.load_with_path().into_iter();
    /// for result in content {
    ///     match result {
    ///         Ok((path, doc)) => println!("{path:?} {doc:?}"),
    ///         Err(e) => eprintln!("Error reading epub: {e}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn load_with_path(self) -> EpubFileLoader<'a, EpubLoaded, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(Loadable::load_with_path)),
            _processor: PhantomData,
        }
    }
}

impl<'a, P> EpubFileLoader<'a, Result<PathBuf, EpubLoaderError>, P>
where
    P: TextProcessor,
{
    /// Concatenates processed chapters without separators, yielding loading or
    /// processing errors. Unavailable chapter text is skipped.
    ///
    /// ```no_run
    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read().into_iter();
    /// for result in content {
    ///     match result {
    ///         Ok(content) => println!("{content}"),
    ///         Err(e) => eprintln!("Error reading epub: {e}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn read(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(|res| {
                let doc = res.load().map(EpubChapterIterator::<P>::from)?;

                Ok(doc
                    .into_iter()
                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
                    .into_iter()
                    .collect::<String>())
            })),
            _processor: PhantomData,
        }
    }

    /// Pairs each path with concatenated processed chapters, yielding loading
    /// or processing errors. Unavailable chapter text is skipped.
    ///
    /// ```no_run
    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?.read_with_path().into_iter();
    /// for result in content {
    ///     match result {
    ///         Ok((path, content)) => println!("{path:?} {content}"),
    ///         Err(e) => eprintln!("Error reading epub: {e}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn read_with_path(
        self,
    ) -> EpubFileLoader<'a, Result<(PathBuf, String), EpubLoaderError>, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(|res| {
                let (path, doc) = res.load_with_path()?;

                let content = EpubChapterIterator::<P>::from(doc)
                    .collect::<Result<Vec<String>, EpubLoaderError>>()?
                    .into_iter()
                    .collect::<String>();
                Ok((path, content))
            })),
            _processor: PhantomData,
        }
    }
}

impl<'a, P> EpubFileLoader<'a, EpubDoc<BufReader<File>>, P>
where
    P: TextProcessor + 'a,
{
    /// Yields processed chapter results across documents in input order,
    /// skipping unavailable chapter text.
    ///
    /// ```no_run
    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
    ///     .load()
    ///     .ignore_errors()
    ///     .by_chapter()
    ///     .into_iter();
    /// for result in content {
    ///     match result {
    ///         Ok(chapter) => println!("{chapter}"),
    ///         Err(e) => eprintln!("Error reading chapter: {e}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn by_chapter(self) -> EpubFileLoader<'a, Result<String, EpubLoaderError>, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.flat_map(EpubChapterIterator::<P>::from)),
            _processor: PhantomData,
        }
    }
}

type ByChapter = (PathBuf, Vec<(usize, Result<String, EpubLoaderError>)>);
impl<'a, P: TextProcessor> EpubFileLoader<'a, (PathBuf, EpubDoc<BufReader<File>>), P> {
    /// Pairs each path with zero-based indices and chapter-processing results.
    /// Indices count available chapter texts, not skipped source entries.
    ///
    /// ```no_run
    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
    ///     .load_with_path()
    ///     .ignore_errors()
    ///     .by_chapter()
    ///     .ignore_errors()
    ///     .into_iter();
    ///
    /// for result in content {
    ///     println!("{result:?}");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn by_chapter(self) -> EpubFileLoader<'a, ByChapter, P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(|doc| {
                let (path, doc) = doc;

                (
                    path,
                    EpubChapterIterator::<P>::from(doc)
                        .enumerate()
                        .collect::<Vec<_>>(),
                )
            })),
            _processor: PhantomData,
        }
    }
}

impl<'a, P> EpubFileLoader<'a, ByChapter, P>
where
    P: TextProcessor,
{
    /// Drops failed chapter results, retaining paths and existing chapter indices.
    /// Documents with no successful chapters remain.
    ///
    /// ```no_run
    /// # use rig_core::loaders::{EpubFileLoader, RawTextProcessor};
    /// # fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let content = EpubFileLoader::<_, RawTextProcessor>::with_glob("tests/data/*.epub")?
    ///     .load_with_path()
    ///     .ignore_errors()
    ///     .by_chapter()
    ///     .ignore_errors();
    /// for (_path, chapters) in content {
    ///     println!("{}", chapters.len())
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn ignore_errors(self) -> EpubFileLoader<'a, (PathBuf, Vec<(usize, String)>), P> {
        EpubFileLoader {
            iterator: Box::new(self.iterator.map(|(path, chapters)| {
                let chapters = chapters
                    .into_iter()
                    .filter_map(|(idx, res)| res.ok().map(|content| (idx, content)))
                    .collect::<Vec<_>>();
                (path, chapters)
            })),
            _processor: PhantomData,
        }
    }
}

loader_scaffold!(EpubFileLoader, EpubLoaderError, dir: all_entries, extra: P);

struct EpubChapterIterator<P> {
    epub: EpubDoc<BufReader<File>>,
    finished: bool,
    _processor: PhantomData<P>,
}

impl<P> From<EpubDoc<BufReader<File>>> for EpubChapterIterator<P> {
    fn from(epub: EpubDoc<BufReader<File>>) -> Self {
        Self::new(epub)
    }
}

impl<P> EpubChapterIterator<P> {
    fn new(epub: EpubDoc<BufReader<File>>) -> Self {
        Self {
            epub,
            finished: false,
            _processor: PhantomData,
        }
    }
}

impl<P> Iterator for EpubChapterIterator<P>
where
    P: TextProcessor,
{
    type Item = Result<String, EpubLoaderError>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            return None;
        }

        // Unavailable chapter text is skipped rather than yielded as an error.
        while !self.finished {
            let chapter = self.epub.get_current_str();

            if !self.epub.go_next() {
                self.finished = true;
            }

            if let Some((text, _)) = chapter {
                return Some(
                    P::process(&text)
                        .map_err(|err| EpubLoaderError::TextProcessorError(Box::new(err))),
                );
            }
        }

        None
    }
}

#[cfg(test)]
mod tests;