use folio_core::{FolioError, Rect, Result};
use folio_cos::{CosDoc, ObjectId, PdfObject};
use indexmap::IndexMap;
use crate::info::DocInfo;
use crate::page::Page;
pub struct PdfDoc {
cos: CosDoc,
}
impl PdfDoc {
pub fn new() -> Result<Self> {
let mut cos = CosDoc::new();
let mut pages_dict = IndexMap::new();
pages_dict.insert(b"Type".to_vec(), PdfObject::Name(b"Pages".to_vec()));
pages_dict.insert(b"Kids".to_vec(), PdfObject::Array(vec![]));
pages_dict.insert(b"Count".to_vec(), PdfObject::Integer(0));
let pages_id = cos.create_indirect(PdfObject::Dict(pages_dict));
let mut catalog_dict = IndexMap::new();
catalog_dict.insert(b"Type".to_vec(), PdfObject::Name(b"Catalog".to_vec()));
catalog_dict.insert(b"Pages".to_vec(), PdfObject::Reference(pages_id));
let catalog_id = cos.create_indirect(PdfObject::Dict(catalog_dict));
cos.trailer_mut()
.insert(b"Root".to_vec(), PdfObject::Reference(catalog_id));
Ok(Self { cos })
}
pub fn open(path: &str) -> Result<Self> {
let cos = CosDoc::open_file(path)?;
Ok(Self { cos })
}
pub fn open_from_bytes(data: Vec<u8>) -> Result<Self> {
let cos = CosDoc::open(data)?;
Ok(Self { cos })
}
pub fn cos(&self) -> &CosDoc {
&self.cos
}
pub fn cos_mut(&mut self) -> &mut CosDoc {
&mut self.cos
}
fn catalog_ref(&self) -> Result<ObjectId> {
self.cos
.trailer()
.get(b"Root".as_slice())
.and_then(|o| o.as_reference())
.ok_or_else(|| FolioError::InvalidObject("Missing /Root in trailer".into()))
}
fn pages_ref(&mut self) -> Result<ObjectId> {
let catalog_ref = self.catalog_ref()?;
let catalog = self
.cos
.get_object(catalog_ref.num)?
.ok_or_else(|| FolioError::InvalidObject("Catalog not found".into()))?
.clone();
catalog
.dict_get(b"Pages")
.and_then(|o| o.as_reference())
.ok_or_else(|| FolioError::InvalidObject("Missing /Pages in catalog".into()))
}
pub fn page_count(&mut self) -> Result<u32> {
let pages_ref = self.pages_ref()?;
let pages = self
.cos
.get_object(pages_ref.num)?
.ok_or_else(|| FolioError::InvalidObject("Pages dict not found".into()))?
.clone();
pages
.dict_get_i64(b"Count")
.map(|c| c as u32)
.ok_or_else(|| FolioError::InvalidObject("Missing /Count in Pages".into()))
}
pub fn get_page(&mut self, page_num: u32) -> Result<Page> {
if page_num == 0 {
return Err(FolioError::InvalidArgument(
"Page numbers are 1-based".into(),
));
}
let page_refs = self.collect_page_refs()?;
let index = (page_num - 1) as usize;
if index >= page_refs.len() {
return Err(FolioError::InvalidArgument(format!(
"Page {} out of range (document has {} pages)",
page_num,
page_refs.len()
)));
}
let page_ref = page_refs[index];
let page_obj = self
.cos
.get_object(page_ref.num)?
.ok_or_else(|| {
FolioError::InvalidObject(format!("Page object {} not found", page_ref.num))
})?
.clone();
let page_obj = self.resolve_inherited_attrs(page_obj)?;
Ok(Page::new(page_ref, page_obj, page_num))
}
fn resolve_inherited_attrs(&mut self, page_obj: PdfObject) -> Result<PdfObject> {
const INHERITABLE: &[&[u8]] = &[b"MediaBox", b"CropBox", b"Rotate", b"Resources"];
let mut dict = match page_obj.as_dict() {
Some(d) => d.clone(),
None => return Ok(page_obj),
};
let missing: Vec<&[u8]> = INHERITABLE
.iter()
.filter(|&&key| !dict.contains_key(key))
.copied()
.collect();
if missing.is_empty() {
return Ok(page_obj);
}
let mut parent_ref = dict
.get(b"Parent".as_slice())
.and_then(|o| o.as_reference());
let mut visited = std::collections::HashSet::new();
while let Some(pref) = parent_ref {
if visited.contains(&pref.num) {
break; }
visited.insert(pref.num);
let parent = match self.cos.get_object(pref.num)? {
Some(obj) => obj.clone(),
None => break,
};
let parent_dict = match parent.as_dict() {
Some(d) => d,
None => break,
};
for &key in &missing {
if !dict.contains_key(key) {
if let Some(value) = parent_dict.get(key) {
dict.insert(key.to_vec(), value.clone());
}
}
}
if INHERITABLE.iter().all(|&key| dict.contains_key(key)) {
break;
}
parent_ref = parent_dict
.get(b"Parent".as_slice())
.and_then(|o| o.as_reference());
}
Ok(PdfObject::Dict(dict))
}
fn collect_page_refs(&mut self) -> Result<Vec<ObjectId>> {
let pages_ref = self.pages_ref()?;
let mut result = Vec::new();
self.collect_pages_recursive(pages_ref, &mut result)?;
Ok(result)
}
fn collect_pages_recursive(
&mut self,
node_ref: ObjectId,
result: &mut Vec<ObjectId>,
) -> Result<()> {
let node = self
.cos
.get_object(node_ref.num)?
.ok_or_else(|| {
FolioError::InvalidObject(format!("Page tree node {} not found", node_ref.num))
})?
.clone();
let type_name = node.dict_get_name(b"Type").unwrap_or(b"");
match type_name {
b"Pages" => {
if let Some(kids) = node.dict_get(b"Kids").and_then(|o| o.as_array()) {
for kid in kids {
if let Some(kid_ref) = kid.as_reference() {
self.collect_pages_recursive(kid_ref, result)?;
}
}
}
}
b"Page" | _ => {
result.push(node_ref);
}
}
Ok(())
}
pub fn create_page(&mut self, media_box: Rect) -> Result<u32> {
let pages_ref = self.pages_ref()?;
let mut page_dict = IndexMap::new();
page_dict.insert(b"Type".to_vec(), PdfObject::Name(b"Page".to_vec()));
page_dict.insert(b"Parent".to_vec(), PdfObject::Reference(pages_ref));
page_dict.insert(
b"MediaBox".to_vec(),
PdfObject::Array(vec![
PdfObject::Real(media_box.x1),
PdfObject::Real(media_box.y1),
PdfObject::Real(media_box.x2),
PdfObject::Real(media_box.y2),
]),
);
let page_id = self.cos.create_indirect(PdfObject::Dict(page_dict));
let pages = self
.cos
.get_object(pages_ref.num)?
.ok_or_else(|| FolioError::InvalidObject("Pages not found".into()))?
.clone();
let mut pages_dict = pages.as_dict().cloned().unwrap_or_default();
let mut kids = pages_dict
.get(b"Kids".as_slice())
.and_then(|o| o.as_array())
.map(|a| a.to_vec())
.unwrap_or_default();
kids.push(PdfObject::Reference(page_id));
pages_dict.insert(b"Kids".to_vec(), PdfObject::Array(kids.clone()));
pages_dict.insert(b"Count".to_vec(), PdfObject::Integer(kids.len() as i64));
self.cos
.update_object(pages_ref.num, PdfObject::Dict(pages_dict));
Ok(kids.len() as u32)
}
pub fn doc_info(&mut self) -> Result<DocInfo> {
let info_ref = self
.cos
.trailer()
.get(b"Info".as_slice())
.and_then(|o| o.as_reference());
match info_ref {
Some(id) => {
let obj = self
.cos
.get_object(id.num)?
.cloned()
.unwrap_or(PdfObject::Null);
Ok(DocInfo::from_dict(
obj.as_dict().cloned().unwrap_or_default(),
))
}
None => Ok(DocInfo::default()),
}
}
pub fn is_modified(&self) -> bool {
self.cos.is_modified()
}
pub fn save_to_bytes(&mut self) -> Result<Vec<u8>> {
self.cos.save_to_bytes()
}
pub fn save(&mut self, path: &str) -> Result<()> {
self.cos.save_to_file(path)
}
}