path: root/app/converter/chunking.py

"""Split extracted book text into TTS-sized chunks."""

import logging
import re
from typing import List, Optional

from . import config

logger = logging.getLogger(__name__)

# Smart-chunking tuning, as fractions of the word limit (not exposed as
# settings): once a smart chunk crosses the target it ends at the next
# clean boundary, so most chunks land between the target and the limit
# with some headroom under the configured size. When the next sentence
# will not fit the running chunk, the break retreats to the previous
# boundary that is clean and at least the floor, so a quotation is not
# severed when a later clean boundary exists nearby.
_TARGET_RATIO = 0.85
_FLOOR_RATIO = 0.7
# A single token with no whitespace cannot be broken at a word
# boundary; if such a token runs longer than this many characters per
# allowed word it is sliced at the hard character cap (a last-resort
# bound for punctuation-free input such as some CJK text or damage from
# text extraction). The slice also honors the word limit itself, since
# every non-spaced-script character counts as one word.
_MAX_CHARS_PER_WORD = 8

_ABBREVIATIONS = frozenset(
    "mr mrs ms miss dr prof rev hon pres gov sen rep capt lt col gen "
    "sgt maj cpl pvt jr sr st mt vs etc eg ie aka fig figs no nos vol "
    "vols ch pp p pg dept univ ave blvd rd inc ltd co corp bros apt "
    "ste est approx min sec hrs msg mme mlle".split())

_EVENT = re.compile(
    r"(?P<end>[.!?…。！？][\"'”’»)}\]]*)"
    r"|(?P<para>\n[ \t]*\n)"
    r"|(?P<open>[“«])"
    r"|(?P<close>[”»])"
    r"|(?P<straight>\")")

_PARA_AFTER = re.compile(r"[ \t\r]*\n[ \t]*\n")

# Scripts whose characters are not whitespace-delimited: count each one
# as a word so a long CJK passage is bounded like English text of the
# same spoken length instead of reading as one giant "word".
_NON_SPACED_RUN = re.compile(r"[\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff"
                             r"\uf900-\ufaff]")


class _Unit:
    """A sentence-sized piece of text plus its packing metadata.

    ``in_quote_end`` is the quotation state after the unit's terminator
    (double-quote open marks and closes are tracked; apostrophes do not
    count as quotes), and ``para_after`` marks a unit that a blank line
    follows.
    """

    __slots__ = ("text", "words", "para_after", "in_quote_end")

    def __init__(self, text: str, words: int, para_after: bool,
                 in_quote_end: bool):
        self.text = text
        self.words = words
        self.para_after = para_after
        self.in_quote_end = in_quote_end


def _count_words(text: str) -> int:
    """Words in TEXT; non-spaced-script characters each count as one."""
    return len(text.split()) + len(_NON_SPACED_RUN.findall(text))


def _is_abbreviation(text: str, dot: int) -> bool:
    """True when the '.' before which a sentence split is being
    considered belongs to an abbreviation or a name initial (Mr. St.
    U.S. J.), where a break would fall mid-sentence."""
    start = dot
    while start > 0 and (text[start - 1].isalnum()
                         or text[start - 1] in ".-'’"):
        start -= 1
    token = text[start:dot]
    if not token:
        return False
    tail = token.split(".")[-1]
    if len(tail) == 1 and tail.isalpha():
        return True
    return token.lower() in _ABBREVIATIONS


def _sentence_continues(text: str, end: int, attach: bool) -> bool:
    """True when the text after a terminator belongs to the same unit.

    Only a closing quote or an ellipsis can pull a lowercase word into
    the unit (dialogue tags: ``“Come here!” she said.``); a plain
    period always ends the sentence, whatever follows.
    """
    match = re.compile(r"\S").search(text, end)
    if match is None:
        return False
    return attach and match.group().islower()


def _scan_units(text: str) -> List[_Unit]:
    """Split TEXT into sentence units, tracking quotes and paragraphs.

    Splits at sentence terminators (.!?…) plus any closing quotes or
    brackets that trail them, suppressing false boundaries for
    abbreviations, initials, and lowercase continuations. Blank lines
    close a paragraph (and reset the quotation state, so one unclosed
    quote cannot silence clean boundaries for the rest of a book).
    Non-whitespace content is preserved; only whitespace between
    boundaries may be normalized.
    """
    units: List[_Unit] = []
    seg_start = 0
    in_quote = False
    depth = 0

    def emit(end: int, para_after: bool) -> None:
        nonlocal seg_start
        segment = text[seg_start:end]
        seg_start = end
        if not segment.strip():
            return
        units.append(_Unit(segment, _count_words(segment), para_after,
                           in_quote or depth > 0))

    for match in _EVENT.finditer(text):
        kind = match.lastgroup
        if kind == "end":
            run = match.group()
            for ch in run[1:]:
                if ch == '"':
                    in_quote = not in_quote
                elif ch in "”»":
                    depth = max(0, depth - 1)
            end = match.end()
            nxt = text[end:end + 1]
            if nxt and not nxt.isspace() \
                    and not _NON_SPACED_RUN.match(nxt):
                continue  # glued punctuation: "3.14" is not a boundary
            if (run[0] == "." and _is_abbreviation(text, match.start())) \
                    or _sentence_continues(text, end, len(run) > 1
                                           or run[0] == "…"):
                continue
            emit(end, _PARA_AFTER.match(text, end) is not None)
        elif kind == "para":
            emit(match.start(), True)
            in_quote = False
            depth = 0
        elif kind == "open":
            depth += 1
        elif kind == "close":
            depth = max(0, depth - 1)
        elif kind == "straight":
            in_quote = not in_quote

    emit(len(text), False)
    return units


def _join(units: List[_Unit]) -> str:
    """Join units into one chunk, keeping paragraph breaks inside it."""
    parts: List[str] = []
    previous = None
    for unit in units:
        text = unit.text.strip()
        if not text:
            continue
        if parts:
            parts.append("\n\n" if previous.para_after else " ")
        parts.append(text)
        previous = unit
    return "".join(parts)


def _retreat_point(cur: List[_Unit], floor: int) -> int:
    """Index in CUR to break at when the next unit will not fit.

    Picks the largest prefix of at least FLOOR words that ends outside
    a quotation; failing that, the largest prefix of at least FLOOR
    words that ends a paragraph; failing both, the whole prefix (the
    chunk keeps as much as fits and carries the open quotation into the
    next chunk).
    """
    best_clean = 0
    best_para = 0
    words = 0
    for k, unit in enumerate(cur, 1):
        words += unit.words
        if words < floor:
            continue
        if not unit.in_quote_end:
            best_clean = k
        elif unit.para_after:
            best_para = k
    return best_clean or best_para or len(cur)


def _split_oversized(text: str, max_words: int) -> List[str]:
    """Break a sentence longer than the limit into packable pieces.

    First at clause punctuation (kept attached for TTS prosody, and
    never at commas without whitespace so tokens like "1,000,000" and
    "12:30" survive), then at word boundaries. A token with no
    whitespace at all is sliced at the hard character cap and the word
    limit so a punctuation-free run cannot exceed either bound.
    """
    parts = re.split(r"(?:(?<=[,;:])\s+|(?<=[，、；：]))", text)
    pieces: List[str] = []
    cap = max_words * _MAX_CHARS_PER_WORD
    for part in parts:
        part = part.strip()
        if not part:
            continue
        if len(part) <= cap and _count_words(part) <= max_words:
            pieces.append(part)
            continue
        pieces.extend(_hard_slices(part, max_words, cap))
    return pieces


def _slice_token(word: str, max_words: int, cap: int) -> List[str]:
    """Slice a single whitespace-free token into packable pieces.

    Cuts before the character cap is exceeded and before the word count
    is exceeded: each non-spaced-script character counts as one word,
    so a punctuation-free CJK token is bounded by the word limit too,
    not only by the character cap.
    """
    slices: List[str] = []
    cur: List[str] = []
    cur_chars = 0
    # Any slice is itself one whitespace-delimited word before its
    # non-spaced-script characters are counted on top.
    cur_words = 1
    for ch in word:
        ch_words = 1 if _NON_SPACED_RUN.match(ch) else 0
        if cur and (cur_chars + 1 > cap or cur_words + ch_words > max_words):
            slices.append("".join(cur))
            cur, cur_chars, cur_words = [], 0, 1
        cur.append(ch)
        cur_chars += 1
        cur_words += ch_words
    if cur:
        slices.append("".join(cur))
    return slices


def _hard_slices(text: str, max_words: int, cap: int) -> List[str]:
    """Slice TEXT at word boundaries, then within tokens when a single
    token (no whitespace to break at) exceeds the character cap or the
    word limit on its own. Buffers never exceed the word limit or the
    character cap, so every piece is packable."""
    pieces: List[str] = []
    buf: List[str] = []
    buf_words = 0
    buf_chars = 0
    for word in text.split():
        if len(word) > cap or _count_words(word) > max_words:
            # A single token too big to buffer: slice it at the hard
            # character cap and the word limit together.
            if buf:
                pieces.append(" ".join(buf))
                buf, buf_words, buf_chars = [], 0, 0
            pieces.extend(_slice_token(word, max_words, cap))
            continue
        words = _count_words(word)
        if buf and buf_words + words <= max_words \
                and buf_chars + 1 + len(word) <= cap:
            buf.append(word)
            buf_words += words
            buf_chars += 1 + len(word)
        else:
            if buf:
                pieces.append(" ".join(buf))
            buf, buf_words, buf_chars = [word], words, len(word)
    if buf:
        pieces.append(" ".join(buf))
    return pieces


def _joined_chars(units: List[_Unit]) -> int:
    """The character length of the text _join would produce."""
    total = 0
    previous = None
    for unit in units:
        text = unit.text.strip()
        if not text:
            continue
        if previous is not None:
            total += 2 if previous.para_after else 1
        total += len(text)
        previous = unit
    return total


def _quote_states(pieces: List[str], in_quote_end: bool) -> List[bool]:
    """The quote state after each piece of an oversized sentence.

    Pieces of one sentence can sit inside a quotation the sentence does
    not close, so ending a chunk at such a piece would sever the quote
    (and marking them clean lets the target flush cut mid-quote). The
    scanner's quote tracking is replayed over the pieces from the state
    the sentence started in; that start state is unknown, so both
    possibilities are replayed and the one whose end state matches the
    unit's known end state wins (a unit rarely starts mid-quote). When
    neither matches, every piece takes the unit's conservative end
    state.
    """
    def simulate(in_quote: bool, depth: int) -> List[bool]:
        states: List[bool] = []
        for piece in pieces:
            for ch in piece:
                if ch == '"':
                    in_quote = not in_quote
                elif ch in "“«":
                    depth += 1
                elif ch in "”»":
                    depth = max(0, depth - 1)
            states.append(in_quote or depth > 0)
        return states

    states = simulate(False, 0)
    if states and states[-1] == in_quote_end:
        return states
    states = simulate(True, 0)
    if states and states[-1] == in_quote_end:
        return states
    return [in_quote_end] * len(pieces)


def _split_smart(text: str, max_words: int) -> List[str]:
    """Smart strategy: fill chunks up to MAX_WORDS at sentence
    boundaries, end them at the target once a quotation has closed or
    a paragraph ends, and retreat cleanly instead of breaking quotes
    when the next sentence does not fit. Every chunk also respects the
    hard character cap, so punctuation-free or non-spaced input stays
    bounded."""
    target = max(1, round(max_words * _TARGET_RATIO))
    floor = max(1, round(max_words * _FLOOR_RATIO))
    char_cap = max_words * _MAX_CHARS_PER_WORD
    chunks: List[str] = []
    cur: List[_Unit] = []
    cur_words = 0
    cur_chars = 0

    def flush() -> None:
        nonlocal cur, cur_words, cur_chars
        if cur:
            chunks.append(_join(cur))
            cur, cur_words, cur_chars = [], 0, 0

    for unit in _scan_units(text):
        if cur and cur_words >= target \
                and (not cur[-1].in_quote_end or cur[-1].para_after):
            flush()
        if unit.words <= max_words \
                and len(unit.text.strip()) <= char_cap \
                and _joined_chars(cur + [unit]) <= char_cap:
            if cur_words + unit.words <= max_words:
                cur.append(unit)
                cur_words += unit.words
                cur_chars = _joined_chars(cur)
                continue
            # The sentence does not fit: break earlier if a clean
            # boundary keeps the chunk above the floor.
            point = _retreat_point(cur, floor)
            if point < len(cur):
                chunks.append(_join(cur[:point]))
                cur = cur[point:]
                cur_words = sum(u.words for u in cur)
                cur_chars = _joined_chars(cur)
                if cur_words + unit.words <= max_words:
                    cur.append(unit)
                    cur_words += unit.words
                    cur_chars = _joined_chars(cur)
                    continue
            flush()
            cur.append(unit)
            cur_words = unit.words
            cur_chars = _joined_chars(cur)
            continue
        # One piece exceeds the limit itself: send its pieces out
        # packed as tightly as the pieces allow, carrying the
        # sentence's quotation state through the fallback so a chunk
        # can end at a clean in-sentence boundary — and a piece still
        # inside a quote is never mistaken for one.
        flush()
        piece_texts = _split_oversized(unit.text, max_words)
        quote_states = _quote_states(piece_texts, unit.in_quote_end)
        for index, piece_text in enumerate(piece_texts):
            if _count_words(piece_text) > max_words \
                    or len(piece_text) > char_cap:
                # A character-sliced piece of degenerate input: its own
                # chunk (already at the hard cap).
                flush()
                chunks.append(piece_text)
                continue
            piece = _Unit(
                piece_text, _count_words(piece_text),
                unit.para_after and index == len(piece_texts) - 1,
                quote_states[index])
            if cur and cur_words + piece.words <= max_words \
                    and cur_chars + 1 + len(piece_text) <= char_cap:
                cur.append(piece)
                cur_words += piece.words
                cur_chars += 1 + len(piece_text)
                continue
            flush()
            cur.append(piece)
            cur_words = piece.words
            cur_chars = len(piece_text)
    flush()
    return [chunk for chunk in chunks if chunk.strip()]


def _split_legacy(text: str, max_words: int) -> List[str]:
    """The pre-smart splitter, kept intact for the Smart chunking=off
    setting. Chunks pack whole sentences up to MAX_WORDS; sentences
    longer than the limit split at clause punctuation, then at word
    boundaries."""
    sentences = re.split(r"(?<=[.!?])\s+", text)
    chunks = []
    current_chunk = ""
    current_words = 0

    for sentence in sentences:
        sentence_words = len(sentence.split())

        if sentence_words > max_words:
            if current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = ""
                current_words = 0

            # Split long sentences at clause boundaries, keeping punctuation.
            # Only split where whitespace already follows the punctuation so
            # tokens are never broken apart or re-joined with added spaces
            # (no spaces are injected into "1,000,000" or "12:30").
            parts = re.split(r"(?<=[,;:])\s+", sentence)
            for part in parts:
                part_words = len(part.split())
                if part_words > max_words:
                    # Last resort: no punctuation split point is available,
                    # so split at word boundaries. Tokens themselves (and
                    # therefore numbers like "1,000,000") stay intact.
                    if current_chunk:
                        chunks.append(current_chunk.strip())
                        current_chunk = ""
                        current_words = 0
                    words = part.split()
                    for start in range(0, len(words), max_words):
                        chunks.append(" ".join(words[start:start + max_words]))
                    continue
                if current_words + part_words <= max_words:
                    current_chunk += part + " "
                    current_words += part_words
                else:
                    if current_chunk:
                        chunks.append(current_chunk.strip())
                    current_chunk = part + " "
                    current_words = part_words
        else:
            if current_words + sentence_words <= max_words:
                current_chunk += sentence + " "
                current_words += sentence_words
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = sentence + " "
                current_words = sentence_words

    if current_chunk.strip():
        chunks.append(current_chunk.strip())

    return [chunk for chunk in chunks if chunk.strip()]


def split_into_chunks(text: str, max_words: Optional[int] = None,
                      smart: Optional[bool] = None) -> List[str]:
    """Split text into chunks of at most ``max_words`` words.

    ``max_words`` defaults to ``config.CHUNK_SIZE`` (read at call time).
    There is no ceiling beyond that setting, but note that the TTS
    servers silently truncate audio when a single generation runs too
    long without reporting an error, so very large values are at your
    own risk (see CHUNK_SIZE in app/converter/config.py).

    With smart chunking on (``config.SMART_CHUNKING``, the default), a
    chunk crossing the target (~85% of the limit) ends at the next
    sentence end outside a quotation or at a paragraph break, so most
    requests stop at natural boundaries instead of pressing against the
    limit. When the next sentence will not fit, the break retreats to
    the last boundary that ends outside a quotation (never smaller than
    ~70% of the limit), so short quotations stay in one request when
    possible; very long or unclosed quotes still split, carrying their
    open state into the next chunk. Malformed punctuation never grows a
    chunk beyond the limit: every break consumes text, and input without
    sentence punctuation or whitespace (some CJK text, extraction
    damage) falls back to clause punctuation, then word boundaries,
    then a hard character cap. Content is preserved: no closing quote
    or punctuation is ever invented or dropped.

    Both strategies split on sentence boundaries. Sentences longer than
    the limit are split further at clause punctuation (which is kept
    attached for TTS prosody). Clause splits only happen at whitespace
    after punctuation, so tokens like "1,000,000" or "12:30" are never
    broken apart.
    """
    if max_words is None:
        max_words = config.CHUNK_SIZE
    if max_words < 1:
        max_words = 1

    if not text.strip():
        return []

    if smart is None:
        smart = bool(getattr(config, "SMART_CHUNKING", False))
    if smart:
        return _split_smart(text, max_words)
    return _split_legacy(text, max_words)