use crate::types::internal::{ElementKind, InternalDocument};
use crate::types::{HierarchicalBlock, PageContent, PageHierarchy};
pub(crate) fn assign_hierarchy_to_pages(pages: &mut [PageContent], doc: &InternalDocument) {
let mut page_hierarchies: std::collections::HashMap<u32, Vec<HierarchicalBlock>> = std::collections::HashMap::new();
for element in &doc.elements {
let page_num = match element.page {
Some(p) => p,
None => continue,
};
match element.kind {
ElementKind::Heading { level } => {
let block = HierarchicalBlock {
text: element.text.clone(),
font_size: 12.0,
level: format!("h{}", level),
bbox: element
.bbox
.map(|b| (b.x0 as f32, b.y0 as f32, b.x1 as f32, b.y1 as f32)),
};
page_hierarchies.entry(page_num).or_default().push(block);
}
ElementKind::Paragraph => {
let block = HierarchicalBlock {
text: element.text.clone(),
font_size: 12.0,
level: "body".to_string(),
bbox: element
.bbox
.map(|b| (b.x0 as f32, b.y0 as f32, b.x1 as f32, b.y1 as f32)),
};
page_hierarchies.entry(page_num).or_default().push(block);
}
_ => {}
}
}
for page in pages.iter_mut() {
if let Some(blocks) = page_hierarchies.remove(&page.page_number) {
let block_count = blocks.len() as u32;
page.hierarchy = Some(PageHierarchy { block_count, blocks });
}
}
}
pub(crate) fn assign_tables_and_images_to_pages(
mut page_contents: Option<Vec<PageContent>>,
tables: &[crate::types::Table],
images: &[crate::types::ExtractedImage],
) -> Option<Vec<PageContent>> {
let pages = page_contents.take()?;
let mut updated_pages = pages;
for table in tables {
if let Some(page) = updated_pages.iter_mut().find(|p| p.page_number == table.page_number) {
page.tables.push(std::sync::Arc::new(table.clone()));
}
}
for (idx, image) in images.iter().enumerate() {
if let Some(page_num) = image.page_number
&& let Some(page) = updated_pages.iter_mut().find(|p| p.page_number == page_num)
{
page.image_indices.push(idx as u32);
}
}
for page in &mut updated_pages {
if !page.tables.is_empty() || !page.image_indices.is_empty() {
page.is_blank = Some(false);
}
}
Some(updated_pages)
}