memvid_core/

extract.rs

use std::fs;
use std::path::Path;

use crate::{Result, error::MemvidError, text::truncate_at_grapheme_boundary};
#[cfg(feature = "extractous")]
use log::LevelFilter;
#[cfg(feature = "extractous")]
use lopdf::Document as LopdfDocument;
use serde_json::{Value, json};

#[cfg(feature = "extractous")]
use extractous::Extractor;
#[cfg(feature = "extractous")]
use std::collections::HashMap;
#[cfg(feature = "extractous")]
use std::sync::{Mutex, OnceLock};

/// Structured result produced by [`DocumentProcessor`] after running
/// Extractous over an input document.
#[derive(Debug, Clone)]
pub struct ExtractedDocument {
    pub text: Option<String>,
    pub metadata: Value,
    pub mime_type: Option<String>,
}

impl ExtractedDocument {
    pub fn empty() -> Self {
        Self {
            text: None,
            metadata: Value::Null,
            mime_type: None,
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ProcessorConfig {
    pub max_text_chars: usize,
}

impl Default for ProcessorConfig {
    fn default() -> Self {
        Self {
            max_text_chars: 2_000_000,
        }
    }
}

// ============================================================================
// DocumentProcessor - only available with extractous feature
// ============================================================================

#[cfg(feature = "extractous")]
#[derive(Debug)]
pub struct DocumentProcessor {
    extractor: Mutex<Extractor>,
    max_length: usize,
}

#[cfg(feature = "extractous")]
impl Default for DocumentProcessor {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

#[cfg(feature = "extractous")]
static EXTRACTION_CACHE: OnceLock<Mutex<HashMap<blake3::Hash, ExtractedDocument>>> =
    OnceLock::new();

#[cfg(feature = "extractous")]
impl DocumentProcessor {
    pub fn new(config: ProcessorConfig) -> Self {
        let capped = config
            .max_text_chars
            .min(i32::MAX as usize)
            .try_into()
            .unwrap_or(i32::MAX);
        let mut extractor = Extractor::new().set_extract_string_max_length(capped);
        extractor = extractor.set_xml_output(false);
        Self {
            extractor: Mutex::new(extractor),
            max_length: config.max_text_chars,
        }
    }

    pub fn extract_from_path(&self, path: &Path) -> Result<ExtractedDocument> {
        let path_str = path.to_str().ok_or_else(|| MemvidError::ExtractionFailed {
            reason: "input path contains invalid UTF-8".into(),
        })?;

        let extraction = {
            let extractor = self.locked()?;
            let _log_guard = ScopedLogLevel::lowered(LevelFilter::Off);
            extractor.extract_file_to_string(path_str)
        };

        match extraction {
            Ok((mut content, metadata)) => {
                if needs_pdf_fallback(&content) {
                    if let Ok(bytes) = fs::read(path) {
                        if let Ok(Some(fallback_text)) = pdf_text_fallback(&bytes) {
                            content = fallback_text;
                        }
                    }
                }
                Ok(self.into_document(content, metadata))
            }
            Err(err) => {
                let primary_reason = err.to_string();
                if let Ok(bytes) = fs::read(path) {
                    match pdf_text_fallback(&bytes) {
                        Ok(Some(fallback_text)) => {
                            return Ok(self.into_document(fallback_text, pdf_fallback_metadata()));
                        }
                        Ok(None) => {}
                        Err(fallback_err) => {
                            let reason = format!(
                                "primary extractor error: {}; PDF fallback error: {}",
                                primary_reason, fallback_err
                            );
                            return Err(MemvidError::ExtractionFailed {
                                reason: reason.into(),
                            });
                        }
                    }
                }
                Err(MemvidError::ExtractionFailed {
                    reason: primary_reason.into(),
                })
            }
        }
    }

    pub fn extract_from_bytes(&self, bytes: &[u8]) -> Result<ExtractedDocument> {
        let hash = blake3::hash(bytes);
        if let Some(cached) = cache_lookup(&hash) {
            tracing::debug!(target = "memvid::extract", reader = "cache", "cache hit");
            return Ok(cached);
        }

        let extraction = {
            let extractor = self.locked()?;
            let _log_guard = ScopedLogLevel::lowered(LevelFilter::Off);
            extractor.extract_bytes_to_string(bytes)
        };

        let document = match extraction {
            Ok((mut content, metadata)) => {
                let pdf_needed = needs_pdf_fallback(&content);
                tracing::debug!(
                    target: "memvid::extract",
                    content_len = content.len(),
                    pdf_fallback_needed = pdf_needed,
                    "extractous returned content"
                );
                if pdf_needed {
                    match pdf_text_fallback(bytes) {
                        Ok(Some(fallback_text)) => {
                            tracing::debug!(
                                target: "memvid::extract",
                                fallback_len = fallback_text.len(),
                                "lopdf fallback succeeded"
                            );
                            content = fallback_text;
                        }
                        Ok(None) => {
                            tracing::debug!(
                                target: "memvid::extract",
                                "lopdf fallback returned None"
                            );
                            // PDF detected but lopdf couldn't extract any text
                            // Return empty rather than raw PDF bytes
                            content = String::new();
                        }
                        Err(e) => {
                            tracing::debug!(
                                target: "memvid::extract",
                                error = %e,
                                "lopdf fallback failed"
                            );
                            // lopdf extraction failed - return empty rather than raw PDF bytes
                            content = String::new();
                        }
                    }
                }
                self.into_document(content, metadata)
            }
            Err(err) => {
                let primary_reason = err.to_string();
                match pdf_text_fallback(bytes) {
                    Ok(Some(fallback_text)) => {
                        self.into_document(fallback_text, pdf_fallback_metadata())
                    }
                    Ok(None) => {
                        return Err(MemvidError::ExtractionFailed {
                            reason: primary_reason.into(),
                        });
                    }
                    Err(fallback_err) => {
                        let reason = format!(
                            "primary extractor error: {}; PDF fallback error: {}",
                            primary_reason, fallback_err
                        );
                        return Err(MemvidError::ExtractionFailed {
                            reason: reason.into(),
                        });
                    }
                }
            }
        };

        cache_store(hash, &document);
        Ok(document)
    }

    fn locked(&self) -> Result<std::sync::MutexGuard<'_, Extractor>> {
        self.extractor
            .lock()
            .map_err(|_| MemvidError::ExtractionFailed {
                reason: "extractor mutex poisoned".into(),
            })
    }

    fn into_document<M>(&self, content: String, metadata: M) -> ExtractedDocument
    where
        M: serde::Serialize,
    {
        let metadata_value = serde_json::to_value(metadata).unwrap_or(Value::Null);
        let mime_type = metadata_value.get("Content-Type").and_then(value_to_mime);

        let text = if content.trim().is_empty() {
            tracing::debug!(
                target: "memvid::extract",
                "into_document: content is empty, returning text=None"
            );
            None
        } else {
            let final_text = if content.len() > self.max_length {
                let end = truncate_at_grapheme_boundary(&content, self.max_length);
                content[..end].to_string()
            } else {
                content
            };
            tracing::debug!(
                target: "memvid::extract",
                text_len = final_text.len(),
                starts_with_pdf = final_text.starts_with("%PDF"),
                "into_document: returning text"
            );
            Some(final_text)
        };

        ExtractedDocument {
            text,
            metadata: metadata_value,
            mime_type,
        }
    }
}

// ============================================================================
// Stub DocumentProcessor when extractous is disabled - returns clear error
// ============================================================================

#[cfg(not(feature = "extractous"))]
#[derive(Debug)]
pub struct DocumentProcessor {
    max_length: usize,
}

#[cfg(not(feature = "extractous"))]
impl Default for DocumentProcessor {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

#[cfg(not(feature = "extractous"))]
impl DocumentProcessor {
    pub fn new(config: ProcessorConfig) -> Self {
        Self {
            max_length: config.max_text_chars,
        }
    }

    pub fn extract_from_path(&self, path: &Path) -> Result<ExtractedDocument> {
        // Without extractous, we can still handle plain text files
        let bytes = fs::read(path).map_err(|e| MemvidError::ExtractionFailed {
            reason: format!("failed to read file: {e}").into(),
        })?;
        self.extract_from_bytes(&bytes)
    }

    pub fn extract_from_bytes(&self, bytes: &[u8]) -> Result<ExtractedDocument> {
        // Without extractous, we can still handle plain text files and common text-based formats
        // Try to interpret as UTF-8 text first
        if let Ok(text) = std::str::from_utf8(bytes) {
            // Check if it's likely text (no null bytes in first 8KB)
            let sample = &bytes[..bytes.len().min(8192)];
            if !sample.contains(&0) {
                let truncate_len = truncate_at_grapheme_boundary(text, self.max_length);
                let truncated = &text[..truncate_len];
                return Ok(ExtractedDocument {
                    text: Some(truncated.to_string()),
                    metadata: json!({}),
                    mime_type: Some("text/plain".to_string()),
                });
            }
        }

        // For binary content (video, audio, images, etc.), return success with no text.
        // This allows binary blobs to be stored without requiring the extractous feature.
        // The caller can still store the blob; there just won't be extracted text for search.
        Ok(ExtractedDocument {
            text: None,
            metadata: json!({}),
            mime_type: Some("application/octet-stream".to_string()),
        })
    }
}

#[cfg(feature = "extractous")]
fn needs_pdf_fallback(content: &str) -> bool {
    if content.trim().is_empty() {
        return true;
    }
    looks_like_pdf_structure_dump(content)
}

#[cfg(feature = "extractous")]
fn pdf_fallback_metadata() -> Value {
    json!({
        "Content-Type": "application/pdf",
        "extraction": "lopdf_fallback",
    })
}

#[cfg(feature = "extractous")]
const PDF_FALLBACK_MAX_BYTES: usize = 64 * 1024 * 1024; // 64 MiB hard cap.
#[cfg(feature = "extractous")]
const PDF_FALLBACK_MAX_PAGES: usize = 4_096;

#[cfg(feature = "extractous")]
fn pdf_text_fallback(bytes: &[u8]) -> Result<Option<String>> {
    if !is_probably_pdf(bytes) {
        return Ok(None);
    }

    if bytes.len() > PDF_FALLBACK_MAX_BYTES {
        return Err(MemvidError::ExtractionFailed {
            reason: format!(
                "pdf fallback aborted: {} bytes exceeds limit of {} bytes",
                bytes.len(),
                PDF_FALLBACK_MAX_BYTES
            )
            .into(),
        });
    }

    let _log_guard = ScopedLogLevel::lowered(LevelFilter::Off);
    let mut document =
        LopdfDocument::load_mem(bytes).map_err(|err| MemvidError::ExtractionFailed {
            reason: format!("pdf fallback failed to load document: {err}").into(),
        })?;

    if document.is_encrypted() {
        if document.decrypt("").is_err() {
            return Err(MemvidError::ExtractionFailed {
                reason: "pdf fallback cannot decrypt password-protected file".into(),
            });
        }
    }

    let _ = document.decompress();

    let mut page_numbers: Vec<u32> = document.get_pages().keys().copied().collect();
    if page_numbers.is_empty() {
        return Ok(None);
    }
    page_numbers.sort_unstable();

    if page_numbers.len() > PDF_FALLBACK_MAX_PAGES {
        return Err(MemvidError::ExtractionFailed {
            reason: format!(
                "pdf fallback aborted: page count {} exceeds limit of {}",
                page_numbers.len(),
                PDF_FALLBACK_MAX_PAGES
            )
            .into(),
        });
    }

    match document.extract_text(&page_numbers) {
        Ok(text) => {
            let trimmed = text.trim();
            if trimmed.is_empty() {
                Ok(None)
            } else {
                Ok(Some(trimmed.to_string()))
            }
        }
        Err(err) => Err(MemvidError::ExtractionFailed {
            reason: format!("pdf fallback failed to extract text: {err}").into(),
        }),
    }
}

#[cfg(feature = "extractous")]
struct ScopedLogLevel {
    previous: LevelFilter,
    changed: bool,
}

#[cfg(feature = "extractous")]
impl ScopedLogLevel {
    fn lowered(level: LevelFilter) -> Self {
        let previous = log::max_level();
        if level < previous {
            log::set_max_level(level);
            Self {
                previous,
                changed: true,
            }
        } else {
            Self {
                previous,
                changed: false,
            }
        }
    }
}

#[cfg(feature = "extractous")]
impl Drop for ScopedLogLevel {
    fn drop(&mut self) {
        if self.changed {
            log::set_max_level(self.previous);
        }
    }
}

#[cfg(feature = "extractous")]
fn is_probably_pdf(bytes: &[u8]) -> bool {
    if bytes.is_empty() {
        return false;
    }
    let mut slice = bytes;
    if slice.starts_with(&[0xEF, 0xBB, 0xBF]) {
        slice = &slice[3..];
    }
    while let Some((first, rest)) = slice.split_first() {
        if *first == 0 || first.is_ascii_whitespace() {
            slice = rest;
        } else {
            break;
        }
    }
    slice.starts_with(b"%PDF")
}

#[cfg(feature = "extractous")]
fn looks_like_pdf_structure_dump(content: &str) -> bool {
    if content.len() < 256 {
        return false;
    }
    let sample_len = content.len().min(8_192);
    // Find the nearest valid character boundary at or before sample_len
    let safe_len = truncate_at_grapheme_boundary(content, sample_len);
    let sample = &content[..safe_len];
    let endobj_hits = sample.matches("endobj").take(3).count();
    if endobj_hits < 2 {
        return false;
    }
    let has_obj =
        sample.contains(" 0 obj") || sample.contains("\n0 obj") || sample.contains("\r0 obj");
    let has_stream = sample.contains("endstream");
    let has_page_type = sample.contains("/Type /Page");
    endobj_hits >= 2 && (has_obj || has_stream || has_page_type)
}

#[cfg(feature = "extractous")]
fn value_to_mime(value: &Value) -> Option<String> {
    if let Some(mime) = value.as_str() {
        return Some(mime.to_string());
    }
    if let Some(array) = value.as_array() {
        for entry in array {
            if let Some(mime) = entry.as_str() {
                return Some(mime.to_string());
            }
        }
    }
    None
}

#[cfg(feature = "extractous")]
fn cache_lookup(hash: &blake3::Hash) -> Option<ExtractedDocument> {
    let cache = EXTRACTION_CACHE.get_or_init(|| Mutex::new(HashMap::new()));
    cache.lock().ok().and_then(|map| map.get(hash).cloned())
}

#[cfg(feature = "extractous")]
fn cache_store(hash: blake3::Hash, document: &ExtractedDocument) {
    let cache = EXTRACTION_CACHE.get_or_init(|| Mutex::new(HashMap::new()));
    if let Ok(mut map) = cache.lock() {
        map.insert(hash, document.clone());
    }
}

#[cfg(all(test, feature = "extractous"))]
mod tests {
    use super::*;

    #[test]
    fn detects_pdf_like_dump() {
        let snippet = "binary %PDF snippet endobj endobj endstream 0 obj /Type /Page endobj ";
        let dump = snippet.repeat(12);
        assert!(looks_like_pdf_structure_dump(&dump));
    }

    #[test]
    fn skips_normal_text() {
        let text = "This is a perfectly normal paragraph that should not trigger the PDF fallback.";
        assert!(!looks_like_pdf_structure_dump(text));
    }

    #[test]
    fn identifies_pdf_magic() {
        let bytes = b"%PDF-1.7 some data";
        assert!(is_probably_pdf(bytes));
        let padded = b"\n\n%PDF-1.5";
        assert!(is_probably_pdf(padded));
        let not_pdf = b"<!doctype html>";
        assert!(!is_probably_pdf(not_pdf));
    }
}

#[cfg(all(test, feature = "extractous"))]
mod pdf_fix_tests {
    use super::*;

    /// Test that PDF extraction via lopdf fallback returns actual text content,
    /// not raw PDF bytes. This test verifies the fix for the bug where PDFs
    /// with extractous returning empty content would have their raw bytes
    /// indexed instead of extracted text.
    #[test]
    fn test_pdf_structure_dump_detection_prevents_raw_indexing() {
        // Create a synthetic PDF-like structure that should trigger the fallback
        let pdf_structure = b"%PDF-1.4\n%\xff\xff\xff\xff\n1 0 obj\n<</Type/Catalog>>\nendobj\n";

        // This should be detected as PDF structure
        assert!(is_probably_pdf(pdf_structure));

        // And content that looks like PDF structure should be detected
        let structure_dump =
            "binary %PDF snippet endobj endobj endstream 0 obj /Type /Page endobj ".repeat(12);
        assert!(looks_like_pdf_structure_dump(&structure_dump));

        // Normal text should NOT be detected as PDF structure
        let normal_text = "This is perfectly normal extracted text from a document.";
        assert!(!looks_like_pdf_structure_dump(normal_text));
    }
}