From 4f46eb74f4dc805ead981aec395321e98bbd234e Mon Sep 17 00:00:00 2001
From: JB <jb-apps@outlook.com>
Date: Fri, 10 Oct 2025 12:08:22 -0700
Subject: [PATCH] feat: Add original text preservation in chunk overlays and
 enhance whitespace handling in EPUB export

---
 abogen/epub3/exporter.py    | 20 +++++++++++++++++---
 tests/test_epub_exporter.py | 13 +++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/abogen/epub3/exporter.py b/abogen/epub3/exporter.py
index b032dfa..e3e29f1 100644
--- a/abogen/epub3/exporter.py
+++ b/abogen/epub3/exporter.py
@@ -18,6 +18,7 @@ from abogen.text_extractor import ExtractedChapter, ExtractionResult
 class ChunkOverlay:
     id: str
     text: str
+    original_text: Optional[str]
     start: Optional[float]
     end: Optional[float]
     speaker_id: str
@@ -259,6 +260,7 @@ class EPUB3PackageBuilder:
                 ChunkOverlay(
                     id=normalized_id,
                     text=text or self.extraction.chapters[chapter_index].text,
+                    original_text=None,
                     start=_safe_float(marker.get("start")),
                     end=_safe_float(marker.get("end")),
                     speaker_id=speaker_id,
@@ -284,11 +286,11 @@ class EPUB3PackageBuilder:
             chunk_html = "<p></p>"
         original_block = ""
         if chapter.chunks:
-            original_text = "".join(chunk.text or "" for chunk in chapter.chunks)
+            original_text = "".join((chunk.original_text if chunk.original_text is not None else (chunk.text or "")) for chunk in chapter.chunks)
             if original_text:
                 safe_original = html.escape(original_text)
                 original_block = (
-                    "      <pre class=\"chapter-original\" hidden>\n"
+                    "      <pre class=\"chapter-original\" hidden=\"hidden\" aria-hidden=\"true\">\n"
                     f"{safe_original}\n"
                     "      </pre>"
                 )
@@ -701,12 +703,24 @@ def _restore_original_chunk_text(chapter_text: str, overlays: List[ChunkOverlay]
         if match is None and cursor:
             match = _search_original_span(chapter_text, candidate, 0)
         if match is None:
+            if chunk.original_text is None:
+                chunk.original_text = chunk.text
+            chunk.text = _prepare_display_text(chunk.text or "")
             continue
         start, end = match
-        chunk.text = chapter_text[start:end]
+        segment = chapter_text[start:end]
+        chunk.original_text = segment
+        chunk.text = _prepare_display_text(segment)
         cursor = end
 
 
+def _prepare_display_text(value: str) -> str:
+    if not value:
+        return ""
+    cleaned = re.sub(r"(?:[ \t]*\r?\n)+\Z", "", value)
+    return cleaned if cleaned else ""
+
+
 def _search_original_span(source: str, normalized: str, start: int) -> Optional[Tuple[int, int]]:
     if not normalized:
         return None
diff --git a/tests/test_epub_exporter.py b/tests/test_epub_exporter.py
index 3df703a..b2e572e 100644
--- a/tests/test_epub_exporter.py
+++ b/tests/test_epub_exporter.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import html
+import re
 import zipfile
 
 from abogen.epub3.exporter import build_epub3_package
@@ -167,5 +169,12 @@ def test_epub3_preserves_original_whitespace(tmp_path) -> None:
     with zipfile.ZipFile(output_path) as archive:
         chapter_doc = archive.read("OEBPS/text/chapter_0001.xhtml").decode("utf-8")
     assert "Line one  with  double spaces." in chapter_doc
-    normalized = chapter_doc.replace("        ", "")
-    assert "Second line\n\nThird paragraph." in normalized
\ No newline at end of file
+
+    chunk_section = chapter_doc.replace("        ", "")
+    assert "Second line" in chunk_section
+    assert "Third paragraph." in chunk_section
+
+    match = re.search(r"<pre class=\"chapter-original\"[^>]*>(.*?)</pre>", chapter_doc, re.DOTALL)
+    assert match is not None
+    original_text = html.unescape(match.group(1))
+    assert "Second line\n\nThird paragraph." in original_text
\ No newline at end of file