unios-tts/app/ssml.py

from __future__ import annotations

import re
from dataclasses import dataclass

# Sentinel, который нормализатор оставляет вместо удалённого выделительного
# разделителя (пауза фиксированной длины).
PAUSE_DELIM = "\uE000"
DEFAULT_PAUSE_MS = 150
MAX_BREAK_MS = 5000

# Явная пользовательская пауза кодируется ОДНИМ символом из PUA:
# код = EXPLICIT_BASE + длительность_в_мс (0..MAX_BREAK_MS). Одно кодо-point
# переживает rutextnorm/translit (в отличие от длинных прогонов и цифр).
_EXPLICIT_BASE = 0xE100

RATE_VALUES = ("x-slow", "slow", "medium", "fast", "x-fast")
PITCH_VALUES = ("x-low", "low", "medium", "high", "x-high")


def _escape(text: str) -> str:
    return (
        text.replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
    )


def explicit_pause_char(ms: int) -> str:
    ms = max(0, min(int(ms), MAX_BREAK_MS))
    return chr(_EXPLICIT_BASE + ms)


def _is_sentinel(char: str) -> bool:
    return _segment_to_break(char) is not None


def strip_sentinels(text: str, replacement: str = "") -> str:
    """Убирает sentinel-символы (для отображения нормализованного текста)."""
    out = [replacement if _is_sentinel(ch) else ch for ch in text]
    return re.sub(r" {2,}", " ", "".join(out)).strip()


def _segment_to_break(char: str) -> int | None:
    """Длительность паузы (мс) для sentinel-символа или None, если это текст."""
    code = ord(char)
    if char == PAUSE_DELIM:
        return DEFAULT_PAUSE_MS
    if _EXPLICIT_BASE <= code <= _EXPLICIT_BASE + MAX_BREAK_MS:
        return code - _EXPLICIT_BASE
    return None


def _sentence_content(text: str) -> str:
    """Текст предложения с заменой sentinel-пауз на <break .../>; сам sentinel
    вырезается, XML-спецсимволы экранируются."""
    parts: list[str] = []
    buf: list[str] = []
    for ch in text:
        ms = _segment_to_break(ch)
        if ms is None:
            buf.append(ch)
            continue
        if buf:
            parts.append(_escape("".join(buf)))
            buf = []
        ms = max(0, min(ms, MAX_BREAK_MS))
        parts.append(f'<break time="{ms}ms"/>')
    if buf:
        parts.append(_escape("".join(buf)))
    return "".join(parts)


def _para_block(segments) -> str:
    inner = "".join(f"<s>{_sentence_content(s.text)}</s>" for s in segments)
    return f"<p>{inner}</p>"


def _content_length(segments) -> int:
    total = 0
    for s in segments:
        total += sum(1 for ch in s.text if _segment_to_break(ch) is None)
    return total


def build_ssml(segments, *, rate: str | None = None, pitch: str | None = None) -> str:
    """Собирает один <speak>-документ: предложения -> <s>, группы предложений
    одного абзаца (одинаковый Segment.para) -> <p>. rate/pitch (опционально)
    оборачивают всё содержимое одним <prosody>."""
    if not segments:
        return ""

    groups: list[list] = []
    for seg in segments:
        key = (seg.lang, getattr(seg, "para", 0))
        if groups and (
            (groups[-1][0].lang, getattr(groups[-1][0], "para", 0)) == key
        ):
            groups[-1].append(seg)
        else:
            groups.append([seg])

    body = "".join(_para_block(g) for g in groups)

    attrs: list[str] = []
    if rate and rate != "medium":
        attrs.append(f'rate="{rate}"')
    if pitch and pitch != "medium":
        attrs.append(f'pitch="{pitch}"')
    if attrs:
        body = f"<prosody {' '.join(attrs)}>{body}</prosody>"

    return f"<speak>{body}</speak>"


def build_ssml_chunks(segments, *, rate: str | None = None, pitch: str | None = None,
                      max_chars: int = 5000) -> list[str]:
    """Режет сегменты по границам абзацев на документы <= max_chars символов
    контента (параграф = атомарный чанк; гигантский параграф режется по
    предложениям). Возвращает список <speak>-документов."""
    if not segments:
        return []

    paragraphs: list[list] = []
    for seg in segments:
        key = (seg.lang, getattr(seg, "para", 0))
        if paragraphs and (
            (paragraphs[-1][0].lang, getattr(paragraphs[-1][0], "para", 0)) == key
        ):
            paragraphs[-1].append(seg)
        else:
            paragraphs.append([seg])

    chunks: list[list] = []
    current: list = []
    current_len = 0
    for para in paragraphs:
        para_len = _content_length(para)
        if para_len > max_chars:
            if current:
                chunks.append(current)
                current = []
                current_len = 0
            piece: list = []
            piece_len = 0
            for seg in para:
                seg_len = _content_length([seg])
                if piece and piece_len + seg_len > max_chars:
                    chunks.append(piece)
                    piece = []
                    piece_len = 0
                piece.append(seg)
                piece_len += seg_len
            if piece:
                current = piece
                current_len = piece_len
            continue
        if current and current_len + para_len > max_chars:
            chunks.append(current)
            current = []
            current_len = 0
        current.extend(para)
        current_len += para_len
    if current:
        chunks.append(current)

    return [build_ssml(ch, rate=rate, pitch=pitch) for ch in chunks]