UmbraTrans/tools/translation/book.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Load the Deal Hunter controller bundle into a book model.

The bundle is a JSONL manifest produced by UmbraAudio. One utterance is a
sentence or a short paragraph; the original DOCX paragraph boundaries are not
in the bundle. This module therefore reconstructs readable paragraphs with a
documented heuristic and keeps the source utterance indices on every
paragraph so a later audio-alignment step can still map back.

Nothing here calls an LLM.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from dataclasses import asdict, dataclass, field
from pathlib import Path

from languages import TargetLanguage, resolve

DEFAULT_BUNDLE = Path("/home/umbra/audio-read/bundles/dh-bundle/manifest.jsonl")

# Chapter titles per target language. Kept here, not inferred, so every render
# is stable and reviewable in one place; the English title stays on the
# paragraph.
#
# The Chinese table is the one the pipeline has always run with. A language
# with no table falls back to the English title rather than to another
# language's title: a Spanish edition with Chinese chapter headings is the
# same mislabelling defect as a Chinese draft under a Spanish name, one
# level down.
TITLES: dict[str, dict[str, str]] = {
    "zh": {
        "Sample Track": "音频版序言：声音介绍",
        "Opening Track": "开篇说明",
        "Copyright": "版权信息",
        "Author's Note": "作者的话",
        "Salvage": "打捞",
        "Family Reunion": "家庭团聚",
        "Selling Salvage": "出售打捞物",
        "Frustration": "挫败",
        "Next Bounty": "下一个赏金",
        "Time for a Change": "改变的时刻",
        "Trade Hub Vince-4": "贸易枢纽文斯-4",
        "Can I Join the Crew": "我能加入船员吗",
        "Salvage from Crashed Fighters": "坠毁战机的残骸",
        "Where Is the Captain?": "船长在哪里？",
        "Turn in Fugitives": "上交逃犯",
        "Her New Ship": "她的新船",
        "Introducing the new XO": "新任副舰长",
        "A New Ship Salvage": "新船打捞",
        "The Poker Game Distraction": "扑克牌局掩护",
        "Deal Hunter Bounty": "Deal Hunter 的赏金",
        "Selecting Bounties": "挑选赏金",
        "Shuttle Scans": "穿梭机扫描",
        "The Pirate Trap": "海盗陷阱",
        "The Vid Interview": "视频采访",
        "Hunting Pirates": "追猎海盗",
        "The Friendly Pirate": "友好的海盗",
        "Aunt Vicky": "维姬姑妈",
        "Preparation": "准备",
        "The Coronation": "加冕典礼",
        "Long Live the King": "国王万岁",
        "Let’s Make It Better": "让它变得更好",
        "Let's Make It Better": "让它变得更好",
        "Closing Track": "后记：感谢收听",
    },
}

# The book's subtitle per target language. The title itself stays in English
# - "Deal Hunter" is also the ship's name and the AI's name, so translating it
# would rename a character - and the subtitle is the translated reading aid.
# A language with no subtitle falls back to none rather than to another
# language's.
SUBTITLES: dict[str, str] = {
    "zh": "交易猎人",
    "es": "Cazador de Ofertas",
    "fr": "Chasseur de Primes",
    "de": "Prämienjäger",
    "it": "Cacciatore di Taglie",
    "pt": "Caçador de Recompensas",
    "ru": "Охотник за Головами",
    "ja": "賞金ハンター",
    "ko": "보상금 헌터",
    "ar": "صياد المكافآت",
}

MAX_NARRATION_SENTENCES = 3
MAX_NARRATION_WORDS = 58


@dataclass
class Utterance:
    index: int
    speaker: str
    type: str
    text: str


@dataclass
class Paragraph:
    id: str
    speaker: str
    type: str
    text: str
    source_indices: list[int] = field(default_factory=list)
    source_ids: list[str] = field(default_factory=list)

    @property
    def words(self) -> int:
        return len(self.text.split())


@dataclass
class Scene:
    id: str
    index: int
    paragraphs: list[Paragraph] = field(default_factory=list)


@dataclass
class Chapter:
    id: str
    index: int
    row: int
    title: str
    title_target: str
    matter: str
    paragraphs: list[Paragraph] = field(default_factory=list)
    scenes: list[Scene] = field(default_factory=list)

    def scenes_json(self) -> list[dict]:
        return [asdict(s) for s in self.scenes]


@dataclass
class Book:
    title: str
    title_target: str
    subtitle_target: str
    author: str
    source_language: str
    target_language: str
    bundle: str
    bundle_sha256: str
    chapters: list[Chapter] = field(default_factory=list)

    def chapter(self, ref: str | int) -> Chapter:
        for c in self.chapters:
            if str(c.index) == str(ref) or c.id == str(ref):
                return c
        raise KeyError(f"chapter not found: {ref}")

    def to_json(self) -> dict:
        return {
            "title": self.title,
            "title_target": self.title_target,
            "subtitle_target": self.subtitle_target,
            "author": self.author,
            "source_language": self.source_language,
            "target_language": self.target_language,
            "bundle": self.bundle,
            "bundle_sha256": self.bundle_sha256,
            "chapters": [
                {
                    **{k: v for k, v in asdict(c).items() if k != "scenes"},
                    "scenes": [asdict(s) for s in c.scenes],
                }
                for c in self.chapters
            ],
        }


def _sentences(text: str) -> int:
    # A conservative sentence count for grouping only. Quotes and ellipses are
    # the places where a naive split goes wrong, so count terminators instead.
    n = len(re.findall(r"[.!?]+(?:[”’\"']|$|\s)", text))
    return max(1, n)


def _classify(index: int, title: str) -> str:
    if index < 0:
        if title in {"Sample Track", "Opening Track"}:
            return "front_audio"
        if title == "Copyright":
            return "front_copyright"
        if title == "Author's Note":
            return "front_note"
        return "front"
    if title == "Closing Track" or index > 26:
        return "back"
    return "body"


def _clean_title(title: str) -> str:
    return title.replace("’", "'").strip()


def _title_for(title: str, target: TargetLanguage) -> str:
    """The chapter title in the target language.

    Two different fallbacks, for two different situations:

    - A language with **no title table** falls back to the English title. It
      does not fall back to another language's title: a Spanish edition
      carrying Chinese chapter headings is the same defect as a Chinese draft
      under a Spanish name, one level down.
    - A language **with** a table, and a chapter missing from it, gets a
      visible marker. An untranslated heading is a gap the reviewer has to
      see, not an English title that looks finished.
    """
    table = TITLES.get(target.base)
    if table is None:
        return title
    key = _clean_title(title)
    if key in table:
        return table[key]
    return target.string("missing_title", "[Untranslated] {title}").format(title=title)


def load_manifest(path: Path = DEFAULT_BUNDLE) -> tuple[list[dict], dict]:
    raw = path.read_bytes()
    digest = hashlib.sha256(raw).hexdigest()
    records = [json.loads(line) for line in raw.decode("utf-8").splitlines() if line.strip()]
    return records, {"path": str(path), "sha256": digest}


def build_book(path: Path = DEFAULT_BUNDLE, target_language: str = "zh-Hans") -> Book:
    target = resolve(target_language)
    records, meta = load_manifest(path)
    book_rec = next(r for r in records if r.get("kind") == "book")
    chapter_recs = sorted(
        (r for r in records if r.get("kind") == "chapter"),
        key=lambda r: r["chapter_index"],
    )
    utterances: dict[str, list[dict]] = {}
    for r in records:
        if r.get("kind") == "utterance":
            utterances.setdefault(r["chapter_id"], []).append(r)
    for rows in utterances.values():
        rows.sort(key=lambda r: r["utterance_index"])

    book = Book(
        title="Deal Hunter",
        title_target=_title_for("Deal Hunter", target),
        subtitle_target=SUBTITLES.get(target.base, ""),
        author="B. D. Murphy",
        source_language="en",
        target_language=target.code,
        bundle=meta["path"],
        bundle_sha256=meta["sha256"],
    )

    for cr in chapter_recs:
        index = int(cr["chapter_index"])
        title = cr["title"]
        matter = _classify(index, title)
        rows = utterances.get(cr["id"], [])
        paragraphs = _paragraphs(cr["id"], title, matter, rows)
        chapter = Chapter(
            id=cr["id"],
            index=index,
            row=int(cr["database_row_id"]),
            title=title,
            title_target=_title_for(title, target),
            matter=matter,
            paragraphs=paragraphs,
        )
        chapter.scenes = _scenes(cr["id"], paragraphs)
        book.chapters.append(chapter)

    return book


def _paragraphs(chapter_id: str, title: str, matter: str, rows: list[dict]) -> list[Paragraph]:
    if not rows:
        return []

    # The parser leaves the spoken chapter announcement in the first utterance
    # of Salvage ("Deal Hunter. Chapter one, Salvage."). It is a production
    # cue, not prose, so it must not be translated into the reading edition.
    if matter == "body" and rows[0]["text"].strip().startswith("Deal Hunter. Chapter"):
        rows = rows[1:]

    # The author's-note heading is an utterance, not a paragraph.
    if (
        matter == "front_note"
        and rows
        and rows[0]["text"].strip().rstrip(".").lower()
        in {
            "author's note",
            "authors note",
            "author’s note",
        }
    ):
        rows = rows[1:]

    paragraphs: list[Paragraph] = []
    scene_no = 0
    buf: list[dict] = []
    buf_words = 0
    buf_sentences = 0

    def flush() -> None:
        nonlocal buf, buf_words, buf_sentences
        if not buf:
            return
        paragraphs.append(
            Paragraph(
                id=f"{chapter_id}-p{len(paragraphs):04d}",
                speaker=buf[0].get("speaker_name_resolved") or "narrator",
                type=buf[0].get("utterance_type") or "narration",
                text=" ".join(r["text"].strip() for r in buf),
                source_indices=[int(r["utterance_index"]) for r in buf],
                source_ids=[r["id"] for r in buf],
            )
        )
        buf = []
        buf_words = 0
        buf_sentences = 0

    for row in rows:
        utype = row.get("utterance_type") or "narration"
        speaker = row.get("speaker_name_resolved") or "narrator"
        text = (row.get("text") or "").strip()
        if utype == "scene_break":
            flush()
            scene_no += 1
            paragraphs.append(
                Paragraph(
                    id=f"{chapter_id}-break{scene_no}",
                    speaker="",
                    type="scene_break",
                    text="* * *",
                    source_indices=[int(row["utterance_index"])],
                    source_ids=[row["id"]],
                )
            )
            continue

        if matter.startswith("front") or matter == "back":
            # Front and back matter are line-oriented: a copyright line and an
            # author's-note bullet are not the same thing as a prose paragraph.
            flush()
            paragraphs.append(
                Paragraph(
                    id=f"{chapter_id}-p{len(paragraphs):04d}",
                    speaker=speaker,
                    type=utype,
                    text=text,
                    source_indices=[int(row["utterance_index"])],
                    source_ids=[row["id"]],
                )
            )
            continue

        if utype == "dialogue":
            # Each speech turn is its own paragraph.
            flush()
            paragraphs.append(
                Paragraph(
                    id=f"{chapter_id}-p{len(paragraphs):04d}",
                    speaker=speaker,
                    type=utype,
                    text=text,
                    source_indices=[int(row["utterance_index"])],
                    source_ids=[row["id"]],
                )
            )
            continue

        # Narration: group consecutive sentences into a readable paragraph.
        if buf and (buf[0].get("speaker_name_resolved") != speaker or buf[0].get("utterance_type") != utype):
            flush()
        sentences = _sentences(text)
        if buf and (
            buf_sentences + sentences > MAX_NARRATION_SENTENCES or buf_words + len(text.split()) > MAX_NARRATION_WORDS
        ):
            flush()
        buf.append(row)
        buf_words += len(text.split())
        buf_sentences += sentences
        if buf_sentences >= MAX_NARRATION_SENTENCES or buf_words >= MAX_NARRATION_WORDS:
            flush()

    flush()
    # Drop a trailing scene-break marker.
    while paragraphs and paragraphs[-1].type == "scene_break":
        paragraphs.pop()
    return paragraphs


def _scenes(chapter_id: str, paragraphs: list[Paragraph]) -> list[Scene]:
    scenes: list[Scene] = []
    current = Scene(id=f"{chapter_id}-s{len(scenes)}", index=len(scenes))
    for p in paragraphs:
        if p.type == "scene_break":
            if current.paragraphs:
                scenes.append(current)
            current = Scene(id=f"{chapter_id}-s{len(scenes)}", index=len(scenes))
            continue
        current.paragraphs.append(p)
    if current.paragraphs:
        scenes.append(current)
    return scenes


def save(book: Book, out: Path) -> None:
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(book.to_json(), ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def load(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def summary(book: Book) -> str:
    lines = [
        f"Book: {book.title} / {book.title_target}",
        f"Author: {book.author}",
        f"Bundle: {book.bundle}",
        f"Bundle sha256: {book.bundle_sha256}",
        "",
        f"{'idx':>4} {'matter':16} {'paras':>6} {'words':>7}  title",
    ]
    for c in book.chapters:
        words = sum(p.words for p in c.paragraphs if p.type != "scene_break")
        lines.append(f"{c.index:>4} {c.matter:16} {len(c.paragraphs):>6} {words:>7}  {c.title} / {c.title_target}")
    body = [c for c in book.chapters if c.matter == "body"]
    lines += [
        "",
        f"Front matter chapters: {sum(1 for c in book.chapters if c.matter.startswith('front'))}",
        f"Body chapters: {len(body)}",
        f"Back matter chapters: {sum(1 for c in book.chapters if c.matter == 'back')}",
        f"Paragraphs: {sum(len(c.paragraphs) for c in book.chapters)}",
        f"Words: {sum(p.words for c in book.chapters for p in c.paragraphs if p.type != 'scene_break')}",
    ]
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser(description="Extract the Deal Hunter bundle into a book model.")
    ap.add_argument("--bundle", type=Path, default=DEFAULT_BUNDLE)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument(
        "--target-language",
        default="zh-Hans",
        help="the language the book is being translated into",
    )
    args = ap.parse_args()
    book = build_book(args.bundle, args.target_language)
    save(book, args.out)
    print(summary(book))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()