use crate::ExportEpubDto;
use crate::ExportEpubResultDto;
use crate::html_render;
use anyhow::{Result, anyhow};
use common::database::QueryUnitOfWork;
use common::database::Store;
use common::entities::{Block, Document, Frame, List, Root, SemanticRole, Table, TableCell};
use common::long_operation::{LongOperation, OperationProgress};
use common::parser_tools::EpubExportOptions;
use common::parser_tools::ExportImages;
use common::types::{EntityId, ROOT_ENTITY_ID};
use epub_builder::{
EpubBuilder, EpubContent, EpubVersion, PageDirection, ReferenceType, ZipLibrary,
};
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
pub trait ExportEpubUnitOfWorkFactoryTrait: Send + Sync {
fn create(&self) -> Box<dyn ExportEpubUnitOfWorkTrait>;
}
#[macros::uow_action(entity = "Root", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Root", action = "GetRelationshipRO", thread_safe = true)]
#[macros::uow_action(entity = "Document", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Document", action = "GetRelationshipRO", thread_safe = true)]
#[macros::uow_action(entity = "Frame", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Frame", action = "GetRelationshipRO", thread_safe = true)]
#[macros::uow_action(entity = "Block", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Block", action = "GetMultiRO", thread_safe = true)]
#[macros::uow_action(entity = "Block", action = "GetRelationshipRO", thread_safe = true)]
#[macros::uow_action(entity = "List", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Table", action = "GetRO", thread_safe = true)]
#[macros::uow_action(entity = "Table", action = "GetRelationshipRO", thread_safe = true)]
#[macros::uow_action(entity = "TableCell", action = "GetMultiRO", thread_safe = true)]
pub trait ExportEpubUnitOfWorkTrait: QueryUnitOfWork + Send + Sync {}
pub struct ExportEpubUseCase {
uow_factory: Box<dyn ExportEpubUnitOfWorkFactoryTrait>,
dto: ExportEpubDto,
}
impl ExportEpubUseCase {
pub fn new(
uow_factory: Box<dyn ExportEpubUnitOfWorkFactoryTrait>,
dto: &ExportEpubDto,
) -> Self {
ExportEpubUseCase {
uow_factory,
dto: dto.clone(),
}
}
}
impl LongOperation for ExportEpubUseCase {
type Output = ExportEpubResultDto;
fn execute(
&self,
progress_callback: Box<dyn Fn(OperationProgress) + Send>,
cancel_flag: Arc<AtomicBool>,
) -> Result<Self::Output> {
let output_path = std::path::Path::new(&self.dto.output_path);
if let Some(parent) = output_path.parent()
&& !parent.as_os_str().is_empty()
&& !parent.exists()
{
return Err(anyhow!(
"Output directory does not exist: '{}'",
parent.display()
));
}
progress_callback(OperationProgress::new(
0.0,
Some("Starting EPUB export...".to_string()),
));
let uow = self.uow_factory.create();
uow.begin_transaction()?;
let build_result = self.build_chapters(
&*uow,
progress_callback.as_ref(),
Some(cancel_flag.as_ref()),
);
uow.end_transaction()?;
let chapters = build_result?;
let chapter_count = chapters.len() as i64;
progress_callback(OperationProgress::new(
85.0,
Some("Packaging EPUB...".to_string()),
));
let epub_bytes = package_epub(
&self.dto.options,
&chapters,
&image_packaging_map(&self.dto.options.images),
)?;
progress_callback(OperationProgress::new(
90.0,
Some("Writing EPUB file...".to_string()),
));
std::fs::write(&self.dto.output_path, &epub_bytes).map_err(|e| {
anyhow!(
"Failed to write output file '{}': {}",
self.dto.output_path,
e
)
})?;
progress_callback(OperationProgress::new(100.0, Some("completed".to_string())));
Ok(ExportEpubResultDto {
file_path: self.dto.output_path.clone(),
chapter_count,
})
}
}
impl ExportEpubUseCase {
pub(crate) fn build_document(&self) -> Result<(Vec<u8>, i64)> {
let uow = self.uow_factory.create();
uow.begin_transaction()?;
let result = self.build_chapters(&*uow, &|_progress| {}, None);
uow.end_transaction()?;
let chapters = result?;
let chapter_count = chapters.len() as i64;
let epub_bytes = package_epub(
&self.dto.options,
&chapters,
&image_packaging_map(&self.dto.options.images),
)?;
Ok((epub_bytes, chapter_count))
}
pub(crate) fn build_chapters(
&self,
uow: &dyn ExportEpubUnitOfWorkTrait,
progress_callback: &dyn Fn(OperationProgress),
cancel_flag: Option<&AtomicBool>,
) -> Result<Vec<Chapter>> {
let root = uow
.get_root(&ROOT_ENTITY_ID)?
.ok_or_else(|| anyhow!("Root entity not found"))?;
let doc_ids = uow.get_root_relationship(
&root.id,
&common::direct_access::root::RootRelationshipField::Document,
)?;
let doc_id = *doc_ids
.first()
.ok_or_else(|| anyhow!("Root has no associated Document"))?;
let frame_ids = uow.get_document_relationship(
&doc_id,
&common::direct_access::document::DocumentRelationshipField::Frames,
)?;
let table_ids = uow.get_document_relationship(
&doc_id,
&common::direct_access::document::DocumentRelationshipField::Tables,
)?;
let mut cell_frame_ids: HashSet<EntityId> = HashSet::new();
for tid in &table_ids {
let cell_ids = uow.get_table_relationship(
tid,
&common::direct_access::table::TableRelationshipField::Cells,
)?;
let cells_opt = uow.get_table_cell_multi(&cell_ids)?;
for cell in cells_opt.into_iter().flatten() {
if let Some(cf_id) = cell.cell_frame {
cell_frame_ids.insert(cf_id);
}
}
}
progress_callback(OperationProgress::new(
10.0,
Some("Walking document tree...".to_string()),
));
let image_hrefs = image_packaging_map(&self.dto.options.images);
let image_policy = html_render::HtmlImagePolicy::Rewrite(&image_hrefs);
let notes = crate::footnotes::Footnotes::build(&uow.store());
let mut units: Vec<RenderUnit> = Vec::new();
let total_frames = frame_ids.len().max(1);
for (frame_idx, frame_id) in frame_ids.iter().enumerate() {
check_cancelled(cancel_flag)?;
if cell_frame_ids.contains(frame_id) {
continue;
}
if notes.is_definition(*frame_id) {
continue;
}
if let Some(f) = uow.get_frame(frame_id)?
&& f.parent_frame.is_some()
{
continue;
}
self.render_frame_units(
uow,
frame_id,
&cell_frame_ids,
¬es,
image_policy,
&mut units,
)?;
let pct = 10.0 + (frame_idx as f32 / total_frames as f32) * 60.0;
progress_callback(OperationProgress::new(
pct,
Some(format!(
"Processing frame {}/{}",
frame_idx + 1,
total_frames
)),
));
}
progress_callback(OperationProgress::new(
75.0,
Some("Splitting into chapters...".to_string()),
));
let mut chapters = split_into_chapters(units, &self.dto.options);
if !notes.is_empty() {
let in_print_order = notes.in_print_order();
for chapter in &mut chapters {
if chapter.footnote_labels.is_empty() {
continue;
}
let mut aside_html = String::new();
for (number, label, frame_id) in &in_print_order {
if !chapter.footnote_labels.contains(label) {
continue;
}
let mut inner: Vec<RenderUnit> = Vec::new();
self.render_frame_units(
uow,
frame_id,
&cell_frame_ids,
¬es,
image_policy,
&mut inner,
)?;
let body: String = inner.into_iter().map(|u| u.html).collect();
let id = html_render::escape_html(label);
aside_html.push_str(&format!(
"<aside epub:type=\"footnote\" role=\"doc-footnote\" id=\"fn-{id}\">\
<a href=\"#fnref-{id}\" role=\"doc-backlink\">{number}</a>. {body}</aside>"
));
}
chapter.body_html.push_str(&aside_html);
}
}
Ok(chapters)
}
fn render_frame_units(
&self,
uow: &dyn ExportEpubUnitOfWorkTrait,
frame_id: &EntityId,
cell_frame_ids: &HashSet<EntityId>,
notes: &crate::footnotes::Footnotes,
image_policy: html_render::HtmlImagePolicy<'_>,
out: &mut Vec<RenderUnit>,
) -> Result<()> {
let frame = uow
.get_frame(frame_id)?
.ok_or_else(|| anyhow!("Frame not found"))?;
if let Some(table_id) = frame.table {
let html = html_render::render_table_html(&uow.store(), table_id, image_policy, notes)?;
if !html.is_empty() {
let footnote_labels = table_footnote_labels(uow, table_id)?;
out.push(RenderUnit::content_with_labels(html, footnote_labels));
}
return Ok(());
}
if !frame.child_order.is_empty() {
return self.render_frame_units_by_child_order(
uow,
&frame,
cell_frame_ids,
notes,
image_policy,
out,
);
}
let block_ids = uow.get_frame_relationship(
frame_id,
&common::direct_access::frame::FrameRelationshipField::Blocks,
)?;
if block_ids.is_empty() {
return Ok(());
}
let blocks_opt = uow.get_block_multi(&block_ids)?;
let mut blocks: Vec<Block> = blocks_opt.into_iter().flatten().collect();
blocks.sort_by_key(|b| b.document_position);
push_block_run_units(&uow.store(), &blocks, image_policy, notes, out);
Ok(())
}
fn render_frame_units_by_child_order(
&self,
uow: &dyn ExportEpubUnitOfWorkTrait,
frame: &Frame,
cell_frame_ids: &HashSet<EntityId>,
notes: &crate::footnotes::Footnotes,
image_policy: html_render::HtmlImagePolicy<'_>,
out: &mut Vec<RenderUnit>,
) -> Result<()> {
let mut pending_blocks: Vec<Block> = Vec::new();
for &entry in &frame.child_order {
if entry > 0 {
let block_id = entry as u64;
if let Some(block) = uow.get_block(&block_id)? {
pending_blocks.push(block);
}
} else {
if !pending_blocks.is_empty() {
push_block_run_units(&uow.store(), &pending_blocks, image_policy, notes, out);
pending_blocks.clear();
}
let sub_frame_id = (-entry) as u64;
if cell_frame_ids.contains(&sub_frame_id) {
continue;
}
let sub_frame = uow.get_frame(&sub_frame_id)?;
if let Some(ref sf) = sub_frame {
if sf.fmt_is_blockquote == Some(true) {
let mut inner: Vec<RenderUnit> = Vec::new();
self.render_frame_units(
uow,
&sub_frame_id,
cell_frame_ids,
notes,
image_policy,
&mut inner,
)?;
let inner_labels =
dedup_labels(inner.iter().flat_map(|u| &u.footnote_labels));
let inner_html: String = inner.into_iter().map(|u| u.html).collect();
if !inner_html.is_empty() {
let semantics = match &sf.fmt_semantic_role {
Some(SemanticRole::Epigraph) => {
r#" epub:type="epigraph" role="doc-epigraph""#
}
None => "",
};
out.push(RenderUnit::content_with_labels(
format!("<blockquote{}>{}</blockquote>", semantics, inner_html),
inner_labels,
));
}
} else {
self.render_frame_units(
uow,
&sub_frame_id,
cell_frame_ids,
notes,
image_policy,
out,
)?;
}
}
}
}
if !pending_blocks.is_empty() {
push_block_run_units(&uow.store(), &pending_blocks, image_policy, notes, out);
}
Ok(())
}
}
pub(crate) struct Chapter {
title: String,
body_html: String,
footnote_labels: Vec<String>,
}
struct RenderUnit {
heading_level: Option<i64>,
heading_text: Option<String>,
html: String,
footnote_labels: Vec<String>,
}
impl RenderUnit {
fn content_with_labels(html: String, footnote_labels: Vec<String>) -> Self {
RenderUnit {
heading_level: None,
heading_text: None,
html,
footnote_labels,
}
}
}
fn dedup_labels<'a, I: IntoIterator<Item = &'a String>>(labels: I) -> Vec<String> {
let mut seen: HashSet<&str> = HashSet::new();
let mut out = Vec::new();
for label in labels {
if seen.insert(label.as_str()) {
out.push(label.clone());
}
}
out
}
fn footnote_labels_in_blocks(store: &Store, blocks: &[Block]) -> Vec<String> {
let refs = store.block_footnote_refs.read();
let mut raw: Vec<String> = Vec::new();
for block in blocks {
let Some(anchors) = refs.get(&block.id) else {
continue;
};
let mut anchors: Vec<_> = anchors.iter().collect();
anchors.sort_by_key(|a| a.byte_offset);
raw.extend(anchors.into_iter().map(|a| a.label.clone()));
}
dedup_labels(&raw)
}
fn table_footnote_labels(
uow: &dyn ExportEpubUnitOfWorkTrait,
table_id: EntityId,
) -> Result<Vec<String>> {
let cell_ids = uow.get_table_relationship(
&table_id,
&common::direct_access::table::TableRelationshipField::Cells,
)?;
let cells_opt = uow.get_table_cell_multi(&cell_ids)?;
let store = uow.store();
let refs = store.block_footnote_refs.read();
let mut raw: Vec<String> = Vec::new();
for cell in cells_opt.into_iter().flatten() {
let Some(cf_id) = cell.cell_frame else {
continue;
};
let block_ids = uow.get_frame_relationship(
&cf_id,
&common::direct_access::frame::FrameRelationshipField::Blocks,
)?;
for block_id in block_ids {
if let Some(anchors) = refs.get(&block_id) {
raw.extend(anchors.iter().map(|a| a.label.clone()));
}
}
}
Ok(dedup_labels(&raw))
}
fn heading_level_for_split(store: &Store, block: &Block) -> Option<i64> {
if block.fmt_is_code_block == Some(true) {
return None;
}
let is_listed = block
.list
.is_some_and(|list_id| store.lists.read().contains_key(&list_id));
if is_listed {
return None;
}
block.fmt_heading_level
}
fn push_block_run_units(
store: &Store,
blocks: &[Block],
image_policy: html_render::HtmlImagePolicy<'_>,
notes: &crate::footnotes::Footnotes,
out: &mut Vec<RenderUnit>,
) {
let mut i = 0;
while i < blocks.len() {
if let Some(level) = heading_level_for_split(store, &blocks[i]) {
let html = html_render::render_blocks_html(
store,
std::slice::from_ref(&blocks[i]),
image_policy,
notes,
);
let text = html_render::block_plain_text(store, &blocks[i]);
let footnote_labels =
footnote_labels_in_blocks(store, std::slice::from_ref(&blocks[i]));
out.push(RenderUnit {
heading_level: Some(level),
heading_text: Some(text),
html,
footnote_labels,
});
i += 1;
continue;
}
let start = i;
while i < blocks.len() && heading_level_for_split(store, &blocks[i]).is_none() {
i += 1;
}
let html = html_render::render_blocks_html(store, &blocks[start..i], image_policy, notes);
if !html.is_empty() {
let footnote_labels = footnote_labels_in_blocks(store, &blocks[start..i]);
out.push(RenderUnit::content_with_labels(html, footnote_labels));
}
}
}
fn split_into_chapters(units: Vec<RenderUnit>, options: &EpubExportOptions) -> Vec<Chapter> {
let front_title = if options.title.trim().is_empty() {
"Untitled".to_string()
} else {
options.title.clone()
};
let Some(target_level) = units.iter().filter_map(|u| u.heading_level).min() else {
let footnote_labels = dedup_labels(units.iter().flat_map(|u| &u.footnote_labels));
let body_html: String = units.into_iter().map(|u| u.html).collect();
return vec![Chapter {
title: front_title,
body_html,
footnote_labels,
}];
};
let mut chapters: Vec<Chapter> = Vec::new();
let mut current_title: Option<String> = None;
let mut current_html = String::new();
let mut current_labels: Vec<String> = Vec::new();
for unit in units {
if unit.heading_level == Some(target_level) {
if !current_html.is_empty() || current_title.is_some() {
chapters.push(Chapter {
title: current_title.take().unwrap_or_else(|| front_title.clone()),
body_html: std::mem::take(&mut current_html),
footnote_labels: dedup_labels(¤t_labels),
});
current_labels.clear();
}
current_title = unit.heading_text.clone();
}
current_labels.extend(unit.footnote_labels.iter().cloned());
current_html.push_str(&unit.html);
}
if !current_html.is_empty() || current_title.is_some() {
chapters.push(Chapter {
title: current_title.unwrap_or(front_title),
body_html: current_html,
footnote_labels: dedup_labels(¤t_labels),
});
}
chapters
}
fn wrap_xhtml(title: &str, lang: &str, rtl: bool, body_html: &str) -> String {
let dir_attr = if rtl { " dir=\"rtl\"" } else { "" };
format!(
"<?xml version=\"1.0\" encoding=\"utf-8\"?>\n\
<!DOCTYPE html>\n\
<html xmlns=\"http://www.w3.org/1999/xhtml\" xmlns:epub=\"http://www.idpf.org/2007/ops\" xml:lang=\"{lang}\" lang=\"{lang}\"{dir_attr}>\n\
<head><meta charset=\"utf-8\"/><title>{title}</title></head><body>{body}</body></html>",
lang = lang,
dir_attr = dir_attr,
title = html_render::escape_html(title),
body = body_html,
)
}
fn image_packaging_map(images: &ExportImages) -> std::collections::BTreeMap<String, String> {
images
.iter()
.enumerate()
.map(|(i, (src, image))| {
(
src.clone(),
format!("images/img_{:03}.{}", i + 1, image.extension()),
)
})
.collect()
}
fn package_epub(
options: &EpubExportOptions,
chapters: &[Chapter],
image_hrefs: &std::collections::BTreeMap<String, String>,
) -> Result<Vec<u8>> {
let lang = if options.language.trim().is_empty() {
"en"
} else {
options.language.trim()
};
let zip = ZipLibrary::new().map_err(|e| anyhow!("EPUB: {e}"))?;
let mut builder = EpubBuilder::new(zip).map_err(|e| anyhow!("EPUB: {e}"))?;
builder.epub_version(EpubVersion::V30);
builder.add_language(lang);
if !options.title.trim().is_empty() {
builder.set_title(options.title.trim());
}
if !options.author.trim().is_empty() {
builder.add_author(options.author.trim());
}
builder.set_generator("Skribisto");
if options.rtl {
builder.epub_direction(PageDirection::Rtl);
builder
.metadata("direction", "rtl")
.map_err(|e| anyhow!("EPUB: {e}"))?;
}
if let Some(cover) = &options.cover {
let href = format!("cover.{}", cover.extension());
builder
.add_cover_image(&href, cover.bytes.as_slice(), cover.mime_type.clone())
.map_err(|e| anyhow!("EPUB: adding cover: {e}"))?;
let alt = html_render::escape_html(if options.title.trim().is_empty() {
"Cover"
} else {
options.title.trim()
});
let body = format!(
"<div epub:type=\"cover\" style=\"text-align:center;margin:0;padding:0;\">\
<img src=\"{href}\" alt=\"{alt}\" style=\"max-width:100%;height:auto;\"/></div>"
);
let xhtml = wrap_xhtml("Cover", lang, options.rtl, &body);
builder
.add_content(
EpubContent::new("cover.xhtml", xhtml.as_bytes()).reftype(ReferenceType::Cover),
)
.map_err(|e| anyhow!("EPUB: adding cover page: {e}"))?;
}
for (src, href) in image_hrefs {
let Some(image) = options.images.get(src) else {
continue;
};
builder
.add_resource(href, image.bytes.as_slice(), image.mime_type.clone())
.map_err(|e| anyhow!("EPUB: adding image {src}: {e}"))?;
}
for (i, chapter) in chapters.iter().enumerate() {
let xhtml = wrap_xhtml(&chapter.title, lang, options.rtl, &chapter.body_html);
let href = format!("chapter_{:03}.xhtml", i + 1);
builder
.add_content(
EpubContent::new(href, xhtml.as_bytes())
.title(chapter.title.clone())
.reftype(ReferenceType::Text),
)
.map_err(|e| anyhow!("EPUB: {e}"))?;
}
builder.inline_toc();
let mut bytes: Vec<u8> = Vec::new();
builder
.generate(&mut bytes)
.map_err(|e| anyhow!("EPUB: {e}"))?;
Ok(bytes)
}
fn check_cancelled(cancel_flag: Option<&AtomicBool>) -> Result<()> {
if let Some(flag) = cancel_flag
&& flag.load(Ordering::Relaxed)
{
return Err(anyhow!("Operation was cancelled"));
}
Ok(())
}