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

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

# Chinese chapter titles. Kept here, not inferred, so every render is stable.
# Reviewable in one place; the English title stays on the paragraph.
TITLE_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": "后记：感谢收听",
}

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_zh: 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_zh: 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_zh": self.title_zh,
            "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_zh(title: str) -> str:
    key = _clean_title(title)
    if key in TITLE_ZH:
        return TITLE_ZH[key]
    # Unknown title: leave a visible marker rather than invent a translation.
    return f"[待译] {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) -> Book:
    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_zh="交易猎人",
        author="B. D. Murphy",
        source_language="en",
        target_language="zh-CN",
        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_zh=_title_zh(title),
            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_zh}",
        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_zh}"
        )
    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)
    args = ap.parse_args()
    book = build_book(args.bundle)
    save(book, args.out)
    print(summary(book))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()