# Option A: uv installs everything automatically (https://docs.astral.sh/uv/)
uv run pdf_to_audiobook.py
 
# Option B: plain Python 3.9+
pip install pymupdf edge-tts
python3 pdf_to_audiobook.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.9"
# dependencies = ["pymupdf", "edge-tts"]
# ///
"""
pdf_to_audiobook.py: turn any PDF into an MP3 audiobook, one file per chapter.
 
Drop this file next to a PDF and run it:
 
    uv run pdf_to_audiobook.py            # uv installs the dependencies for you
    # or
    pip install pymupdf edge-tts && python3 pdf_to_audiobook.py
 
It finds the PDF in the current folder (it asks you to pick if there are several),
splits it into chapters using the PDF's table of contents, strips page numbers and
running headers, and narrates each chapter with a natural-sounding neural voice.
 
Options:
    python3 pdf_to_audiobook.py "My Book.pdf"     # choose a specific PDF
    python3 pdf_to_audiobook.py --all             # convert every PDF in the folder
    python3 pdf_to_audiobook.py --text-only       # write cleaned .txt files instead (to proofread)
    python3 pdf_to_audiobook.py --voice en-GB-RyanNeural
    python3 pdf_to_audiobook.py --list-voices
    python3 pdf_to_audiobook.py --engine openai   # needs OPENAI_API_KEY
"""
import argparse
import asyncio
import os
import re
import shutil
import sys
import tempfile
import time
from collections import Counter
from pathlib import Path
 
try:
    import fitz  # PyMuPDF
except ImportError:
    sys.exit("Missing dependency. Run:  pip install pymupdf edge-tts   (or use: uv run pdf_to_audiobook.py)")
 
DEFAULT_EDGE_VOICE = "en-US-AvaNeural"
DEFAULT_OPENAI_VOICE = "onyx"
SKIP_SECTIONS = ("bibliography", "references", "index", "contents", "table of contents",
                 "list of figures", "list of tables", "copyright")
 
CHAR_FIXES = {
    "fi": "fi", "fl": "fl", "ff": "ff", "ffi": "ffi", "ffl": "ffl",
    "’": "'", "‘": "'", "“": '"', "”": '"', "–": "-", "—": " - ", "\u00ad": "",
}
 
 
# ---------------------------------------------------------------- text cleanup
 
def _signature(line):
    """Normalise a line so 'Chapter 3 · 41' and 'Chapter 3 · 42' look identical."""
    return re.sub(r"\d+", "#", line.strip().lower())
 
 
def find_running_headers(doc):
    """Lines that repeat at the top/bottom of many pages are headers or footers."""
    counts = Counter()
    for page in doc:
        lines = [l for l in page.get_text().splitlines() if l.strip()]
        for line in set(lines[:3] + lines[-3:]):
            counts[_signature(line)] += 1
    threshold = max(3, len(doc) * 0.15)
    return {sig for sig, n in counts.items() if n >= threshold}
 
 
def clean_page(text, headers):
    kept = []
    for line in text.splitlines():
        line = line.strip()
        if not line or _signature(line) in headers:
            continue
        if re.fullmatch(r"\d+|[ivxlcdm]+", line.lower()):   # page numbers
            continue
        if line in ("•", "*", "■", "♦", "◦"):                # stray bullets
            continue
        for bad, good in CHAR_FIXES.items():
            line = line.replace(bad, good)
        kept.append(line)
 
    # Re-join lines, healing words hyphenated across a line break ("compu-" + "ter").
    out = ""
    for line in kept:
        if out.endswith("-") and len(out) > 1 and out[-2].isalpha() and line[:1].islower():
            out = out[:-1] + line
        else:
            out = f"{out} {line}" if out else line
    return out
 
 
def chunk_text(text, max_chars):
    """Split text on sentence boundaries so each TTS request stays under max_chars."""
    chunks, current = [], ""
    for sentence in re.split(r"(?<=[.!?])\s+", text):
        while len(sentence) > max_chars:              # pathological sentence: hard split
            chunks.append(sentence[:max_chars])
            sentence = sentence[max_chars:]
        if len(current) + len(sentence) + 1 > max_chars and current:
            chunks.append(current)
            current = ""
        current = f"{current} {sentence}" if current else sentence
    if current.strip():
        chunks.append(current)
    return chunks
 
 
# ---------------------------------------------------------------- chapters
 
def safe_name(title):
    name = re.sub(r"[^\w\s-]", "", title).strip()
    return re.sub(r"[-\s]+", "_", name)[:80] or "Untitled"
 
 
def get_chapters(doc, toc_level=None, pages_per_part=20):
    """
    Use the PDF's bookmarks to find chapters. By default we pick the shallowest
    outline level that has at least 3 entries (level 1 is often "Part I, Part II").
    PDFs without bookmarks are split into fixed-size parts instead.
    """
    toc = [(lvl, title.strip(), page) for lvl, title, page in doc.get_toc() if page > 0]
    if toc:
        levels = Counter(lvl for lvl, _, _ in toc)
        level = toc_level or next((l for l in sorted(levels) if levels[l] >= 3), min(levels))
        # Bookmarks are sometimes out of order or junk ("cp"), so tidy them up first.
        entries = sorted(((t, p) for lvl, t, p in toc if lvl == level and len(t) >= 3), key=lambda e: e[1])
        chapters = []
        for i, (title, start) in enumerate(entries):
            end = entries[i + 1][1] - 1 if i + 1 < len(entries) else len(doc)
            plain = title.lower().strip(" .:0123456789")
            if plain in SKIP_SECTIONS or plain.endswith(" index") or end < start:
                continue
            chapters.append({"title": title, "start": start, "end": end})
        if chapters:
            return chapters
 
    if len(doc) <= pages_per_part * 1.5:
        return [{"title": "Full Text", "start": 1, "end": len(doc)}]
    return [{"title": f"Part {n + 1}", "start": s + 1, "end": min(s + pages_per_part, len(doc))}
            for n, s in enumerate(range(0, len(doc), pages_per_part))]
 
 
def extract_text(doc, chapter, headers):
    pages = (clean_page(doc[i].get_text(), headers) for i in range(chapter["start"] - 1, chapter["end"]))
    return "\n\n".join(p for p in pages if p)
 
 
# ---------------------------------------------------------------- speech
 
async def synth_edge(text, out_path, voice):
    import edge_tts
    with open(out_path, "wb") as out:
        for chunk in chunk_text(text, 5000):
            for attempt in range(3):                       # the free endpoint occasionally hiccups
                try:
                    async for msg in edge_tts.Communicate(chunk, voice).stream():
                        if msg["type"] == "audio":
                            out.write(msg["data"])
                    break
                except Exception:
                    if attempt == 2:
                        raise
                    await asyncio.sleep(2 ** attempt)
 
 
async def synth_openai(text, out_path, voice):
    from openai import OpenAI
    client = OpenAI()  # reads OPENAI_API_KEY from the environment
    with open(out_path, "wb") as out:
        for chunk in chunk_text(text, 4000):
            def call(c=chunk):
                return client.audio.speech.create(model="tts-1", voice=voice, input=c).content
            out.write(await asyncio.to_thread(call))
 
 
async def convert_chapter(num, chapter, text, staging, final_dir, engine, voice, sem):
    name = f"{num:02d}_{safe_name(chapter['title'])}.mp3"
    if (final_dir / name).exists():
        print(f"  ↷ {name} already exists, skipping")
        return True
    async with sem:
        print(f"  ♪ Narrating {num:02d}: {chapter['title']} (pages {chapter['start']}-{chapter['end']})")
        started = time.time()
        tmp = staging / name
        try:
            await (synth_openai if engine == "openai" else synth_edge)(text, tmp, voice)
            # Write locally first, then move: cloud-synced folders (Google Drive, Dropbox)
            # can silently drop files during many concurrent writes.
            shutil.move(str(tmp), str(final_dir / name))
            print(f"  ✓ {name} ({time.time() - started:.0f}s)")
            return True
        except Exception as e:
            print(f"  ✗ {name} failed: {e}")
            return False
 
 
async def convert_pdf(pdf, args):
    doc = fitz.open(pdf)
    headers = find_running_headers(doc)
    chapters = get_chapters(doc, args.toc_level)
    if args.dry_run:
        chapters = chapters[:1]
 
    out_dir = pdf.parent / f"{pdf.stem}_audiobook"
    out_dir.mkdir(exist_ok=True)
    print(f"\n📖 {pdf.name}: {len(doc)} pages, {len(chapters)} chapter(s) -> {out_dir.name}/")
 
    texts = [extract_text(doc, ch, headers) for ch in chapters]
 
    if args.text_only:
        for num, (ch, text) in enumerate(zip(chapters, texts), 1):
            path = out_dir / f"{num:02d}_{safe_name(ch['title'])}.txt"
            path.write_text(text, encoding="utf-8")
            print(f"  ✓ {path.name} ({len(text.split()):,} words)")
        return
 
    sem = asyncio.Semaphore(args.workers)
    with tempfile.TemporaryDirectory(prefix="audiobook_") as staging:
        jobs = [convert_chapter(n, ch, t, Path(staging), out_dir, args.engine, args.voice, sem)
                for n, (ch, t) in enumerate(zip(chapters, texts), 1) if t.strip()]
        results = await asyncio.gather(*jobs)
    print(f"Done: {sum(results)}/{len(jobs)} chapters saved in {out_dir}")
 
 
# ---------------------------------------------------------------- main
 
def pick_pdfs(args):
    if args.pdf:
        path = Path(args.pdf).expanduser()
        if not path.exists():
            sys.exit(f"Can't find '{path}'.")
        return [path.absolute()]
 
    here = Path.cwd()
    pdfs = sorted(here.glob("*.pdf")) or sorted(Path(__file__).resolve().parent.glob("*.pdf"))
    if not pdfs:
        sys.exit("No PDF found. Put this script in the same folder as your PDF, or pass a path.")
    if len(pdfs) == 1 or args.all:
        return pdfs
 
    print("Found several PDFs:")
    for i, p in enumerate(pdfs, 1):
        print(f"  {i:>2}. {p.name}")
    choice = input("Which one? (number, or 'all'): ").strip().lower()
    if choice == "all":
        return pdfs
    if choice.isdigit() and 1 <= int(choice) <= len(pdfs):
        return [pdfs[int(choice) - 1]]
    sys.exit("No valid choice made.")
 
 
async def main():
    parser = argparse.ArgumentParser(description="Turn a PDF into an MP3 audiobook, one file per chapter.")
    parser.add_argument("pdf", nargs="?", help="PDF to convert (default: the PDF in this folder)")
    parser.add_argument("--all", action="store_true", help="convert every PDF in the folder")
    parser.add_argument("--engine", choices=["edge", "openai"], default="edge",
                        help="edge = free Microsoft neural voices (default); openai = paid, needs OPENAI_API_KEY")
    parser.add_argument("--voice", help=f"voice name (default: {DEFAULT_EDGE_VOICE} / {DEFAULT_OPENAI_VOICE})")
    parser.add_argument("--list-voices", action="store_true", help="show available English edge voices")
    parser.add_argument("--workers", type=int, default=4, help="chapters narrated in parallel (default: 4)")
    parser.add_argument("--toc-level", type=int, help="table-of-contents depth to treat as chapters")
    parser.add_argument("--text-only", action="store_true", help="write cleaned text files instead of audio")
    parser.add_argument("--dry-run", action="store_true", help="only process the first chapter")
    args = parser.parse_args()
 
    if args.list_voices:
        import edge_tts
        for v in sorted(await edge_tts.list_voices(), key=lambda v: v["ShortName"]):
            if v["Locale"].startswith("en-"):
                print(f"{v['ShortName']:<32} {v['Gender']}")
        return
 
    args.voice = args.voice or (DEFAULT_OPENAI_VOICE if args.engine == "openai" else DEFAULT_EDGE_VOICE)
    if args.engine == "openai" and not os.environ.get("OPENAI_API_KEY"):
        sys.exit("Set OPENAI_API_KEY to use the OpenAI engine (or drop --engine openai to use the free voices).")
 
    started = time.time()
    for pdf in pick_pdfs(args):
        await convert_pdf(pdf, args)
    print(f"\nFinished in {(time.time() - started) / 60:.1f} min.")
 
 
if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        sys.exit("\nStopped. Run again to pick up where you left off.")