kcode-doc-extraction 0.1.0

//! Local PDF, DOCX, spreadsheet, and plain-text extraction.

#![deny(missing_docs)]
#![forbid(unsafe_code)]

use std::{
    error::Error as StdError,
    fmt,
    io::{Cursor, Read},
};

use calamine::{Reader as _, open_workbook_auto_from_rs};
use quick_xml::{Reader as XmlReader, events::Event as XmlEvent};

/// Supported local document representation.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DocumentFormat {
    /// Searchable PDF text extraction without OCR.
    Pdf,
    /// Microsoft Word Open XML document.
    Docx,
    /// XLSX, XLS, XLSB, or ODS workbook.
    Spreadsheet,
    /// UTF-8-family text, including CSV, TSV, Markdown, JSON, YAML, and XML.
    PlainText,
}

impl DocumentFormat {
    /// Returns the stable format label.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Pdf => "pdf",
            Self::Docx => "docx",
            Self::Spreadsheet => "spreadsheet",
            Self::PlainText => "text",
        }
    }
}

/// Stable document-extraction failure classification.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ErrorKind {
    /// Input metadata, size, or bytes were invalid.
    InvalidInput,
    /// The filename and media type do not identify a supported format.
    UnsupportedFormat,
    /// A local parser could not extract the document.
    ExtractionFailed,
    /// The parser found no readable text.
    EmptyText,
}

/// A local document validation or extraction failure.
#[derive(Debug)]
pub struct Error {
    kind: ErrorKind,
    message: String,
}

impl Error {
    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: message.into(),
        }
    }

    /// Returns the stable failure classification.
    pub const fn kind(&self) -> ErrorKind {
        self.kind
    }

    /// Returns the parser or validation detail.
    pub fn message(&self) -> &str {
        &self.message
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl StdError for Error {}

/// Result type returned by document extraction.
pub type Result<T> = std::result::Result<T, Error>;

/// Input and output resource limits.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExtractionLimits {
    /// Maximum in-memory document bytes.
    pub max_bytes: usize,
    /// Maximum returned Unicode characters.
    pub max_characters: usize,
}

impl Default for ExtractionLimits {
    fn default() -> Self {
        Self {
            max_bytes: 20 * 1024 * 1024,
            max_characters: 1_000_000,
        }
    }
}

/// One complete in-memory document.
#[derive(Clone, Eq, PartialEq)]
pub struct DocumentInput {
    /// Caller-supplied filename used for format detection and returned safely.
    pub file_name: String,
    /// Caller-supplied media type used as a format-detection fallback.
    pub content_type: String,
    /// Complete document bytes.
    pub data: Vec<u8>,
}

impl fmt::Debug for DocumentInput {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("DocumentInput")
            .field("file_name", &self.file_name)
            .field("content_type", &self.content_type)
            .field("bytes", &self.data.len())
            .finish()
    }
}

/// Successfully extracted normalized text and metadata.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExtractedDocument {
    /// Safe display filename.
    pub file_name: String,
    /// Original normalized media type.
    pub content_type: String,
    /// Detected document format.
    pub format: DocumentFormat,
    /// Normalized, non-empty extracted text.
    pub text: String,
    /// Returned Unicode character count.
    pub characters: usize,
    /// Whether the character limit truncated the text.
    pub truncated: bool,
}

/// Stateless local document extractor with fixed limits.
#[derive(Clone, Debug)]
pub struct DocumentExtractor {
    limits: ExtractionLimits,
}

impl DocumentExtractor {
    /// Constructs an extractor after validating nonzero limits.
    pub fn new(limits: ExtractionLimits) -> Result<Self> {
        if limits.max_bytes == 0 || limits.max_characters == 0 {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "document byte and character limits must be greater than zero",
            ));
        }
        Ok(Self { limits })
    }

    /// Extracts and normalizes one complete document in the current thread.
    pub fn extract(&self, input: DocumentInput) -> Result<ExtractedDocument> {
        if input.data.is_empty() || input.data.len() > self.limits.max_bytes {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                format!(
                    "document must contain between 1 and {} bytes",
                    self.limits.max_bytes
                ),
            ));
        }
        let content_type = input
            .content_type
            .split(';')
            .next()
            .unwrap_or(&input.content_type)
            .trim()
            .to_ascii_lowercase();
        let format = detect_format(&input.file_name, &content_type).ok_or_else(|| {
            Error::new(
                ErrorKind::UnsupportedFormat,
                "supported documents are PDF, DOCX, XLSX, XLS, XLSB, ODS, CSV, TSV, and plain text",
            )
        })?;
        let file_name = safe_filename(&input.file_name, format);
        let extracted = extract_text(format, &input.data).map_err(|message| {
            Error::new(
                ErrorKind::ExtractionFailed,
                format!("document could not be converted to text: {message}"),
            )
        })?;
        if extracted.is_empty() {
            let message = if format == DocumentFormat::Pdf {
                "PDF contains no extractable text and may require OCR"
            } else {
                "document contains no extractable text"
            };
            return Err(Error::new(ErrorKind::EmptyText, message));
        }
        let (text, truncated) = truncate_characters(&extracted, self.limits.max_characters);
        let characters = text.chars().count();
        Ok(ExtractedDocument {
            file_name,
            content_type,
            format,
            text,
            characters,
            truncated,
        })
    }
}

impl Default for DocumentExtractor {
    fn default() -> Self {
        Self::new(ExtractionLimits::default()).expect("default document limits are valid")
    }
}

fn detect_format(file_name: &str, content_type: &str) -> Option<DocumentFormat> {
    let extension = file_name
        .rsplit_once('.')
        .map(|(_, extension)| extension.to_ascii_lowercase());
    match extension.as_deref() {
        Some("pdf") => Some(DocumentFormat::Pdf),
        Some("docx") => Some(DocumentFormat::Docx),
        Some("xlsx" | "xls" | "xlsb" | "ods") => Some(DocumentFormat::Spreadsheet),
        Some("csv" | "tsv" | "txt" | "md" | "json" | "yaml" | "yml" | "xml") => {
            Some(DocumentFormat::PlainText)
        }
        _ if content_type == "application/pdf" => Some(DocumentFormat::Pdf),
        _ if content_type
            == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" =>
        {
            Some(DocumentFormat::Docx)
        }
        _ if matches!(
            content_type,
            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                | "application/vnd.ms-excel"
                | "application/vnd.ms-excel.sheet.binary.macroenabled.12"
                | "application/vnd.oasis.opendocument.spreadsheet"
        ) =>
        {
            Some(DocumentFormat::Spreadsheet)
        }
        _ if content_type.starts_with("text/")
            || matches!(
                content_type,
                "application/json" | "application/xml" | "application/yaml" | "application/x-yaml"
            ) =>
        {
            Some(DocumentFormat::PlainText)
        }
        _ => None,
    }
}

fn safe_filename(value: &str, format: DocumentFormat) -> String {
    let cleaned = value
        .chars()
        .filter(|character| {
            !character.is_control() && !matches!(character, '/' | '\\' | ':' | '\0')
        })
        .take(200)
        .collect::<String>();
    if cleaned.trim().is_empty() {
        format!("document.{}", format.as_str())
    } else {
        cleaned
    }
}

fn local_xml_name(name: &[u8]) -> &[u8] {
    name.rsplit(|byte| *byte == b':').next().unwrap_or(name)
}

fn extract_docx_text(bytes: &[u8]) -> std::result::Result<String, String> {
    let mut archive =
        zip::ZipArchive::new(Cursor::new(bytes)).map_err(|error| error.to_string())?;
    let mut document = archive
        .by_name("word/document.xml")
        .map_err(|error| format!("word/document.xml: {error}"))?;
    let mut xml = String::new();
    document
        .read_to_string(&mut xml)
        .map_err(|error| error.to_string())?;
    let mut reader = XmlReader::from_str(&xml);
    let mut output = String::new();
    let mut in_text = false;
    let mut in_cell = false;
    loop {
        match reader.read_event() {
            Ok(XmlEvent::Start(event)) => match local_xml_name(event.name().as_ref()) {
                b"t" => in_text = true,
                b"tc" => in_cell = true,
                _ => {}
            },
            Ok(XmlEvent::Empty(event)) => match local_xml_name(event.name().as_ref()) {
                b"tab" => output.push('\t'),
                b"br" | b"cr" => output.push('\n'),
                _ => {}
            },
            Ok(XmlEvent::Text(text)) if in_text => {
                let decoded = text.decode().map_err(|error| error.to_string())?;
                output.push_str(&decoded);
            }
            Ok(XmlEvent::CData(text)) if in_text => {
                let decoded = text.decode().map_err(|error| error.to_string())?;
                output.push_str(&decoded);
            }
            Ok(XmlEvent::GeneralRef(reference)) if in_text => {
                let decoded = reference.decode().map_err(|error| error.to_string())?;
                let escaped = format!("&{decoded};");
                let resolved =
                    quick_xml::escape::unescape(&escaped).map_err(|error| error.to_string())?;
                output.push_str(&resolved);
            }
            Ok(XmlEvent::End(event)) => match local_xml_name(event.name().as_ref()) {
                b"t" => in_text = false,
                b"p" if in_cell => output.push_str(" / "),
                b"p" => output.push('\n'),
                b"tc" => {
                    if output.ends_with(" / ") {
                        output.truncate(output.len() - 3);
                    }
                    output.push('\t');
                    in_cell = false;
                }
                b"tr" => {
                    if output.ends_with('\t') {
                        output.pop();
                    }
                    output.push('\n');
                }
                _ => {}
            },
            Ok(XmlEvent::Eof) => break,
            Err(error) => return Err(error.to_string()),
            _ => {}
        }
    }
    Ok(output)
}

fn extract_spreadsheet_text(bytes: &[u8]) -> std::result::Result<String, String> {
    let mut workbook = open_workbook_auto_from_rs(Cursor::new(bytes.to_vec()))
        .map_err(|error| error.to_string())?;
    let sheet_names = workbook.sheet_names().to_vec();
    let mut output = String::new();
    for (sheet_index, sheet_name) in sheet_names.iter().enumerate() {
        if sheet_index > 0 {
            output.push('\n');
        }
        output.push_str("Sheet: ");
        output.push_str(sheet_name);
        output.push('\n');
        let range = workbook
            .worksheet_range(sheet_name)
            .map_err(|error| format!("{sheet_name}: {error}"))?;
        for row in range.rows() {
            let line = row
                .iter()
                .map(ToString::to_string)
                .map(|cell| cell.replace(['\t', '\r', '\n'], " "))
                .collect::<Vec<_>>()
                .join("\t");
            output.push_str(line.trim_end());
            output.push('\n');
        }
    }
    Ok(output)
}

fn normalize_text(value: String) -> String {
    let normalized = value
        .replace("\r\n", "\n")
        .replace('\r', "\n")
        .replace('\0', "");
    let mut output = String::with_capacity(normalized.len());
    let mut blank_lines = 0;
    for line in normalized.lines() {
        let line = line.trim_end();
        if line.trim().is_empty() {
            blank_lines += 1;
            if blank_lines > 1 {
                continue;
            }
        } else {
            blank_lines = 0;
        }
        output.push_str(line);
        output.push('\n');
    }
    output.trim().to_owned()
}

fn extract_text(format: DocumentFormat, bytes: &[u8]) -> std::result::Result<String, String> {
    let extracted = match format {
        DocumentFormat::Pdf => {
            pdf_extract::extract_text_from_mem(bytes).map_err(|error| error.to_string())?
        }
        DocumentFormat::Docx => extract_docx_text(bytes)?,
        DocumentFormat::Spreadsheet => extract_spreadsheet_text(bytes)?,
        DocumentFormat::PlainText => String::from_utf8_lossy(bytes)
            .trim_start_matches('\u{feff}')
            .to_owned(),
    };
    Ok(normalize_text(extracted))
}

fn truncate_characters(value: &str, limit: usize) -> (String, bool) {
    let mut iter = value.char_indices();
    let Some((boundary, _)) = iter.nth(limit) else {
        return (value.to_owned(), false);
    };
    (value[..boundary].trim_end().to_owned(), true)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn formats_are_selected_from_extensions_and_media_types() {
        assert_eq!(
            detect_format("report.PDF", "application/octet-stream"),
            Some(DocumentFormat::Pdf)
        );
        assert_eq!(
            detect_format("book.xlsx", "application/octet-stream"),
            Some(DocumentFormat::Spreadsheet)
        );
        assert_eq!(
            detect_format("notes", "text/plain; charset=utf-8"),
            Some(DocumentFormat::PlainText)
        );
        assert_eq!(detect_format("archive.zip", "application/zip"), None);
    }

    #[test]
    fn plain_text_is_normalized_and_bounded() {
        let extractor = DocumentExtractor::new(ExtractionLimits {
            max_bytes: 100,
            max_characters: 5,
        })
        .unwrap();
        let extracted = extractor
            .extract(DocumentInput {
                file_name: "notes.txt".into(),
                content_type: "text/plain".into(),
                data: b"hello\r\n\r\n\r\nworld\0".to_vec(),
            })
            .unwrap();
        assert_eq!(extracted.text, "hello");
        assert!(extracted.truncated);
    }
}