codecraft/server/src/docs.rs

//! Text extraction for documents the read tool cannot show as text:
//! PDF (poppler-utils), spreadsheets (calamine), DOCX and PPTX (zip + XML).

use crate::types::ImageData;
use anyhow::{Context, Result, bail};
use base64::Engine;
use std::io::Read as _;
use std::path::Path;
use tokio::process::Command;

pub const DOC_EXTENSIONS: &[&str] = &["pdf", "docx", "pptx", "xlsx", "xlsm", "xls", "ods"];

const MAX_PDF_PAGES: u32 = 30;
const MAX_RENDERED_PAGES: usize = 4;
const MAX_ROWS: usize = 2000;

pub fn is_document(path: &Path) -> bool {
    path.extension()
        .and_then(|e| e.to_str())
        .is_some_and(|e| DOC_EXTENSIONS.contains(&e.to_ascii_lowercase().as_str()))
}

/// Extract readable content. `pages` ("3" or "2-5") applies to PDFs, `sheet` to spreadsheets.
pub async fn extract(path: &Path, pages: Option<&str>, sheet: Option<&str>) -> Result<(String, Vec<ImageData>)> {
    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_ascii_lowercase();
    match ext.as_str() {
        "pdf" => pdf(path, pages).await,
        "docx" => {
            let p = path.to_path_buf();
            Ok((tokio::task::spawn_blocking(move || docx(&p)).await??, vec![]))
        }
        "pptx" => {
            let p = path.to_path_buf();
            Ok((tokio::task::spawn_blocking(move || pptx(&p)).await??, vec![]))
        }
        _ => {
            let p = path.to_path_buf();
            let s = sheet.map(String::from);
            Ok((tokio::task::spawn_blocking(move || spreadsheet(&p, s.as_deref())).await??, vec![]))
        }
    }
}

fn parse_range(spec: Option<&str>, total: u32) -> (u32, u32) {
    let Some(spec) = spec.map(str::trim).filter(|s| !s.is_empty()) else {
        return (1, total.min(MAX_PDF_PAGES));
    };
    let (a, b) = spec.split_once('-').unwrap_or((spec, spec));
    let first = a.trim().parse::<u32>().unwrap_or(1).clamp(1, total.max(1));
    let last = b.trim().parse::<u32>().unwrap_or(first).clamp(first, total.max(first));
    (first, last.min(first + MAX_PDF_PAGES - 1))
}

async fn tool_output(cmd: &str, args: &[&str]) -> Result<std::process::Output> {
    let out = Command::new(cmd)
        .args(args)
        .kill_on_drop(true)
        .output()
        .await
        .with_context(|| format!("'{cmd}' is not installed (poppler-utils)"))?;
    if !out.status.success() {
        bail!("{cmd} failed: {}", String::from_utf8_lossy(&out.stderr).trim());
    }
    Ok(out)
}

async fn pdf(path: &Path, pages: Option<&str>) -> Result<(String, Vec<ImageData>)> {
    let p = path.to_string_lossy();
    let info = tool_output("pdfinfo", &[&p]).await?;
    let info = String::from_utf8_lossy(&info.stdout);
    let total: u32 = info
        .lines()
        .find_map(|l| l.strip_prefix("Pages:").map(|v| v.trim().parse().unwrap_or(1)))
        .unwrap_or(1);
    let (first, last) = parse_range(pages, total);
    let (f, l) = (first.to_string(), last.to_string());
    let out = tool_output("pdftotext", &["-layout", "-f", &f, "-l", &l, &p, "-"]).await?;
    let text = String::from_utf8_lossy(&out.stdout);

    let mut result = format!("PDF with {total} pages – showing pages {first}-{last}.\n");
    let mut images = Vec::new();
    // pdftotext separates pages with form feeds.
    for (i, page) in text.split('\u{c}').enumerate() {
        let n = first + i as u32;
        if n > last {
            break;
        }
        result.push_str(&format!("\n--- Page {n} ---\n"));
        if page.trim().is_empty() {
            // Scanned page: render it so a vision model can read it.
            if images.len() < MAX_RENDERED_PAGES {
                match render_page(&p, n).await {
                    Ok(img) => {
                        images.push(img);
                        result.push_str("(no text layer – rendered as image)\n");
                    }
                    Err(e) => result.push_str(&format!("(no text layer; rendering failed: {e})\n")),
                }
            } else {
                result.push_str("(no text layer – read this page separately with pages)\n");
            }
        } else {
            result.push_str(page.trim_end());
            result.push('\n');
        }
    }
    if last < total {
        result.push_str(&format!("\n[{} more pages – use pages=\"{}-{}\"]\n", total - last, last + 1, (last + MAX_PDF_PAGES).min(total)));
    }
    Ok((result, images))
}

async fn render_page(pdf: &str, page: u32) -> Result<ImageData> {
    let dir = std::env::temp_dir().join(format!("cc-pdf-{}", crate::types::new_id()));
    tokio::fs::create_dir_all(&dir).await?;
    let prefix = dir.join("p");
    let n = page.to_string();
    let res = tool_output("pdftoppm", &["-png", "-r", "110", "-f", &n, "-l", &n, "-singlefile", pdf, &prefix.to_string_lossy()]).await;
    let file = dir.join("p.png");
    let bytes = match res {
        Ok(_) => tokio::fs::read(&file).await,
        Err(e) => {
            let _ = tokio::fs::remove_dir_all(&dir).await;
            return Err(e);
        }
    };
    let _ = tokio::fs::remove_dir_all(&dir).await;
    let img = image::load_from_memory(&bytes?)?;
    let img = img.resize(1568, 1568, image::imageops::FilterType::Triangle);
    let mut buf = Vec::new();
    img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)?;
    Ok(ImageData { blob: None, size: None, media_type: "image/png".into(), data: base64::engine::general_purpose::STANDARD.encode(buf) })
}

fn spreadsheet(path: &Path, sheet: Option<&str>) -> Result<String> {
    use calamine::Reader;
    let mut wb = calamine::open_workbook_auto(path).context("cannot open spreadsheet")?;
    let names = wb.sheet_names();
    let selected: Vec<String> = match sheet {
        Some(s) => {
            if !names.iter().any(|n| n == s) {
                bail!("no sheet '{s}' – sheets: {}", names.join(", "));
            }
            vec![s.to_string()]
        }
        None => names.clone(),
    };
    let mut out = format!("Spreadsheet with {} sheet(s): {}\n", names.len(), names.join(", "));
    for name in selected {
        let range = wb.worksheet_range(&name).with_context(|| format!("cannot read sheet {name}"))?;
        let (rows, cols) = range.get_size();
        out.push_str(&format!("\n--- Sheet \"{name}\" ({rows} rows × {cols} columns, tab-separated) ---\n"));
        for (i, row) in range.rows().enumerate() {
            if i >= MAX_ROWS {
                out.push_str(&format!("[... {} more rows ...]\n", rows - MAX_ROWS));
                break;
            }
            let line: Vec<String> = row.iter().map(|c| c.to_string().replace(['\t', '\n'], " ")).collect();
            out.push_str(line.join("\t").trim_end());
            out.push('\n');
        }
    }
    Ok(out)
}

fn zip_text(path: &Path, entry: &str) -> Result<String> {
    let file = std::fs::File::open(path)?;
    let mut zip = zip::ZipArchive::new(file).context("not a valid Office file")?;
    let mut f = zip.by_name(entry).with_context(|| format!("missing {entry}"))?;
    let mut s = String::new();
    f.read_to_string(&mut s)?;
    Ok(s)
}

fn unescape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut rest = s;
    while let Some(i) = rest.find('&') {
        out.push_str(&rest[..i]);
        let tail = &rest[i..];
        let Some(end) = tail.find(';').filter(|e| *e < 10) else {
            out.push('&');
            rest = &tail[1..];
            continue;
        };
        let ent = &tail[1..end];
        let ch = match ent {
            "amp" => Some('&'),
            "lt" => Some('<'),
            "gt" => Some('>'),
            "quot" => Some('"'),
            "apos" => Some('\''),
            _ if ent.starts_with("#x") => u32::from_str_radix(&ent[2..], 16).ok().and_then(char::from_u32),
            _ if ent.starts_with('#') => ent[1..].parse().ok().and_then(char::from_u32),
            _ => None,
        };
        match ch {
            Some(c) => out.push(c),
            None => out.push_str(&tail[..=end]),
        }
        rest = &tail[end + 1..];
    }
    out.push_str(rest);
    out
}

/// Collect text from WordprocessingML / DrawingML: `w:t`/`a:t` runs,
/// paragraphs, tabs, breaks and table cells.
pub fn xml_text(xml: &str) -> String {
    let mut out = String::new();
    let mut rest = xml;
    let mut in_text = false;
    while let Some(lt) = rest.find('<') {
        if in_text {
            out.push_str(&unescape(&rest[..lt]));
        }
        let Some(gt) = rest[lt..].find('>') else { break };
        let tag = &rest[lt + 1..lt + gt];
        rest = &rest[lt + gt + 1..];
        let closing = tag.starts_with('/');
        let self_closing = tag.ends_with('/');
        let name = tag.trim_start_matches('/').split(|c: char| c.is_whitespace() || c == '/').next().unwrap_or("");
        match name {
            "w:t" | "a:t" => in_text = !closing && !self_closing,
            "w:tab" if !closing => out.push('\t'),
            "w:br" | "a:br" | "w:cr" if !closing => out.push('\n'),
            "w:p" | "a:p" if closing => out.push('\n'),
            "w:tc" if closing => out.push_str(" | "),
            "w:tr" if closing => out.push('\n'),
            _ => {}
        }
    }
    // Collapse runs of empty lines.
    let mut cleaned = String::new();
    let mut blank = 0;
    for line in out.lines() {
        if line.trim().is_empty() {
            blank += 1;
            if blank > 1 {
                continue;
            }
        } else {
            blank = 0;
        }
        cleaned.push_str(line.trim_end());
        cleaned.push('\n');
    }
    cleaned
}

fn docx(path: &Path) -> Result<String> {
    let xml = zip_text(path, "word/document.xml")?;
    Ok(format!("Word document:\n\n{}", xml_text(&xml)))
}

fn pptx(path: &Path) -> Result<String> {
    let file = std::fs::File::open(path)?;
    let zip = zip::ZipArchive::new(file).context("not a valid PowerPoint file")?;
    let mut slides: Vec<(u32, String)> = zip
        .file_names()
        .filter_map(|n| n.ok().map(|n| n.to_string()))
        .filter_map(|n| {
            let num = n.strip_prefix("ppt/slides/slide")?.strip_suffix(".xml")?.parse().ok()?;
            Some((num, n))
        })
        .collect();
    slides.sort();
    let mut out = format!("Presentation with {} slides:\n", slides.len());
    for (num, name) in slides {
        let xml = zip_text(path, &name)?;
        out.push_str(&format!("\n--- Slide {num} ---\n{}", xml_text(&xml)));
    }
    Ok(out)
}

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

    #[test]
    fn extracts_word_xml() {
        let xml = r#"<w:document><w:body><w:p><w:r><w:t>Hello &amp; welcome</w:t></w:r><w:r><w:tab/><w:t xml:space="preserve">to CodeCraft</w:t></w:r></w:p><w:tbl><w:tr><w:tc><w:p><w:r><w:t>A1</w:t></w:r></w:p></w:tc><w:tc><w:p><w:r><w:t>B1</w:t></w:r></w:p></w:tc></w:tr></w:tbl></w:body></w:document>"#;
        let t = xml_text(xml);
        assert!(t.contains("Hello & welcome\tto CodeCraft"));
        assert!(t.contains("A1"));
        assert!(t.contains("B1"));
    }

    #[test]
    fn page_ranges() {
        assert_eq!(parse_range(None, 100), (1, 30));
        assert_eq!(parse_range(Some("5"), 100), (5, 5));
        assert_eq!(parse_range(Some("3-7"), 100), (3, 7));
        assert_eq!(parse_range(Some("90-200"), 100), (90, 100));
    }

    #[test]
    fn unescapes_entities() {
        assert_eq!(unescape("a &lt;b&gt; &#228; &#x263A; &bogus; &"), "a <b> ä ☺ &bogus; &");
    }
}