spectre_pdf 1.0.0

//! Persistent PDF handle. Open once, query many. Each method is the
//! same logic as the matching free function but reuses one parse.

use crate::geom::Rect;
use crate::positioned::{
    extract_blocks_from_doc, extract_blocks_streaming_from_doc, extract_text_positioned_from_doc,
    extract_words_from_doc,
};
use crate::search::{search_in_words_pub, SearchOptions};
use crate::structure::{
    extract_annotations_from_sp, extract_document_info_from_sp, extract_images_from_sp,
    extract_links_from_sp, extract_page_info_from_sp, extract_toc_from_sp,
};
use crate::{
    Annotation, DocumentInfo, ExtractError, ImageInfo, Link, PageInfo, PositionedPage, SearchHit,
    TextBlock, TocEntry, Word,
};

/// A parsed PDF document. Open once, query many times.
///
/// Use this when pulling more than one surface from the same PDF; the
/// free functions reparse on every call.
///
/// ```no_run
/// use spectre_pdf::Document;
/// let bytes = std::fs::read("doc.pdf").unwrap();
/// let doc = Document::open(&bytes).unwrap();
/// let words = doc.words(None).unwrap();
/// let toc = doc.toc().unwrap();
/// ```
pub struct Document {
    parse: spectre_parse::Document,
}

impl Document {
    /// Parse PDF bytes. Eager work is just header + xref + trailer;
    /// object bodies, fonts, and CMaps materialize on first access.
    /// Tries the empty password — many encrypted-for-edit PDFs set an
    /// owner password but no user password.
    pub fn open(pdf_bytes: &[u8]) -> Result<Self, ExtractError> {
        Self::open_with_password(pdf_bytes, b"")
    }

    /// Parse a password-protected PDF. Supports every Standard Security
    /// Handler revision: RC4 40/128-bit, AES-128 CBC, AES-256 R=5/R=6.
    pub fn open_with_password(
        pdf_bytes: &[u8],
        password: &[u8],
    ) -> Result<Self, ExtractError> {
        let parse = open_sp_with_password(pdf_bytes, password)?;
        Ok(Self { parse })
    }

    /// Document-level summary.
    pub fn info(&self) -> DocumentInfo {
        extract_document_info_from_sp(&self.parse)
    }

    /// Number of pages in the document.
    pub fn page_count(&self) -> u32 {
        self.parse.get_pages().len() as u32
    }

    /// Per-page geometry (`mediabox`, `cropbox`, `rotation`, width, height).
    pub fn pages(&self, page: Option<u32>) -> Vec<PageInfo> {
        extract_page_info_from_sp(&self.parse, page)
    }

    /// Document outline (bookmarks / TOC).
    pub fn toc(&self) -> Result<Vec<TocEntry>, ExtractError> {
        extract_toc_from_sp(&self.parse)
    }

    /// Hyperlinks (filter to one page with `page = Some(n)`, 1-indexed).
    pub fn links(&self, page: Option<u32>) -> Vec<Link> {
        extract_links_from_sp(&self.parse, page)
    }

    /// Annotations (read-only).
    pub fn annotations(&self, page: Option<u32>) -> Vec<Annotation> {
        extract_annotations_from_sp(&self.parse, page)
    }

    /// Embedded-image inventory (no decoding).
    pub fn images(&self, page: Option<u32>) -> Vec<ImageInfo> {
        extract_images_from_sp(&self.parse, page)
    }

    /// Positional spans (page, text, font, font_size, bbox).
    pub fn text_positioned(
        &self,
        page: Option<u32>,
    ) -> Result<Vec<PositionedPage>, ExtractError> {
        extract_text_positioned_from_doc(&self.parse, page)
    }

    /// Words with bounding boxes.
    pub fn words(&self, page: Option<u32>) -> Result<Vec<Word>, ExtractError> {
        extract_words_from_doc(&self.parse, page)
    }

    /// Text blocks with bounding boxes.
    pub fn blocks(&self, page: Option<u32>) -> Result<Vec<TextBlock>, ExtractError> {
        extract_blocks_from_doc(&self.parse, page)
    }

    /// Block view used by the markdown renderer. Same per-span
    /// segmentation as [`Self::blocks`] but without the column pre-pass,
    /// so horizontally-spanning headings stay whole. Not re-exposed
    /// through `blocks()` — that surface is tuned for DocLayNet token
    /// recall and measures better with the column-aware version.
    pub fn blocks_for_markdown(
        &self,
        page: Option<u32>,
    ) -> Result<Vec<TextBlock>, ExtractError> {
        extract_blocks_streaming_from_doc(&self.parse, page)
    }

    /// Search for `query`; returns one [`SearchHit`] per match with the
    /// union-bbox of the contributing words.
    pub fn search(
        &self,
        query: &str,
        page: Option<u32>,
        options: Option<SearchOptions>,
    ) -> Result<Vec<SearchHit>, ExtractError> {
        let words = self.words(page)?;
        Ok(search_in_words_pub(&words, query, options.unwrap_or_default()))
    }

    /// Plain-text extraction. Same shape as [`crate::extract_text`].
    pub fn text(&self) -> Result<String, ExtractError> {
        let mut pages: Vec<u32> = self.parse.get_pages().keys().copied().collect();
        pages.sort();
        let mut chunks = Vec::with_capacity(pages.len());
        for n in pages {
            let s = self
                .parse
                .extract_text(&[n])
                .map_err(|e| ExtractError::PageExtractFailed {
                    page: n,
                    source: e.to_string(),
                })?;
            chunks.push(s);
        }
        Ok(chunks.join("\n\n").trim().to_string())
    }

    /// Lenient text extraction: failing pages return empty strings
    /// rather than erroring the whole document.
    pub fn text_lenient(&self) -> Result<String, ExtractError> {
        let mut pages: Vec<u32> = self.parse.get_pages().keys().copied().collect();
        pages.sort();
        let chunks: Vec<String> = pages
            .iter()
            .map(|n| self.parse.extract_text(&[*n]).unwrap_or_default())
            .collect();
        Ok(chunks.join("\n\n").trim().to_string())
    }

    /// Union of every page's media box.
    pub fn bbox(&self) -> Rect {
        self.pages(None)
            .iter()
            .map(|p| p.mediabox)
            .reduce(|a, b| a.union(b))
            .unwrap_or(Rect::ZERO)
    }

    /// `(page, font resource name)` → `(is_bold, is_italic)` for every
    /// font referenced from any page's `/Resources/Font`. Resolves
    /// `/BaseFont` + `/FontDescriptor/Flags` + `/FontWeight` and applies
    /// the PostScript naming convention. The markdown renderer keys
    /// `**bold**` / `*italic*` runs off this map.
    pub fn font_styles(&self) -> std::collections::HashMap<(u32, String), (bool, bool)> {
        let mut out: std::collections::HashMap<(u32, String), (bool, bool)> =
            std::collections::HashMap::new();
        let pages = self.parse.get_pages();
        for (page_num, page_id) in pages {
            let Ok(fonts) = self.parse.get_page_fonts(page_id) else {
                continue;
            };
            for (name_bytes, font_dict) in fonts.iter() {
                let resource_name = String::from_utf8_lossy(name_bytes).into_owned();
                let base_font = font_dict
                    .get_optional(b"BaseFont")
                    .and_then(|o| o.as_name().ok())
                    .map(|n| String::from_utf8_lossy(n).into_owned())
                    .unwrap_or_default();
                // Subset prefix: `ABCDEF+RealName` — strip if present.
                let stripped = base_font
                    .split_once('+')
                    .map(|(_, rest)| rest.to_string())
                    .unwrap_or(base_font.clone());
                let lower = stripped.to_lowercase();
                let mut bold = lower.contains("bold")
                    || lower.contains("black")
                    || lower.contains("heavy");
                let mut italic = lower.contains("italic") || lower.contains("oblique");
                // Walk /FontDescriptor if present.
                if let Some(desc_obj) = font_dict.get_optional(b"FontDescriptor") {
                    let desc_dict = match desc_obj {
                        spectre_parse::Object::Dictionary(d) => Some(d.clone()),
                        spectre_parse::Object::Reference(id) => {
                            self.parse.get_dictionary(*id).ok()
                        }
                        _ => None,
                    };
                    if let Some(desc) = desc_dict {
                        if let Some(flags) =
                            desc.get_optional(b"Flags").and_then(|o| o.as_i64().ok())
                        {
                            // PDF spec §9.8.2: bit 7 (mask 0x40) = italic.
                            if flags & 0x40 != 0 {
                                italic = true;
                            }
                        }
                        if let Some(weight) = desc
                            .get_optional(b"FontWeight")
                            .and_then(|o| o.as_i64().ok())
                        {
                            if weight >= 600 {
                                bold = true;
                            }
                        }
                    }
                }
                out.insert((page_num as u32, resource_name), (bold, italic));
            }
        }
        out
    }

    /// AcroForm fields (terminal nodes only — interior tree nodes are
    /// collapsed into the fully qualified name of each leaf). Empty
    /// `Vec` for PDFs without an `/AcroForm` entry.
    pub fn widgets(&self, page: Option<u32>) -> Vec<Widget> {
        let raw = self.parse.get_form_fields();
        let mapped: Vec<Widget> = raw.into_iter().map(Widget::from).collect();
        match page {
            Some(p) => mapped.into_iter().filter(|w| w.page == Some(p)).collect(),
            None => mapped,
        }
    }

 /// Decode one image to its canonical container bytes. `xref_id` is
 /// the `xref` field on [`ImageInfo`] returned by [`Self::images`].
 /// Most embedded images are JPEG (DCTDecode), which we pass through
 /// without re-encoding. Raw-pixel formats (FlateDecode + DeviceRGB /
 /// DeviceGray / DeviceCMYK) get wrapped in a minimal PNG so the bytes
 /// land on disk as a viewable file. JBIG2/CCITT raw bytes are
    /// surfaced for callers who want to wrap them themselves.
    pub fn image_bytes(&self, xref_id: u32) -> Option<DecodedImage> {
        let id = (xref_id, 0u16);
        self.parse.get_image_bytes(id).map(|(enc, bytes)| DecodedImage {
            container: enc.into(),
            bytes,
        })
    }
}

pub(crate) fn open_sp_with_password(
    pdf_bytes: &[u8],
    password: &[u8],
) -> Result<spectre_parse::Document, ExtractError> {
    spectre_parse::Document::open_with_password(pdf_bytes, password).map_err(|e| {
        let s = e.to_string();
        if s.contains("password did not authenticate") || s.contains("/Encrypt") {
            ExtractError::Encrypted
        } else {
            ExtractError::LoadFailed(s)
        }
    })
}

// ── Form-field surface ──────────────────────────────────────────────────────

/// One AcroForm field. `name` is the fully qualified, dot-separated path
/// from the field tree root. `value` is the current field value as a
/// PDF text string; for choice fields with multi-select, comma-joined.
#[derive(Debug, Clone)]
pub struct Widget {
    pub name: String,
    pub value: String,
    pub default_value: String,
    pub field_type: FormFieldType,
    pub flags: i64,
    pub rect: Option<Rect>,
    pub page: Option<u32>,
}

impl From<spectre_parse::FormField> for Widget {
    fn from(f: spectre_parse::FormField) -> Self {
        let rect = f.rect.map(|[x0, y0, x1, y1]| Rect { x0, y0, x1, y1 });
        Self {
            name: f.name,
            value: f.value,
            default_value: f.default_value,
            field_type: f.field_type.into(),
            flags: f.flags,
            rect,
            page: f.page,
        }
    }
}

/// AcroForm field types (PDF spec §12.7.4 `/FT`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FormFieldType {
    Text,
    Button,
    Choice,
    Signature,
    Unknown,
}

impl FormFieldType {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Text => "text",
            Self::Button => "button",
            Self::Choice => "choice",
            Self::Signature => "signature",
            Self::Unknown => "unknown",
        }
    }
}

impl From<spectre_parse::FormFieldType> for FormFieldType {
    fn from(t: spectre_parse::FormFieldType) -> Self {
        use spectre_parse::FormFieldType as P;
        match t {
            P::Text => Self::Text,
            P::Button => Self::Button,
            P::Choice => Self::Choice,
            P::Signature => Self::Signature,
            P::Unknown => Self::Unknown,
        }
    }
}

/// Re-export of [`spectre_parse::FormField`] for tests/internal callers.
pub use spectre_parse::FormField;

// ── Image-bytes surface ────────────────────────────────────────────────────

/// One image's decoded byte buffer.
#[derive(Debug, Clone)]
pub struct DecodedImage {
    pub container: ImageContainer,
    pub bytes: Vec<u8>,
}

/// Output container for [`Document::image_bytes`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageContainer {
    Jpeg,
    Jpeg2000,
    Png,
    Jbig2,
    Ccitt,
}

impl ImageContainer {
    pub fn extension(self) -> &'static str {
        match self {
            Self::Jpeg => "jpg",
            Self::Jpeg2000 => "jp2",
            Self::Png => "png",
            Self::Jbig2 => "jb2",
            Self::Ccitt => "tiff",
        }
    }
}

impl From<spectre_parse::ImageEncoding> for ImageContainer {
    fn from(e: spectre_parse::ImageEncoding) -> Self {
        use spectre_parse::ImageEncoding as E;
        match e {
            E::Jpeg => Self::Jpeg,
            E::Jpeg2000 => Self::Jpeg2000,
            E::Png => Self::Png,
            E::Jbig2 => Self::Jbig2,
            E::Ccitt => Self::Ccitt,
        }
    }
}