feat(converter): add per-chapter and m4b output

4 files changed, 353 insertions, 26 deletions

diff --git a/converter/audio.py b/converter/audio.py
index e7fa0bb..2f40777 100644
--- a/converter/audio.py
+++ b/converter/audio.py
@@ -47,8 +47,16 @@ def _concat_escape(path: str) -> str:
     return path.replace("'", "'\\''")
 
 
+def _encode_args(output_format: str) -> List[str]:
+    """Return ffmpeg output codec/bitrate args for the requested container."""
+    if output_format == "m4b":
+        return ["-c:a", "aac", "-b:a", config.AUDIO_BITRATE]
+    return ["-b:a", config.AUDIO_BITRATE]
+
+
 def combine_chunks(total_chunks: int, output_path: Path,
-                   results: Optional[Dict[int, bool]] = None, speed: float = 1.0) -> bool:
+                   results: Optional[Dict[int, bool]] = None, speed: float = 1.0,
+                   output_format: str = "mp3") -> bool:
     """Combine audio chunks into the final audiobook using ffmpeg's concat demuxer.
 
     ``results`` maps chunk numbers to success flags; failed chunks are
@@ -88,20 +96,21 @@ def combine_chunks(total_chunks: int, output_path: Path,
                 list_file.write(f"file '{_concat_escape(str(chunk_file))}'\n")
 
         filters = atempo_filters(speed)
+        encode = _encode_args(output_format)
         if filters:
             speed_path = output_path.with_name(f"{output_path.stem}_{speed:g}{output_path.suffix}")
             cmd = [
                 "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
                 "-filter_complex",
                 f"[0:a]split=2[base][spd];[spd]{filters}[spdout]",
-                "-map", "[base]", "-b:a", config.AUDIO_BITRATE, str(output_path),
-                "-map", "[spdout]", "-b:a", config.AUDIO_BITRATE, str(speed_path),
+                "-map", "[base]", *encode, str(output_path),
+                "-map", "[spdout]", *encode, str(speed_path),
             ]
         else:
             speed_path = None
             cmd = [
                 "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
-                "-b:a", config.AUDIO_BITRATE, str(output_path),
+                *encode, str(output_path),
             ]
 
         proc = subprocess.run(cmd, capture_output=True, text=True)
@@ -136,19 +145,138 @@ def combine_chunks(total_chunks: int, output_path: Path,
 
 
 def cleanup_chunks() -> None:
-    """Remove temporary chunk files from the scratch folder."""
+    """Remove temporary chunk and chapter files from the scratch folder."""
     try:
         chunk_count = 0
-        for chunk_file in config.CHUNKS_FOLDER.glob("chunk_*"):
-            try:
-                if chunk_file.is_file():
-                    chunk_file.unlink()
-                    chunk_count += 1
-            except Exception as exc:
-                logger.warning("Failed to delete %s: %s", chunk_file, exc)
+        for pattern in ("chunk_*", "chapter_*"):
+            for chunk_file in config.CHUNKS_FOLDER.glob(pattern):
+                try:
+                    if chunk_file.is_file():
+                        chunk_file.unlink()
+                        chunk_count += 1
+                except Exception as exc:
+                    logger.warning("Failed to delete %s: %s", chunk_file, exc)
 
         if chunk_count > 0:
             logger.info("Cleaned up %d chunk files", chunk_count)
             print(f"[INFO] Cleaned up {chunk_count} chunk files")
     except Exception as exc:
         logger.warning("Cleanup failed: %s", exc)
+
+
+def probe_duration_ms(path: Path) -> int:
+    """Return audio duration in milliseconds using ffprobe."""
+    result = subprocess.run(
+        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
+         "-of", "default=noprint_wrappers=1:nokey=1", str(path)],
+        capture_output=True, text=True,
+    )
+    if result.returncode != 0:
+        logger.warning("ffprobe failed for %s: %s", path, result.stderr[-200:])
+        return 0
+    try:
+        return max(0, int(round(float(result.stdout.strip()) * 1000.0)))
+    except ValueError:
+        logger.warning("Could not parse ffprobe duration for %s", path)
+        return 0
+
+
+def build_ffmetadata(chapters: List[tuple], path: Path) -> None:
+    """Write an ffmpeg FFMETADATA file with ``[CHAPTER]`` entries.
+
+    ``chapters`` is a list of ``(start_ms, end_ms, title)`` tuples.
+    """
+    with open(path, "w", encoding="utf-8") as meta_file:
+        meta_file.write(";FFMETADATA1\n")
+        for start_ms, end_ms, title in chapters:
+            meta_file.write("[CHAPTER]\n")
+            meta_file.write("TIMEBASE=1/1000\n")
+            meta_file.write(f"START={int(start_ms)}\n")
+            meta_file.write(f"END={int(end_ms)}\n")
+            meta_file.write(f"title={title}\n")
+
+
+def combine_chapters_to_m4b(chapter_files: List[Path], titles: List[str],
+                            output_path: Path, speed: float = 1.0) -> bool:
+    """Concatenate per-chapter audio into a single m4b with embedded chapter markers.
+
+    Chapter start/end times are derived from each chapter file's duration and
+    written as ffmpeg chapter metadata. When ``speed`` differs from 1.0, a
+    speed-adjusted copy (with rescaled chapter markers) is written alongside
+    the normal-speed file.
+    """
+    if shutil.which("ffmpeg") is None or shutil.which("ffprobe") is None:
+        logger.error("ffmpeg and ffprobe are required to build an m4b with chapters")
+        return False
+
+    if not chapter_files:
+        logger.error("No chapter files provided")
+        return False
+
+    concat_list = config.CHUNKS_FOLDER / "_concat_list.txt"
+    metadata_file = config.CHUNKS_FOLDER / "_chapters.txt"
+    speed_metadata_file = config.CHUNKS_FOLDER / "_chapters_speed.txt"
+    try:
+        chapters = []
+        start_ms = 0
+        with open(concat_list, "w", encoding="utf-8") as list_file:
+            for chapter_file, title in zip(chapter_files, titles):
+                list_file.write(f"file '{_concat_escape(str(chapter_file))}'\n")
+                duration_ms = probe_duration_ms(chapter_file)
+                end_ms = start_ms + duration_ms
+                chapters.append((start_ms, end_ms, title or "Chapter"))
+                start_ms = end_ms
+
+        build_ffmetadata(chapters, metadata_file)
+
+        filters = atempo_filters(speed)
+        encode = _encode_args("m4b")
+        if filters:
+            speed_path = output_path.with_name(f"{output_path.stem}_{speed:g}{output_path.suffix}")
+            scaled = [(int(s / speed), int(e / speed), t) for s, e, t in chapters]
+            build_ffmetadata(scaled, speed_metadata_file)
+            # The speed-adjusted stream needs rescaled chapter markers, so the
+            # rescaled metadata is passed as a third input.
+            cmd = [
+                "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
+                "-i", str(metadata_file), "-i", str(speed_metadata_file),
+                "-filter_complex",
+                f"[0:a]split=2[base][spd];[spd]{filters}[spdout]",
+                "-map", "[base]", "-map_metadata", "1", *encode, str(output_path),
+                "-map", "[spdout]", "-map_metadata", "2", *encode, str(speed_path),
+            ]
+        else:
+            speed_path = None
+            cmd = [
+                "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", str(concat_list),
+                "-i", str(metadata_file),
+                "-map", "0:a", "-map_metadata", "1", *encode, str(output_path),
+            ]
+
+        proc = subprocess.run(cmd, capture_output=True, text=True)
+        if proc.returncode != 0:
+            logger.error("ffmpeg failed: %s", proc.stderr[-2000:])
+            return False
+
+        logger.info("Audiobook saved: %s (%d chapters)", output_path, len(chapter_files))
+        print(f"[INFO] Saved audiobook: {output_path.name} ({len(chapter_files)} chapters)")
+
+        if speed_path is not None:
+            logger.info("Saved speed-adjusted audiobook (%gx): %s", speed, speed_path)
+            print(f"[INFO] Saved speed-adjusted audiobook: {speed_path.name} ({speed:g}x)")
+
+        return True
+
+    except FileNotFoundError:
+        logger.error("ffmpeg/ffprobe not found on PATH (install ffmpeg and try again)")
+        return False
+    except Exception as exc:
+        logger.error("Failed to combine chapters: %s", exc)
+        logger.error(traceback.format_exc())
+        return False
+    finally:
+        for scratch in (concat_list, metadata_file, speed_metadata_file):
+            try:
+                scratch.unlink(missing_ok=True)
+            except OSError:
+                pass
diff --git a/converter/config.py b/converter/config.py
index 1b203d2..5542bae 100644
--- a/converter/config.py
+++ b/converter/config.py
@@ -61,8 +61,8 @@ VOICE_CLONE_API_URL = "http://127.0.0.1:7861"
 # PROCESSING SETTINGS
 # =============================================================================
 
-BOOKS_FOLDER = BASE_DIR / "book_to_convert"  # Input folder
-AUDIOBOOKS_FOLDER = BASE_DIR / "audiobooks"  # Output folder
+BOOKS_FOLDER = BASE_DIR / "input"  # Input folder
+AUDIOBOOKS_FOLDER = BASE_DIR / "output"  # Output folder
 CHUNKS_FOLDER = BASE_DIR / "chunks"  # Scratch space for per-chunk audio (cleaned per book)
 LOGS_FOLDER = BASE_DIR / "logs"
 
@@ -74,7 +74,8 @@ HEARTBEAT_INTERVAL_SECONDS = 30  # Print "still working" this often during a chu
 # AUDIO OUTPUT SETTINGS
 # =============================================================================
 
-AUDIO_FORMAT = "mp3"
+AUDIO_FORMAT = "mp3"  # Default output container ("mp3" or "m4b")
+AUDIO_FORMATS = ("mp3", "m4b")
 AUDIO_BITRATE = "128k"
 
 SUPPORTED_FORMATS = [".txt", ".pdf", ".epub"]
diff --git a/converter/converter.py b/converter/converter.py
index 65efb68..ce2c351 100644
--- a/converter/converter.py
+++ b/converter/converter.py
@@ -1,6 +1,7 @@
 """Orchestrates book-to-audiobook conversion."""
 
 import logging
+import re
 import sys
 import time
 import traceback
@@ -42,12 +43,16 @@ class AudiobookConverter:
 
     def __init__(self, voice_mode: str = "custom_voice", voice_clone_ref_audio: Optional[str] = None,
                  voice_clone_ref_text: Optional[str] = None, skip_transcription: bool = False,
-                 speed: float = 1.0):
+                 speed: float = 1.0, single_file: bool = False, output_format: str = "mp3"):
         if speed <= 0:
             raise ValueError(f"Speed must be a positive number, got {speed}")
+        if output_format not in config.AUDIO_FORMATS:
+            raise ValueError(f"Unsupported output format: {output_format}")
         self.voice_mode = voice_mode
         self.voice_clone_ref_audio = voice_clone_ref_audio
         self.speed = speed
+        self.single_file = single_file
+        self.output_format = output_format
         self._validate_configuration()
         self.tts = QwenTTSClient(
             voice_mode=voice_mode,
@@ -70,8 +75,15 @@ class AudiobookConverter:
                     f"Reference audio file not found: {self.voice_clone_ref_audio}"
                 )
 
+    @staticmethod
+    def _sanitize_filename(name: str) -> str:
+        """Make a chapter title safe to use as part of a file name."""
+        cleaned = re.sub(r'[\\/:*?"<>|]', " ", name)
+        cleaned = re.sub(r"\s+", " ", cleaned).strip().strip(".")
+        return cleaned[:80] or "chapter"
+
     def convert_book(self, file_path: Path, output_name: Optional[str] = None) -> bool:
-        """Convert a single book to an audiobook."""
+        """Convert a single book to one or more audiobook files."""
         logger.info("Converting: %s", file_path.name)
         start_time = time.time()
 
@@ -80,13 +92,70 @@ class AudiobookConverter:
             # affect this run
             audio.cleanup_chunks()
 
-            # Extract text
             logger.info("Extracting text...")
-            text = extractors.extract_text(file_path)
-            if not text.strip():
+            sections = extractors.extract_sections(file_path)
+            if not sections or all(not s.text.strip() for s in sections):
                 logger.error("No text extracted")
                 return False
 
+            stem = output_name or file_path.stem
+            embed_chapters = (self.output_format == "m4b" and self.single_file
+                              and len(sections) > 1)
+            if embed_chapters:
+                return self._convert_single_m4b_with_chapters(sections, stem, start_time)
+
+            if self.single_file or len(sections) == 1:
+                text = "\n\n".join(section.text for section in sections)
+                output_path = config.AUDIOBOOKS_FOLDER / f"{stem}.{self.output_format}"
+                return self._convert_text(text, output_path, start_time)
+
+            success = True
+            for index, section in enumerate(sections, 1):
+                chapter_name = f"{stem}_{index:02d}_{self._sanitize_filename(section.title)}"
+                output_path = config.AUDIOBOOKS_FOLDER / f"{chapter_name}.{self.output_format}"
+                success = self._convert_text(section.text, output_path, start_time) and success
+            return success
+
+        except Exception as exc:
+            logger.error("Conversion failed: %s", exc)
+            logger.error(traceback.format_exc())
+            return False
+        finally:
+            # Always cleanup, even on failure or interrupt
+            audio.cleanup_chunks()
+
+    def _convert_single_m4b_with_chapters(self, sections, stem: str, start_time: float) -> bool:
+        """Convert each chapter to audio, then assemble a single m4b with
+        embedded chapter markers."""
+        chapter_files = []
+        titles = []
+        for index, section in enumerate(sections, 1):
+            chapter_path = config.CHUNKS_FOLDER / f"chapter_{index:04d}.{self.output_format}"
+            if not self._convert_text(section.text, chapter_path, start_time, speed=1.0):
+                logger.warning("Skipping chapter %d (%s) due to conversion failure",
+                               index, section.title)
+                continue
+            chapter_files.append(chapter_path)
+            titles.append(section.title or f"Chapter {index}")
+
+        if not chapter_files:
+            logger.error("No chapters were successfully converted")
+            return False
+
+        output_path = config.AUDIOBOOKS_FOLDER / f"{stem}.{self.output_format}"
+        return audio.combine_chapters_to_m4b(chapter_files, titles, output_path, speed=self.speed)
+
+    def _convert_text(self, text: str, output_path: Path, start_time: float,
+                      speed: Optional[float] = None) -> bool:
+        """Chunk, synthesize, and assemble ``text`` into ``output_path``."""
+        if speed is None:
+            speed = self.speed
+
+        try:
+            if not text.strip():
+                logger.error("No text to convert for %s", output_path.name)
+                return False
+
             logger.info("Extracted %d characters (%d words)", len(text), len(text.split()))
 
             # Split into chunks
@@ -143,8 +212,8 @@ class AudiobookConverter:
                                successful_chunks, total_chunks)
 
             # Combine chunks (only the successful ones)
-            output_path = config.AUDIOBOOKS_FOLDER / f"{output_name or file_path.stem}.{config.AUDIO_FORMAT}"
-            success = audio.combine_chunks(total_chunks, output_path, results, speed=self.speed)
+            success = audio.combine_chunks(total_chunks, output_path, results,
+                                           speed=speed, output_format=self.output_format)
 
             if success:
                 duration = time.time() - start_time
@@ -161,9 +230,6 @@ class AudiobookConverter:
             logger.error("Conversion failed: %s", exc)
             logger.error(traceback.format_exc())
             return False
-        finally:
-            # Always cleanup, even on failure or interrupt
-            audio.cleanup_chunks()
 
     def run(self) -> bool:
         """Main conversion process. Returns True if all books converted."""
@@ -183,7 +249,9 @@ class AudiobookConverter:
         elif self.voice_mode == "voice_clone":
             print(f"Reference audio: {Path(self.voice_clone_ref_audio).name}")
             print(f"Language: {config.VOICE_CLONE_LANGUAGE}")
-        print(f"Output format: {config.AUDIO_FORMAT}")
+        print(f"Output format: {self.output_format}")
+        if self.single_file:
+            print("Chapter mode: single file (--single-file)")
         if abs(self.speed - 1.0) >= 1e-6:
             print(f"Playback speed: {self.speed:g}x")
         print("=" * 70)
diff --git a/converter/extractors.py b/converter/extractors.py
index e6e27ad..9139626 100644
--- a/converter/extractors.py
+++ b/converter/extractors.py
@@ -6,6 +6,7 @@ import re
 import zipfile
 from html import unescape
 from pathlib import Path
+from typing import List, NamedTuple
 
 try:
     from bs4 import BeautifulSoup
@@ -16,6 +17,13 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
+class Section(NamedTuple):
+    """A titled chunk of a book (e.g. an EPUB chapter)."""
+
+    title: str
+    text: str
+
+
 def extract_text(file_path: Path) -> str:
     """Extract text from a book file based on its extension."""
     extension = file_path.suffix.lower()
@@ -28,6 +36,128 @@ def extract_text(file_path: Path) -> str:
     raise ValueError(f"Unsupported file format: {extension}")
 
 
+def extract_sections(file_path: Path) -> List[Section]:
+    """Extract the book's text as titled sections (chapters).
+
+    EPUB files are split on their spine documents so they can be converted
+    one chapter at a time. TXT and PDF files have no chapter structure and
+    always yield a single section.
+    """
+    extension = file_path.suffix.lower()
+    if extension == ".epub":
+        chapters = _extract_epub_chapters(file_path)
+        if len(chapters) > 1:
+            return chapters
+        return [Section(file_path.stem, extract_epub(file_path))]
+
+    return [Section(file_path.stem, extract_text(file_path))]
+
+
+def _extract_epub_chapters(file_path: Path) -> List[Section]:
+    """Return one Section per EPUB spine document (chapter), in reading order."""
+    import ebooklib
+
+    book = None
+    for method in (_read_epub_ebooklib, _read_epub_zipfile, _read_epub_manual):
+        try:
+            book = method(file_path)
+        except Exception as exc:
+            logger.warning("EPUB chapter method %s failed: %s", method.__name__, exc)
+            continue
+        if book:
+            break
+
+    if book is None:
+        return []
+
+    chapters = []
+    for title, text in book:
+        cleaned = clean_html(text)
+        if cleaned.strip():
+            chapters.append(Section(title or file_path.stem, cleaned))
+    return chapters
+
+
+def _toc_titles(book) -> dict:
+    """Flatten an ebooklib TOC into a ``{href: title}`` mapping."""
+    titles = {}
+
+    def walk(nodes) -> None:
+        for node in nodes:
+            if isinstance(node, (tuple, list)):
+                walk(node[1] if len(node) > 1 else [])
+                continue
+            href = getattr(node, "href", None)
+            title = getattr(node, "title", None)
+            if href and title:
+                titles[href.split("#")[0]] = title
+
+    walk(book.toc)
+    return titles
+
+
+def _read_epub_ebooklib(file_path: Path):
+    """Read EPUB spine documents as (title, html) pairs via ebooklib."""
+    import ebooklib
+    from ebooklib import epub
+
+    book = epub.read_epub(str(file_path))
+    titles = _toc_titles(book)
+    items = []
+    for entry in book.spine:
+        item_id = entry[0] if isinstance(entry, (tuple, list)) else entry
+        try:
+            item = book.get_item_with_id(item_id)
+        except Exception as exc:
+            logger.debug("Skipping EPUB spine item %r: %s", item_id, exc)
+            continue
+        if not item or item.get_type() != ebooklib.ITEM_DOCUMENT:
+            continue
+        if isinstance(item, epub.EpubNav):
+            continue
+        content = item.get_body_content()
+        if content:
+            if isinstance(content, bytes):
+                content = content.decode("utf-8", errors="ignore")
+            title = (titles.get(item.file_name)
+                     or titles.get(item.get_name())
+                     or getattr(item, "title", None)
+                     or item.get_name())
+            items.append((title, str(content)))
+    return items
+
+
+def _read_epub_zipfile(file_path: Path):
+    """Read EPUB HTML members as (title, html) pairs, ordered by filename."""
+    items = []
+    with zipfile.ZipFile(file_path, "r") as epub_zip:
+        for file_name in sorted(epub_zip.namelist(), key=_natural_key):
+            if file_name.lower().endswith((".html", ".xhtml", ".htm")):
+                try:
+                    content = epub_zip.read(file_name).decode("utf-8", errors="ignore")
+                    items.append((Path(file_name).stem, content))
+                except Exception as exc:
+                    logger.debug("Skipping EPUB member %r: %s", file_name, exc)
+    return items
+
+
+def _read_epub_manual(file_path: Path):
+    """Last-resort read of any markup-looking EPUB member."""
+    skipped_extensions = (".jpg", ".jpeg", ".png", ".gif", ".css", ".js")
+    items = []
+    with zipfile.ZipFile(file_path, "r") as epub_zip:
+        for file_name in sorted(epub_zip.namelist(), key=_natural_key):
+            if file_name.lower().endswith(skipped_extensions):
+                continue
+            try:
+                content = epub_zip.read(file_name).decode("utf-8", errors="ignore")
+                if "<" in content and len(content.strip()) > 100:
+                    items.append((Path(file_name).stem, content))
+            except Exception as exc:
+                logger.debug("Skipping EPUB member %r: %s", file_name, exc)
+    return items
+
+
 def clean_text(text: str) -> str:
     """Normalize whitespace and strip standalone page numbers.