text-extract/src/lib.rs

//! `text.extract` module — extract plain text from PDF or DOCX
//! documents.
//!
//! # Capability
//!
//! - **Input:** `document` — `bytes` payload, MIME type
//!   `application/pdf` or
//!   `application/vnd.openxmlformats-officedocument.wordprocessingml.document`.
//! - **Output:** `extracted` — `json` payload conforming to the
//!   `ExtractedDocument` schema below.
//!
//! # Output schema
//!
//! ```json
//! {
//!   "engine": "pdf-extract",
//!   "engine_version": "0.7",
//!   "pages": [
//!     { "number": 1, "text": "..." }
//!   ]
//! }
//! ```
//!
//! Per-page `confidence` and `bbox` fields are reserved for a
//! future service-mode variant (Phase 0.5+, host-services-backed)
//! that uses pdfium or mupdf and exposes layout information.
//! v0.1.0 does not populate them; the schema accepts them as
//! optional so that the v0.2.0 service variant can add them
//! without an interface break.
//!
//! # Stability
//!
//! Module version `0.1.0` ships as `alpha` in the store index:
//! installable, but the PDF extraction quality is layout-naive
//! (single-column ok; multi-column or table-heavy documents may
//! produce reordered text). Bumping to `1.0.0` requires either
//! the service-mode variant or a confirmed-good extraction
//! quality on a representative document set.

#![allow(clippy::result_large_err)]

use chain_module_sdk::prelude::*;
use serde::Serialize;

const PDF_MIME: &str = "application/pdf";
const PDF_MIME_ALT: &str = "application/x-pdf";
const DOCX_MIME: &str = "application/vnd.openxmlformats-officedocument.wordprocessingml.document";

const ENGINE_PDF: &str = "pdf-extract";
const ENGINE_PDF_VERSION: &str = "0.7";
const ENGINE_DOCX: &str = "zip+quick-xml";
const ENGINE_DOCX_VERSION: &str = "0.1";

#[derive(Debug, Serialize)]
struct ExtractedDocument {
    engine: String,
    engine_version: String,
    pages: Vec<Page>,
}

#[derive(Debug, Serialize)]
struct Page {
    number: u32,
    text: String,
    /// Reserved for a future service-mode variant. Populated only
    /// when the underlying engine reports a per-page confidence.
    #[serde(skip_serializing_if = "Option::is_none")]
    confidence: Option<f64>,
    /// Reserved for a future service-mode variant. Populated only
    /// when the underlying engine reports per-page bounding boxes.
    #[serde(skip_serializing_if = "Option::is_none")]
    bbox: Option<BoundingBox>,
}

#[derive(Debug, Serialize)]
struct BoundingBox {
    x: f64,
    y: f64,
    width: f64,
    height: f64,
}

#[chain_module]
pub fn invoke(_ctx: Context, inputs: Inputs) -> Result<Outputs, ModuleError> {
    let bytes = inputs.require_bytes("document")?;
    let extracted = match bytes.mime_type.as_str() {
        PDF_MIME | PDF_MIME_ALT => extract_pdf(&bytes.data)?,
        DOCX_MIME => extract_docx(&bytes.data)?,
        other => {
            return Err(ModuleError::invalid_input(format!(
                "unsupported MIME type: {other} (supported: {PDF_MIME}, {DOCX_MIME})",
            )));
        }
    };
    Outputs::new().with_json("extracted", &extracted)
}

fn extract_pdf(data: &[u8]) -> Result<ExtractedDocument, ModuleError> {
    let text = pdf_extract::extract_text_from_mem(data)
        .map_err(|e| ModuleError::invalid_input(format!("PDF parsing failed: {e}")))?;
    Ok(ExtractedDocument {
        engine: ENGINE_PDF.to_string(),
        engine_version: ENGINE_PDF_VERSION.to_string(),
        pages: split_pdf_text_into_pages(&text),
    })
}

/// pdf-extract's `extract_text_from_mem` returns a single string
/// with form-feed (`\x0c`) characters between pages. Split on
/// that boundary so the consumer keeps the page structure.
fn split_pdf_text_into_pages(combined: &str) -> Vec<Page> {
    let trimmed = combined.trim_matches('\x0c');
    if trimmed.is_empty() {
        return Vec::new();
    }
    trimmed
        .split('\x0c')
        .enumerate()
        .map(|(i, page_text)| Page {
            number: u32::try_from(i + 1).unwrap_or(u32::MAX),
            text: page_text.trim().to_string(),
            confidence: None,
            bbox: None,
        })
        .collect()
}

fn extract_docx(data: &[u8]) -> Result<ExtractedDocument, ModuleError> {
    let cursor = std::io::Cursor::new(data);
    let mut archive = zip::ZipArchive::new(cursor)
        .map_err(|e| ModuleError::invalid_input(format!("DOCX is not a valid ZIP: {e}")))?;
    let mut document_xml = archive.by_name("word/document.xml").map_err(|_| {
        ModuleError::invalid_input(
            "DOCX archive does not contain word/document.xml — not a valid DOCX",
        )
    })?;
    let mut buf = Vec::with_capacity(usize::try_from(document_xml.size()).unwrap_or(0));
    std::io::Read::read_to_end(&mut document_xml, &mut buf)
        .map_err(|e| ModuleError::internal(format!("failed reading word/document.xml: {e}")))?;
    let text = extract_docx_text_from_xml(&buf)?;
    Ok(ExtractedDocument {
        engine: ENGINE_DOCX.to_string(),
        engine_version: ENGINE_DOCX_VERSION.to_string(),
        pages: vec![Page {
            number: 1,
            text,
            confidence: None,
            bbox: None,
        }],
    })
}

/// Extract visible text from a DOCX `document.xml` body.
///
/// DOCX paragraphs nest text runs as `<w:r><w:t>...</w:t></w:r>`.
/// We collect text from `<w:t>` elements, insert a newline at
/// each paragraph boundary, and skip everything else.
fn extract_docx_text_from_xml(xml: &[u8]) -> Result<String, ModuleError> {
    use quick_xml::Reader;
    use quick_xml::events::Event;

    let mut reader = Reader::from_reader(xml);
    reader.config_mut().trim_text(false);

    let mut out = String::new();
    let mut in_text = false;

    loop {
        match reader.read_event() {
            Ok(Event::Start(e)) => {
                if local_name_matches(e.name().as_ref(), b"t") {
                    in_text = true;
                }
            }
            Ok(Event::End(e)) => {
                if local_name_matches(e.name().as_ref(), b"t") {
                    in_text = false;
                } else if local_name_matches(e.name().as_ref(), b"p") {
                    out.push('\n');
                }
            }
            Ok(Event::Text(e)) if in_text => {
                let unescaped = e
                    .unescape()
                    .map_err(|err| ModuleError::invalid_input(format!("DOCX XML decode: {err}")))?;
                out.push_str(&unescaped);
            }
            Ok(Event::Eof) => break,
            Err(err) => {
                return Err(ModuleError::invalid_input(format!(
                    "DOCX XML parse error at position {}: {err}",
                    reader.error_position()
                )));
            }
            _ => {}
        }
    }

    Ok(out.trim_end().to_string())
}

fn local_name_matches(qname: &[u8], local: &[u8]) -> bool {
    match qname.iter().rposition(|b| *b == b':') {
        Some(idx) => &qname[idx + 1..] == local,
        None => qname == local,
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::panic)]
    use super::*;

    #[test]
    fn split_pdf_text_into_pages_handles_form_feed() {
        let combined = "Page one text\x0cPage two text\x0c";
        let pages = split_pdf_text_into_pages(combined);
        assert_eq!(pages.len(), 2);
        assert_eq!(pages[0].number, 1);
        assert_eq!(pages[0].text, "Page one text");
        assert_eq!(pages[1].number, 2);
        assert_eq!(pages[1].text, "Page two text");
    }

    #[test]
    fn split_pdf_text_into_pages_returns_empty_on_empty_input() {
        assert!(split_pdf_text_into_pages("").is_empty());
        assert!(split_pdf_text_into_pages("\x0c\x0c").is_empty());
    }

    #[test]
    fn local_name_matches_qualified_and_unqualified() {
        assert!(local_name_matches(b"w:t", b"t"));
        assert!(local_name_matches(b"t", b"t"));
        assert!(local_name_matches(b"namespace:p", b"p"));
        assert!(!local_name_matches(b"w:t", b"p"));
        assert!(!local_name_matches(b"w:tbl", b"t"));
    }

    #[test]
    fn extract_docx_text_handles_paragraphs() {
        let xml = br#"<?xml version="1.0"?>
<w:document xmlns:w="urn:test">
  <w:body>
    <w:p><w:r><w:t>Hello</w:t></w:r><w:r><w:t> World</w:t></w:r></w:p>
    <w:p><w:r><w:t>Second paragraph</w:t></w:r></w:p>
  </w:body>
</w:document>"#;
        let text = extract_docx_text_from_xml(xml).unwrap();
        assert_eq!(text, "Hello World\nSecond paragraph");
    }

    #[test]
    fn extract_docx_text_skips_non_t_elements() {
        let xml = br#"<w:document xmlns:w="urn:test">
  <w:p>
    <w:r><w:rPr><w:b/></w:rPr><w:t>Bold</w:t></w:r>
    <w:r><w:t> normal</w:t></w:r>
  </w:p>
</w:document>"#;
        let text = extract_docx_text_from_xml(xml).unwrap();
        assert_eq!(text, "Bold normal");
    }

    #[test]
    fn extract_docx_text_preserves_whitespace_in_t() {
        // <w:t xml:space="preserve"> usually carries whitespace.
        // Our extractor passes text through untouched.
        let xml =
            br#"<w:document xmlns:w="urn:test"><w:p><w:r><w:t>  spaced  </w:t></w:r></w:p></w:document>"#;
        let text = extract_docx_text_from_xml(xml).unwrap();
        assert_eq!(text, "  spaced");
    }
}