use crate::models::content::ContentElement;
#[derive(Debug, Clone, PartialEq)]
pub struct TocEntry {
pub title: String,
pub level: u32,
pub page_number: u32,
pub children: Vec<TocEntry>,
}
#[derive(Debug, Clone, Default)]
pub struct TableOfContents {
pub entries: Vec<TocEntry>,
}
impl TableOfContents {
pub fn from_pages(pages: &[Vec<ContentElement>]) -> Self {
let mut flat: Vec<TocEntry> = Vec::new();
for (page_idx, page) in pages.iter().enumerate() {
let page_num = (page_idx + 1) as u32;
for elem in page {
if let Some(entry) = extract_heading(elem, page_num) {
flat.push(entry);
}
}
}
Self {
entries: nest_entries(flat),
}
}
pub fn total_entries(&self) -> usize {
count_entries(&self.entries)
}
pub fn to_markdown(&self) -> String {
let mut out = String::new();
render_markdown(&self.entries, &mut out);
out
}
pub fn to_html(&self) -> String {
if self.entries.is_empty() {
return String::new();
}
let mut out = String::new();
render_html(&self.entries, &mut out, 0);
out
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
fn extract_heading(elem: &ContentElement, page_number: u32) -> Option<TocEntry> {
match elem {
ContentElement::Heading(h) => {
let level = h.heading_level.unwrap_or(1);
let title = h.base.base.value().trim().to_string();
if title.is_empty() {
return None;
}
Some(TocEntry {
title,
level,
page_number,
children: Vec::new(),
})
}
_ => None,
}
}
fn nest_entries(flat: Vec<TocEntry>) -> Vec<TocEntry> {
if flat.is_empty() {
return Vec::new();
}
let mut result: Vec<TocEntry> = Vec::new();
let mut stack: Vec<TocEntry> = Vec::new();
for entry in flat {
while let Some(top) = stack.last() {
if top.level >= entry.level {
let popped = stack.pop().unwrap();
if let Some(parent) = stack.last_mut() {
parent.children.push(popped);
} else {
result.push(popped);
}
} else {
break;
}
}
stack.push(entry);
}
while let Some(popped) = stack.pop() {
if let Some(parent) = stack.last_mut() {
parent.children.push(popped);
} else {
result.push(popped);
}
}
result
}
fn count_entries(entries: &[TocEntry]) -> usize {
entries.iter().map(|e| 1 + count_entries(&e.children)).sum()
}
fn render_markdown(entries: &[TocEntry], out: &mut String) {
for entry in entries {
let indent = " ".repeat((entry.level - 1) as usize);
out.push_str(&format!(
"{}- {} (p. {})\n",
indent, entry.title, entry.page_number
));
render_markdown(&entry.children, out);
}
}
fn render_html(entries: &[TocEntry], out: &mut String, depth: usize) {
let indent = " ".repeat(depth);
out.push_str(&format!("{}<ul>\n", indent));
for entry in entries {
out.push_str(&format!(
"{} <li>{} (p. {})",
indent, entry.title, entry.page_number
));
if !entry.children.is_empty() {
out.push('\n');
render_html(&entry.children, out, depth + 2);
out.push_str(&format!("{} </li>\n", indent));
} else {
out.push_str("</li>\n");
}
}
out.push_str(&format!("{}</ul>\n", indent));
}
#[cfg(test)]
mod tests {
use super::*;
fn make_entry(title: &str, level: u32, page: u32) -> TocEntry {
TocEntry {
title: title.to_string(),
level,
page_number: page,
children: Vec::new(),
}
}
#[test]
fn test_nest_flat_entries() {
let flat = vec![
make_entry("Chapter 1", 1, 1),
make_entry("Section 1.1", 2, 2),
make_entry("Section 1.2", 2, 3),
make_entry("Chapter 2", 1, 5),
];
let nested = nest_entries(flat);
assert_eq!(nested.len(), 2);
assert_eq!(nested[0].title, "Chapter 1");
assert_eq!(nested[0].children.len(), 2);
assert_eq!(nested[0].children[0].title, "Section 1.1");
assert_eq!(nested[1].title, "Chapter 2");
assert!(nested[1].children.is_empty());
}
#[test]
fn test_total_entries() {
let toc = TableOfContents {
entries: vec![TocEntry {
title: "Ch1".to_string(),
level: 1,
page_number: 1,
children: vec![make_entry("S1.1", 2, 2), make_entry("S1.2", 2, 3)],
}],
};
assert_eq!(toc.total_entries(), 3);
}
#[test]
fn test_to_markdown() {
let toc = TableOfContents {
entries: vec![TocEntry {
title: "Intro".to_string(),
level: 1,
page_number: 1,
children: vec![make_entry("Overview", 2, 2)],
}],
};
let md = toc.to_markdown();
assert!(md.contains("- Intro (p. 1)"));
assert!(md.contains(" - Overview (p. 2)"));
}
#[test]
fn test_to_html() {
let toc = TableOfContents {
entries: vec![make_entry("Title", 1, 1)],
};
let html = toc.to_html();
assert!(html.contains("<ul>"));
assert!(html.contains("<li>Title (p. 1)</li>"));
assert!(html.contains("</ul>"));
}
#[test]
fn test_empty_toc() {
let toc = TableOfContents::default();
assert!(toc.is_empty());
assert_eq!(toc.total_entries(), 0);
assert_eq!(toc.to_markdown(), "");
assert_eq!(toc.to_html(), "");
}
}