src/parser.rs

use anyhow::{anyhow, Context, Result};
use memmap2::Mmap;
use quick_xml::events::Event;
use quick_xml::reader::Reader;
use std::fs::File;
use std::io::{BufRead, Cursor, Read};
use std::path::Path;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum ParserError {
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Unsupported file format: {0}")]
    UnsupportedFormat(String),
    #[error("Corrupt document: {0}")]
    CorruptDocument(String),
}

#[derive(Debug, Clone)]
pub struct ParsedDocument {
    pub path: String,
    pub title: String,
    pub content: String,
    pub extension: String,
    pub size_bytes: u64,
}

pub struct DocumentParser;

impl DocumentParser {
    /// Ingests a file with zero-copy mmap for large files or standard read for small files.
    pub fn parse_file<P: AsRef<Path>>(path: P) -> Result<ParsedDocument> {
        let path_ref = path.as_ref();
        let metadata = std::fs::metadata(path_ref)
            .with_context(|| format!("Failed to read metadata for {:?}", path_ref))?;
        let size_bytes = metadata.len();

        if size_bytes == 0 {
            return Ok(ParsedDocument {
                path: path_ref.to_string_lossy().to_string(),
                title: Self::extract_title(path_ref),
                content: String::new(),
                extension: Self::get_extension(path_ref),
                size_bytes: 0,
            });
        }

        let extension = Self::get_extension(path_ref);
        let title = Self::extract_title(path_ref);

        let file = File::open(path_ref)?;
        
        // Zero-copy memory map for files >= 16 KB; heap read for smaller files
        let content = if size_bytes >= 16 * 1024 {
            let mmap = unsafe { Mmap::map(&file)? };
            Self::extract_content_from_bytes(&mmap, &extension)?
        } else {
            let mut buffer = Vec::with_capacity(size_bytes as usize);
            let mut reader = std::io::BufReader::new(file);
            reader.read_to_end(&mut buffer)?;
            Self::extract_content_from_bytes(&buffer, &extension)?
        };

        Ok(ParsedDocument {
            path: path_ref.to_string_lossy().to_string(),
            title,
            content,
            extension,
            size_bytes,
        })
    }

    fn extract_content_from_bytes(bytes: &[u8], extension: &str) -> Result<String> {
        match extension.to_lowercase().as_str() {
            "txt" | "md" | "markdown" | "log" | "rs" | "py" | "js" | "ts" | "toml" | "yaml" | "yml" => {
                match std::str::from_utf8(bytes) {
                    Ok(valid_str) => Ok(valid_str.to_string()),
                    Err(_) => Ok(String::from_utf8_lossy(bytes).into_owned()),
                }
            }
            "json" => {
                let val: serde_json::Value = serde_json::from_slice(bytes)
                    .context("Invalid JSON document")?;
                if let Some(text) = val.as_str() {
                    Ok(text.to_string())
                } else {
                    Ok(serde_json::to_string_pretty(&val)?)
                }
            }
            "csv" => Self::parse_csv(bytes),
            "docx" => Self::parse_docx(bytes),
            "pdf" => Self::parse_pdf(bytes),
            ext => Err(ParserError::UnsupportedFormat(ext.to_string()).into()),
        }
    }

    fn parse_csv(bytes: &[u8]) -> Result<String> {
        let mut rdr = csv::ReaderBuilder::new()
            .flexible(true)
            .has_headers(true)
            .from_reader(Cursor::new(bytes));

        let mut output = String::with_capacity(bytes.len());
        
        if let Ok(headers) = rdr.headers() {
            output.push_str(&headers.iter().collect::<Vec<_>>().join(" | "));
            output.push('\n');
        }

        for result in rdr.records() {
            let record = result?;
            output.push_str(&record.iter().collect::<Vec<_>>().join(" | "));
            output.push('\n');
        }

        Ok(output)
    }

    fn parse_docx(bytes: &[u8]) -> Result<String> {
        let cursor = Cursor::new(bytes);
        let mut archive = zip::ZipArchive::new(cursor)
            .context("Failed to open DOCX as ZIP archive")?;

        let mut document_xml = archive
            .by_name("word/document.xml")
            .context("Missing word/document.xml inside DOCX archive")?;

        let mut xml_bytes = Vec::new();
        document_xml.read_to_end(&mut xml_bytes)?;

        let mut reader = Reader::from_reader(Cursor::new(xml_bytes));
        reader.config_mut().trim_text(true);

        let mut txt = String::new();
        let mut buf = Vec::new();
        let mut in_text_node = false;

        loop {
            match reader.read_event_into(&mut buf) {
                Ok(Event::Start(ref e)) if e.name().as_ref() == b"w:t" => {
                    in_text_node = true;
                }
                Ok(Event::End(ref e)) if e.name().as_ref() == b"w:t" => {
                    in_text_node = false;
                }
                Ok(Event::End(ref e)) if e.name().as_ref() == b"w:p" => {
                    txt.push('\n');
                }
                Ok(Event::Text(e)) if in_text_node => {
                    if let Ok(s) = e.unescape() {
                        txt.push_str(&s);
                    }
                }
                Ok(Event::Eof) => break,
                Err(e) => return Err(anyhow!("XML Parsing error in DOCX: {}", e)),
                _ => (),
            }
            buf.clear();
        }

        Ok(txt)
    }

    fn parse_pdf(bytes: &[u8]) -> Result<String> {
        pdf_extract::extract_text_from_mem(bytes)
            .context("Failed to extract text stream from PDF")
    }

    fn extract_title(path: &Path) -> String {
        path.file_name()
            .map(|f| f.to_string_lossy().to_string())
            .unwrap_or_else(|| "Untitled".to_string())
    }

    fn get_extension(path: &Path) -> String {
        path.extension()
            .map(|e| e.to_string_lossy().to_string())
            .unwrap_or_else(|| "txt".to_string())
    }
}