paperforge-pdf 0.1.0

//! Document-level editing: merging documents and extracting page ranges.
//!
//! Both operations work at the object level: they renumber objects so indirect
//! references stay consistent, then rebuild the page tree. No stream decoding
//! is required — content, resources, fonts, and images are copied verbatim.

use std::collections::BTreeSet;

use crate::error::{PdfError, PdfResult};
use crate::object::*;
use crate::parser::Document;

/// Returns the indirect ids of every page (`/Type /Page`) in `doc`, in document
/// order, by walking the catalog's page tree.
pub fn page_ids(doc: &Document) -> Vec<ObjectId> {
    let mut pages = Vec::new();
    let Some(catalog) = doc.catalog() else {
        return pages;
    };
    let Some(PdfObject::Dictionary(catalog_dict)) = doc.get_object(catalog) else {
        return pages;
    };
    let Some(PdfObject::Reference(root)) = catalog_dict.get("Pages") else {
        return pages;
    };

    let mut stack = vec![*root];
    while let Some(id) = stack.pop() {
        let Some(PdfObject::Dictionary(dict)) = doc.get_object(id) else {
            continue;
        };
        match dict.get_name("Type").map(|n| n.as_str()) {
            Some("Pages") => {
                if let Some(PdfObject::Array(kids)) = dict.get("Kids") {
                    for kid in kids.0.iter().rev() {
                        if let PdfObject::Reference(r) = kid {
                            stack.push(*r);
                        }
                    }
                }
            }
            Some("Page") => pages.push(id),
            _ => {}
        }
    }
    pages
}

/// Merges `other` after `base`: every page of `other` is appended to `base`'s
/// page tree, and all objects reachable from those pages (content streams,
/// resources, fonts, images) are copied into the result with renumbered ids.
pub fn merge(base: &Document, other: &Document) -> PdfResult<Document> {
    let mut out = Document::new();
    for (id, obj) in base.objects() {
        out.add_object(*id, obj.clone());
    }
    if let Some(c) = base.catalog() {
        out.set_catalog(c);
    }
    if let Some(i) = base.info() {
        out.set_info(i);
    }

    let other_pages = page_ids(other);
    if other_pages.is_empty() {
        return Ok(out);
    }

    // Copy everything reachable from the incoming pages, skipping pages-tree
    // nodes (base has its own tree). References are shifted by `offset`.
    let reachable = collect_reachable(other, &other_pages);
    let offset = base
        .objects()
        .keys()
        .map(|id| id.number)
        .max()
        .unwrap_or(0)
        .saturating_add(1);

    for id in &reachable {
        let Some(obj) = other.get_object(*id) else {
            continue;
        };
        // The other doc's own /Pages nodes stay behind; pages get reparented
        // to base's tree below.
        if is_pages_node(obj) {
            continue;
        }
        out.add_object(
            ObjectId::new(id.number + offset, id.generation),
            remap_object(obj, offset),
        );
    }

    // Reparent the incoming pages under base's pages root.
    let base_pages_root = pages_root(&out)
        .ok_or_else(|| PdfError::InvalidObject("base document has no page tree".into()))?;
    for pid in &other_pages {
        let new_id = ObjectId::new(pid.number + offset, pid.generation);
        if let Some(PdfObject::Dictionary(dict)) = out.get_object_mut(new_id) {
            dict.insert("Parent", PdfObject::Reference(base_pages_root));
        }
        append_kid(&mut out, base_pages_root, new_id);
    }

    Ok(out)
}

/// Extracts pages `first..=last` (1-based, inclusive) into a new document. All
/// objects are copied; only the page tree is rewritten so the selected pages
/// form the new document's pages. Orphaned objects are harmless and may be
/// dropped by a future optimizer.
pub fn extract_pages(doc: &Document, first: usize, last: usize) -> PdfResult<Document> {
    let pages = page_ids(doc);
    if pages.is_empty() {
        return Err(PdfError::InvalidPage("document has no pages".into()));
    }
    let first = first.max(1);
    let last = last.min(pages.len());
    if first > last {
        return Err(PdfError::InvalidPage(format!(
            "page range {first}-{last} is empty (document has {} pages)",
            pages.len()
        )));
    }
    let selected: Vec<ObjectId> = pages[first - 1..last].to_vec();

    let mut out = Document::new();
    for (id, obj) in doc.objects() {
        out.add_object(*id, obj.clone());
    }
    if let Some(i) = doc.info() {
        out.set_info(i);
    }

    // Fresh flat page tree: one /Pages node holding the selected pages.
    let max = doc.objects().keys().map(|id| id.number).max().unwrap_or(0);
    let new_pages = ObjectId::new(max + 1, 0);
    let new_catalog = ObjectId::new(max + 2, 0);

    let mut kids = PdfArray::new();
    for pid in &selected {
        kids.push(PdfObject::Reference(*pid));
        if let Some(PdfObject::Dictionary(dict)) = out.get_object_mut(*pid) {
            dict.insert("Parent", PdfObject::Reference(new_pages));
        }
    }
    let mut pages_dict = PdfDictionary::new();
    pages_dict.insert("Type", PdfObject::Name(PdfName::new("Pages")));
    pages_dict.insert("Count", PdfObject::Integer(selected.len() as i64));
    pages_dict.insert("Kids", PdfObject::Array(kids));
    out.add_object(new_pages, PdfObject::Dictionary(pages_dict));

    let mut catalog = PdfDictionary::new();
    catalog.insert("Type", PdfObject::Name(PdfName::new("Catalog")));
    catalog.insert("Pages", PdfObject::Reference(new_pages));
    out.add_object(new_catalog, PdfObject::Dictionary(catalog));
    out.set_catalog(new_catalog);

    Ok(out)
}

fn append_kid(doc: &mut Document, pages_id: ObjectId, kid: ObjectId) {
    if let Some(PdfObject::Dictionary(dict)) = doc.get_object_mut(pages_id) {
        let mut kids = dict
            .get_array("Kids")
            .cloned()
            .unwrap_or_else(PdfArray::new);
        kids.push(PdfObject::Reference(kid));
        dict.insert("Kids", PdfObject::Array(kids));
        let count = dict.get_integer("Count").unwrap_or(0) + 1;
        dict.insert("Count", PdfObject::Integer(count));
    }
}

fn pages_root(doc: &Document) -> Option<ObjectId> {
    let catalog = doc.catalog()?;
    let PdfObject::Dictionary(dict) = doc.get_object(catalog)? else {
        return None;
    };
    dict.get("Pages").and_then(|r| r.as_reference())
}

fn is_pages_node(obj: &PdfObject) -> bool {
    obj.as_dict()
        .and_then(|d| d.get_name("Type"))
        .map(|n| n.as_str() == "Pages")
        .unwrap_or(false)
}

/// Breadth-first closure of all indirect references reachable from `roots`.
fn collect_reachable(doc: &Document, roots: &[ObjectId]) -> BTreeSet<ObjectId> {
    let mut seen = BTreeSet::new();
    let mut stack: Vec<ObjectId> = roots.to_vec();
    while let Some(id) = stack.pop() {
        if !seen.insert(id) {
            continue;
        }
        if let Some(obj) = doc.get_object(id) {
            let mut refs = Vec::new();
            collect_refs(obj, &mut refs);
            stack.extend(refs);
        }
    }
    seen
}

fn collect_refs(obj: &PdfObject, out: &mut Vec<ObjectId>) {
    match obj {
        PdfObject::Reference(id) => out.push(*id),
        PdfObject::Array(arr) => {
            for item in &arr.0 {
                collect_refs(item, out);
            }
        }
        PdfObject::Dictionary(dict) => {
            for (_, value) in dict.iter() {
                collect_refs(value, out);
            }
        }
        PdfObject::Stream(stream) => {
            let dict = PdfObject::Dictionary(stream.dictionary.clone());
            collect_refs(&dict, out);
        }
        _ => {}
    }
}

/// Deep-clones `obj`, shifting every indirect reference's object number by
/// `offset` so copied objects stay self-consistent in the target document.
fn remap_object(obj: &PdfObject, offset: u32) -> PdfObject {
    match obj {
        PdfObject::Reference(id) => {
            PdfObject::Reference(ObjectId::new(id.number + offset, id.generation))
        }
        PdfObject::Array(arr) => PdfObject::Array(PdfArray(
            arr.0
                .iter()
                .map(|item| remap_object(item, offset))
                .collect(),
        )),
        PdfObject::Dictionary(dict) => {
            let mut out = PdfDictionary::new();
            for (key, value) in dict.iter() {
                out.insert(&key.0, remap_object(value, offset));
            }
            PdfObject::Dictionary(out)
        }
        PdfObject::Stream(stream) => {
            let dict = match remap_object(&PdfObject::Dictionary(stream.dictionary.clone()), offset)
            {
                PdfObject::Dictionary(d) => d,
                _ => unreachable!(),
            };
            PdfObject::Stream(PdfStream::with_dict(dict, stream.data.clone()))
        }
        other => other.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Parser, Serializer};

    fn two_page_doc() -> Document {
        // catalog(1) -> pages(2) -> page(3), page(4), each with a content stream.
        let mut doc = Document::new();
        doc.set_catalog(ObjectId::new(1, 0));
        doc.add_object(
            ObjectId::new(1, 0),
            PdfObject::Dictionary({
                let mut d = PdfDictionary::new();
                d.insert("Type", PdfObject::Name(PdfName::new("Catalog")));
                d.insert("Pages", PdfObject::Reference(ObjectId::new(2, 0)));
                d
            }),
        );
        doc.add_object(
            ObjectId::new(2, 0),
            PdfObject::Dictionary({
                let mut d = PdfDictionary::new();
                d.insert("Type", PdfObject::Name(PdfName::new("Pages")));
                d.insert("Count", PdfObject::Integer(2));
                let mut kids = PdfArray::new();
                kids.push(PdfObject::Reference(ObjectId::new(3, 0)));
                kids.push(PdfObject::Reference(ObjectId::new(4, 0)));
                d.insert("Kids", PdfObject::Array(kids));
                d
            }),
        );
        for (n, text) in [(3u32, b"BT (one) Tj ET".as_slice()), (4, b"BT (two) Tj ET")] {
            doc.add_object(
                ObjectId::new(n, 0),
                PdfObject::Dictionary({
                    let mut d = PdfDictionary::new();
                    d.insert("Type", PdfObject::Name(PdfName::new("Page")));
                    d.insert("Parent", PdfObject::Reference(ObjectId::new(2, 0)));
                    let mut media = PdfArray::new();
                    media.push(PdfObject::Integer(0));
                    media.push(PdfObject::Integer(0));
                    media.push(PdfObject::Integer(595));
                    media.push(PdfObject::Integer(842));
                    d.insert("MediaBox", PdfObject::Array(media));
                    d.insert("Contents", PdfObject::Reference(ObjectId::new(n + 10, 0)));
                    d
                }),
            );
            doc.add_object(
                ObjectId::new(n + 10, 0),
                PdfObject::Stream(PdfStream::with_dict(PdfDictionary::new(), text.to_vec())),
            );
        }
        doc
    }

    fn serialize(doc: &Document) -> Vec<u8> {
        let mut buf = std::io::Cursor::new(Vec::new());
        Serializer::new().serialize(doc, &mut buf).unwrap();
        buf.into_inner()
    }

    #[test]
    fn page_ids_walks_the_tree_in_order() {
        let doc = two_page_doc();
        let ids = page_ids(&doc);
        assert_eq!(ids, vec![ObjectId::new(3, 0), ObjectId::new(4, 0)]);
    }

    #[test]
    fn merge_appends_pages_and_remaps_references() {
        let a = two_page_doc();
        let b = two_page_doc();
        let merged = merge(&a, &b).unwrap();
        assert_eq!(page_ids(&merged).len(), 4);

        // Both content streams survive; re-parsing proves the file is valid.
        let bytes = serialize(&merged);
        let reparsed = Parser::new().parse(&bytes).unwrap();
        assert_eq!(page_ids(&reparsed).len(), 4);

        // The second document's content stream was renumbered past base's max.
        let mut saw_two = false;
        for (_, obj) in reparsed.objects() {
            if let PdfObject::Stream(s) = obj {
                if s.data == b"BT (two) Tj ET" {
                    saw_two = true;
                }
            }
        }
        assert!(saw_two, "second doc's content must be present");
    }

    #[test]
    fn extract_pages_picks_a_subrange() {
        let doc = two_page_doc();
        let sub = extract_pages(&doc, 2, 2).unwrap();
        let ids = page_ids(&sub);
        assert_eq!(ids, vec![ObjectId::new(4, 0)]);

        let bytes = serialize(&sub);
        let reparsed = Parser::new().parse(&bytes).unwrap();
        assert_eq!(page_ids(&reparsed), vec![ObjectId::new(4, 0)]);
    }

    #[test]
    fn extract_pages_rejects_bad_ranges() {
        let doc = two_page_doc();
        assert!(extract_pages(&doc, 3, 2).is_err());
        assert!(extract_pages(&doc, 1, 99).is_ok()); // clamped, not an error
    }

    #[test]
    fn merge_with_empty_second_doc_is_identity() {
        let a = two_page_doc();
        let mut b = Document::new();
        b.set_catalog(ObjectId::new(1, 0));
        b.add_object(
            ObjectId::new(1, 0),
            PdfObject::Dictionary({
                let mut d = PdfDictionary::new();
                d.insert("Type", PdfObject::Name(PdfName::new("Catalog")));
                d
            }),
        );
        let merged = merge(&a, &b).unwrap();
        assert_eq!(page_ids(&merged).len(), 2);
    }
}