tts/utils/text.py

"""Text segmentation utilities for mantra-aware synthesis.

Mantras are identified as ALL-CAPS runs of two or more characters.
The segmentation pipeline has two stages:

1. inline_single_caps()
   Isolated single ALL-CAPS words (e.g. "the TAM syllable") are replaced
   with their phonetic spelling IN PLACE. They remain part of the surrounding
   prose and are synthesized at prose speed.

2. split_segments()
   Multi-word ALL-CAPS runs (e.g. OM AH HUM) are split out as dedicated
   mantra segments and synthesized at mantra speed.

Always call inline_single_caps() BEFORE split_segments().

ElevenLabs markup tags — `<break time="1.5s" />` (v2/Flash) and bracketed
audio tags like `[pause]`, `[laughs]` (v3) — must reach the API byte-for-byte.
_TAG_SPAN marks those spans so mantra detection, sentence splitting, and
glossary substitution all treat them as opaque.
"""

import re
from ..lexicon.mantra_syllables import MANTRA_SYLLABLES

# Matches an ElevenLabs markup tag span: <break time="1.5s" /> or [pause].
_TAG_SPAN = re.compile(r'<[^>\n]*>|\[[^\]\n]*\]')

# Matches one or more consecutive ALL-CAPS words (minimum 2 characters each),
# or a tag span (which is skipped, not touched, by both functions below).
_TAG_OR_MANTRA = re.compile(
    r'(?P<tag>' + _TAG_SPAN.pattern + r')'
    r'|(?P<mantra>\b[A-Z]{2,}(?:\s+[A-Z]{2,})*\b)'
)


def find_tag_spans(text: str) -> list[tuple[int, int]]:
    """Return (start, end) index ranges of every ElevenLabs markup tag in text."""
    return [m.span() for m in _TAG_SPAN.finditer(text)]


def protect_tags(text: str, fn) -> str:
    """Apply fn to the parts of text outside markup tag spans; tags pass through unchanged."""
    parts = []
    last = 0
    for start, end in find_tag_spans(text):
        if start > last:
            parts.append(fn(text[last:start]))
        parts.append(text[start:end])
        last = end
    if last < len(text):
        parts.append(fn(text[last:]))
    return "".join(parts)


def inline_single_caps(text: str) -> str:
    """Replace isolated single ALL-CAPS words with their phonetic spelling.

    Multi-word ALL-CAPS runs (e.g. OM AH HUM) are left intact so that
    split_segments() can handle them as dedicated mantra segments. Markup
    tag spans (<break .../>, [pause], etc.) are left untouched.

    Example:
        "The TAM syllable rotates."
        → "The tahm syllable rotates."   (stays in prose flow)

        "Recite OM AH HUM three times."
        → unchanged                       (left for split_segments)
    """
    def _replace(m):
        if m.group("tag") is not None:
            return m.group("tag")
        mantra = m.group("mantra")
        words = mantra.split()
        if len(words) == 1:
            # Single word: use phonetic form, or lowercase as fallback
            return MANTRA_SYLLABLES.get(mantra, mantra.lower())
        # Multi-word run: leave for split_segments()
        return mantra

    return _TAG_OR_MANTRA.sub(_replace, text)


def split_segments(text: str) -> list[tuple[str, bool]]:
    """Split text into (segment, is_mantra) pairs, preserving order.

    Multi-word ALL-CAPS runs become mantra segments (is_mantra=True).
    Everything else, including markup tag spans, is prose (is_mantra=False).

    Call inline_single_caps() first so single-word caps are already
    converted and won't be incorrectly split out here.

    Example:
        "Recite OM AH HUM three times."
        → [("Recite ", False), ("OM AH HUM", True), (" three times.", False)]
    """
    segments = []
    last = 0
    for m in _TAG_OR_MANTRA.finditer(text):
        if m.group("tag") is not None:
            continue  # tag spans stay merged into the surrounding prose
        if m.start() > last:
            segments.append((text[last:m.start()], False))
        segments.append((m.group("mantra"), True))
        last = m.end()
    if last < len(text):
        segments.append((text[last:], False))
    return segments or [(text, False)]