1use std::collections::BTreeMap;
2use std::path::Path;
3
4use lopdf::{Dictionary, Document, Object, ObjectId};
5
6use crate::error::{PdfConvertError, Result};
7pub use crate::models::{Bookmark, ChunkMetadata, PdfInfo};
8
9impl PdfInfo {
10 pub fn load(pdf_path: &Path) -> Result<Self> {
11 let doc = Document::load(pdf_path)
12 .map_err(|e| PdfConvertError::parse_error("PDF document", e.to_string()))?;
13 Ok(Self::from_document(doc))
14 }
15
16 pub fn load_from_bytes(data: &[u8]) -> Result<Self> {
17 let doc = Document::load_mem(data)
18 .map_err(|e| PdfConvertError::parse_error("PDF document", e.to_string()))?;
19 Ok(Self::from_document(doc))
20 }
21
22 fn from_document(doc: Document) -> Self {
23 let total_pages = doc.get_pages().len() as u32;
24 let outlines = Self::extract_outlines(&doc);
25
26 PdfInfo {
27 total_pages,
28 outlines,
29 }
30 }
31
32 fn extract_outlines(doc: &Document) -> Vec<Bookmark> {
33 let mut outlines = Vec::new();
34 if let Ok(catalog) = doc.catalog() {
35 if let Ok(Object::Reference(outlines_ref)) = catalog.get(b"Outlines") {
36 if let Ok(outlines_dict) =
37 doc.get_object(*outlines_ref).and_then(|obj| obj.as_dict())
38 {
39 if let Ok(Object::Reference(first_ref)) = outlines_dict.get(b"First") {
40 let page_map = Self::build_page_map(&doc);
41 Self::traverse_outlines(&doc, *first_ref, 0, &page_map, &mut outlines);
42 }
43 }
44 }
45 }
46
47 outlines
48 }
49
50 fn build_page_map(doc: &Document) -> BTreeMap<ObjectId, u32> {
51 let mut page_map = BTreeMap::new();
52 for (idx, &page_id) in doc.get_pages().values().enumerate() {
53 page_map.insert(page_id, (idx + 1) as u32);
54 }
55 page_map
56 }
57
58 fn traverse_outlines(
59 doc: &Document,
60 current_ref: ObjectId,
61 level: usize,
62 page_map: &BTreeMap<ObjectId, u32>,
63 outlines: &mut Vec<Bookmark>,
64 ) {
65 let mut current_obj_ref = Some(current_ref);
66 while let Some(obj_ref) = current_obj_ref {
67 if let Ok(dict) = doc.get_object(obj_ref).and_then(|obj| obj.as_dict()) {
68 if let Some(bookmark) = Self::parse_bookmark(doc, dict, level, page_map) {
69 outlines.push(bookmark);
70 }
71
72 if let Ok(Object::Reference(first_child)) = dict.get(b"First") {
73 Self::traverse_outlines(doc, *first_child, level + 1, page_map, outlines);
74 }
75
76 current_obj_ref = match dict.get(b"Next") {
77 Ok(Object::Reference(next_ref)) => Some(*next_ref),
78 _ => None,
79 };
80 } else {
81 break;
82 }
83 }
84 }
85
86 fn parse_bookmark(
87 doc: &Document,
88 dict: &Dictionary,
89 level: usize,
90 page_map: &BTreeMap<ObjectId, u32>,
91 ) -> Option<Bookmark> {
92 let title = match dict.get(b"Title") {
93 Ok(Object::String(s, _)) => String::from_utf8_lossy(s).into_owned(),
94 _ => return None,
95 };
96
97 let page_number = if let Ok(dest) = dict.get(b"Dest") {
98 Self::resolve_dest(doc, dest, page_map)
99 } else if let Ok(Object::Dictionary(action)) = dict.get(b"A") {
100 if let Ok(Object::Name(name)) = action.get(b"S") {
101 if name == b"GoTo" {
102 if let Ok(dest) = action.get(b"D") {
103 Self::resolve_dest(doc, dest, page_map)
104 } else {
105 None
106 }
107 } else {
108 None
109 }
110 } else {
111 None
112 }
113 } else {
114 None
115 };
116
117 page_number.map(|pn| Bookmark {
118 title,
119 page_number: pn,
120 level,
121 })
122 }
123
124 fn resolve_dest(
125 doc: &Document,
126 dest: &Object,
127 page_map: &BTreeMap<ObjectId, u32>,
128 ) -> Option<u32> {
129 let dest = match dest {
130 Object::Reference(r) => doc.get_object(*r).ok()?,
131 _ => dest,
132 };
133
134 match dest {
135 Object::Array(arr) => {
136 if let Some(Object::Reference(page_ref)) = arr.get(0) {
137 page_map.get(page_ref).copied()
138 } else if let Some(Object::Integer(page_idx)) = arr.get(0) {
139 Some((*page_idx as u32) + 1)
140 } else {
141 None
142 }
143 }
144 Object::Name(_) => None,
147 _ => None,
148 }
149 }
150
151 pub fn calculate_hybrid_ranges(&self, pages_per_chunk: u32) -> Vec<ChunkMetadata> {
152 if self.outlines.is_empty() {
153 return self.calculate_page_ranges(pages_per_chunk);
154 }
155
156 let mut boundary_points: Vec<(u32, Option<usize>, Option<String>)> = self
157 .outlines
158 .iter()
159 .enumerate()
160 .map(|(idx, b)| (b.page_number, Some(idx + 1), Some(b.title.clone())))
161 .collect();
162
163 if !boundary_points.iter().any(|(p, _, _)| *p == 1) {
164 boundary_points.push((1, None, None));
165 }
166 boundary_points.push((self.total_pages + 1, None, None));
167 boundary_points.sort_by_key(|(p, _, _)| *p);
168
169 let mut ranges = Vec::new();
170 for i in 0..boundary_points.len() - 1 {
171 let (start, b_idx, b_title) = &boundary_points[i];
172 let (next_boundary, _, _) = &boundary_points[i + 1];
173
174 let mut current = *start;
175 while current < *next_boundary {
176 let end = std::cmp::min(current + pages_per_chunk - 1, *next_boundary - 1);
177 ranges.push(ChunkMetadata {
178 start_page: current,
179 end_page: end,
180 bookmark_index: b_idx.clone(),
181 bookmark_title: b_title.clone(),
182 });
183 current = end + 1;
184 }
185 }
186
187 log::debug!("Calculated hybrid page ranges: {:?}", ranges);
188 ranges
189 }
190
191 #[cfg_attr(not(test), allow(dead_code))]
192 pub fn calculate_bookmark_ranges(&self) -> Vec<ChunkMetadata> {
193 if self.outlines.is_empty() {
194 return vec![ChunkMetadata {
195 start_page: 1,
196 end_page: self.total_pages,
197 bookmark_index: None,
198 bookmark_title: None,
199 }];
200 }
201
202 let mut points: Vec<(u32, usize, String)> = self
203 .outlines
204 .iter()
205 .enumerate()
206 .map(|(idx, b)| (b.page_number, idx + 1, b.title.clone()))
207 .collect();
208 points.sort_by_key(|p| p.0);
209
210 let mut ranges = Vec::new();
211 for i in 0..points.len() {
212 let (start, b_idx, b_title) = &points[i];
213 let end = if i + 1 < points.len() {
214 points[i + 1].0 - 1
215 } else {
216 self.total_pages
217 };
218
219 if *start <= end {
220 ranges.push(ChunkMetadata {
221 start_page: *start,
222 end_page: end,
223 bookmark_index: Some(*b_idx),
224 bookmark_title: Some(b_title.clone()),
225 });
226 }
227 }
228
229 if ranges.is_empty() {
230 ranges.push(ChunkMetadata {
231 start_page: 1,
232 end_page: self.total_pages,
233 bookmark_index: None,
234 bookmark_title: None,
235 });
236 }
237
238 ranges
239 }
240
241 pub fn calculate_page_ranges(&self, pages_per_chunk: u32) -> Vec<ChunkMetadata> {
242 let mut ranges = Vec::new();
243
244 for start in (0..self.total_pages).step_by(pages_per_chunk as usize) {
245 let actual_start = start + 1;
246 let actual_end = std::cmp::min(start + pages_per_chunk, self.total_pages);
247 ranges.push(ChunkMetadata {
248 start_page: actual_start,
249 end_page: actual_end,
250 bookmark_index: None,
251 bookmark_title: None,
252 });
253 }
254
255 log::debug!("Calculated page ranges: {:?}", ranges);
256
257 ranges
258 }
259}
260
261#[allow(dead_code)]
262pub fn extract_filename(path: &Path) -> Option<String> {
263 path.file_name()
264 .and_then(|name| name.to_str())
265 .map(|s| s.to_string())
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 #[test]
273 fn test_calculate_page_ranges_even_division() {
274 let ranges = calculate_ranges_helper(10, 5);
275 assert_eq!(ranges, vec![(1, 5), (6, 10)]);
276 }
277
278 #[test]
279 fn test_calculate_page_ranges_uneven_division() {
280 let ranges = calculate_ranges_helper(12, 5);
281 assert_eq!(ranges, vec![(1, 5), (6, 10), (11, 12)]);
282 }
283
284 #[test]
285 fn test_calculate_page_ranges_single_page() {
286 let ranges = calculate_ranges_helper(1, 5);
287 assert_eq!(ranges, vec![(1, 1)]);
288 }
289
290 #[test]
291 fn test_calculate_page_ranges_chunk_larger_than_total() {
292 let ranges = calculate_ranges_helper(3, 10);
293 assert_eq!(ranges, vec![(1, 3)]);
294 }
295
296 #[test]
297 fn test_calculate_bookmark_ranges_empty() {
298 let info = PdfInfo {
299 total_pages: 10,
300 outlines: vec![],
301 };
302 let ranges = info.calculate_bookmark_ranges();
303 assert_eq!(
304 ranges,
305 vec![ChunkMetadata {
306 start_page: 1,
307 end_page: 10,
308 bookmark_index: None,
309 bookmark_title: None
310 }]
311 );
312 }
313
314 #[test]
315 fn test_calculate_bookmark_ranges_simple() {
316 let info = PdfInfo {
317 total_pages: 10,
318 outlines: vec![
319 Bookmark {
320 title: "C1".into(),
321 page_number: 1,
322 level: 0,
323 },
324 Bookmark {
325 title: "C2".into(),
326 page_number: 5,
327 level: 0,
328 },
329 ],
330 };
331 let ranges = info.calculate_bookmark_ranges();
332 assert_eq!(
333 ranges,
334 vec![
335 ChunkMetadata {
336 start_page: 1,
337 end_page: 4,
338 bookmark_index: Some(1),
339 bookmark_title: Some("C1".into())
340 },
341 ChunkMetadata {
342 start_page: 5,
343 end_page: 10,
344 bookmark_index: Some(2),
345 bookmark_title: Some("C2".into())
346 },
347 ]
348 );
349 }
350
351 #[test]
352 fn test_calculate_bookmark_ranges_nested() {
353 let info = PdfInfo {
354 total_pages: 10,
355 outlines: vec![
356 Bookmark {
357 title: "C1".into(),
358 page_number: 1,
359 level: 0,
360 },
361 Bookmark {
362 title: "C1.1".into(),
363 page_number: 2,
364 level: 1,
365 },
366 Bookmark {
367 title: "C2".into(),
368 page_number: 5,
369 level: 0,
370 },
371 ],
372 };
373 let ranges = info.calculate_bookmark_ranges();
374 assert_eq!(
375 ranges,
376 vec![
377 ChunkMetadata {
378 start_page: 1,
379 end_page: 1,
380 bookmark_index: Some(1),
381 bookmark_title: Some("C1".into())
382 },
383 ChunkMetadata {
384 start_page: 2,
385 end_page: 4,
386 bookmark_index: Some(2),
387 bookmark_title: Some("C1.1".into())
388 },
389 ChunkMetadata {
390 start_page: 5,
391 end_page: 10,
392 bookmark_index: Some(3),
393 bookmark_title: Some("C2".into())
394 },
395 ]
396 );
397 }
398
399 #[test]
400 fn test_calculate_bookmark_ranges_unsorted() {
401 let info = PdfInfo {
402 total_pages: 10,
403 outlines: vec![
404 Bookmark {
405 title: "C2".into(),
406 page_number: 5,
407 level: 0,
408 },
409 Bookmark {
410 title: "C1".into(),
411 page_number: 1,
412 level: 0,
413 },
414 ],
415 };
416 let ranges = info.calculate_bookmark_ranges();
417 assert_eq!(
418 ranges,
419 vec![
420 ChunkMetadata {
421 start_page: 1,
422 end_page: 4,
423 bookmark_index: Some(2),
424 bookmark_title: Some("C1".into())
425 },
426 ChunkMetadata {
427 start_page: 5,
428 end_page: 10,
429 bookmark_index: Some(1),
430 bookmark_title: Some("C2".into())
431 },
432 ]
433 );
434 }
435
436 #[test]
437 fn test_calculate_hybrid_ranges_simple() {
438 let pdf = PdfInfo {
439 total_pages: 20,
440 outlines: vec![
441 Bookmark {
442 title: "Ch1".to_string(),
443 page_number: 1,
444 level: 0,
445 },
446 Bookmark {
447 title: "Ch2".to_string(),
448 page_number: 10,
449 level: 0,
450 },
451 ],
452 };
453
454 let ranges = pdf.calculate_hybrid_ranges(5);
455 assert_eq!(
456 ranges,
457 vec![
458 ChunkMetadata {
459 start_page: 1,
460 end_page: 5,
461 bookmark_index: Some(1),
462 bookmark_title: Some("Ch1".into())
463 },
464 ChunkMetadata {
465 start_page: 6,
466 end_page: 9,
467 bookmark_index: Some(1),
468 bookmark_title: Some("Ch1".into())
469 },
470 ChunkMetadata {
471 start_page: 10,
472 end_page: 14,
473 bookmark_index: Some(2),
474 bookmark_title: Some("Ch2".into())
475 },
476 ChunkMetadata {
477 start_page: 15,
478 end_page: 19,
479 bookmark_index: Some(2),
480 bookmark_title: Some("Ch2".into())
481 },
482 ChunkMetadata {
483 start_page: 20,
484 end_page: 20,
485 bookmark_index: Some(2),
486 bookmark_title: Some("Ch2".into())
487 },
488 ]
489 );
490 }
491
492 #[test]
493 fn test_calculate_hybrid_ranges_no_bookmarks() {
494 let pdf = PdfInfo {
495 total_pages: 10,
496 outlines: vec![],
497 };
498 let ranges = pdf.calculate_hybrid_ranges(5);
499 assert_eq!(
500 ranges,
501 vec![
502 ChunkMetadata {
503 start_page: 1,
504 end_page: 5,
505 bookmark_index: None,
506 bookmark_title: None
507 },
508 ChunkMetadata {
509 start_page: 6,
510 end_page: 10,
511 bookmark_index: None,
512 bookmark_title: None
513 },
514 ]
515 );
516 }
517
518 #[test]
519 fn test_calculate_hybrid_ranges_overlap_and_boundary() {
520 let pdf = PdfInfo {
521 total_pages: 10,
522 outlines: vec![
523 Bookmark {
524 title: "Start".into(),
525 page_number: 1,
526 level: 0,
527 },
528 Bookmark {
529 title: "Mid".into(),
530 page_number: 5,
531 level: 0,
532 },
533 Bookmark {
534 title: "End".into(),
535 page_number: 10,
536 level: 0,
537 },
538 ],
539 };
540
541 let ranges = pdf.calculate_hybrid_ranges(100);
542 assert_eq!(
543 ranges,
544 vec![
545 ChunkMetadata {
546 start_page: 1,
547 end_page: 4,
548 bookmark_index: Some(1),
549 bookmark_title: Some("Start".into())
550 },
551 ChunkMetadata {
552 start_page: 5,
553 end_page: 9,
554 bookmark_index: Some(2),
555 bookmark_title: Some("Mid".into())
556 },
557 ChunkMetadata {
558 start_page: 10,
559 end_page: 10,
560 bookmark_index: Some(3),
561 bookmark_title: Some("End".into())
562 },
563 ]
564 );
565 }
566
567 #[test]
568 fn test_calculate_hybrid_ranges_duplicate_pages() {
569 let pdf = PdfInfo {
570 total_pages: 10,
571 outlines: vec![
572 Bookmark {
573 title: "A".into(),
574 page_number: 5,
575 level: 0,
576 },
577 Bookmark {
578 title: "B".into(),
579 page_number: 5,
580 level: 1,
581 },
582 ],
583 };
584
585 let ranges = pdf.calculate_hybrid_ranges(10);
586 assert_eq!(ranges.len(), 2);
587 assert_eq!(ranges[0].start_page, 1);
588 assert_eq!(ranges[0].end_page, 4);
589 assert_eq!(ranges[1].start_page, 5);
590 assert_eq!(ranges[1].end_page, 10);
591 }
592
593 fn calculate_ranges_helper(total_pages: u32, pages_per_chunk: u32) -> Vec<(u32, u32)> {
594 let info = PdfInfo {
595 total_pages,
596 outlines: vec![],
597 };
598 info.calculate_page_ranges(pages_per_chunk)
599 .into_iter()
600 .map(|m| (m.start_page, m.end_page))
601 .collect()
602 }
603}