dude_tts/interleave.py

"""Dialogue text frontend: XML -> dependency DAG -> dense, ordered lanes.

No timestamps or XML tokens enter the lanes. Start anchors are token boundaries.
Same-channel turns are serialized in XML order; a delayed turn carries all of
its anchors with it, so descendant insertions follow the actual placement.
"""
from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
import heapq
import re
from typing import Generic, TypeVar
import xml.etree.ElementTree as ET

T = TypeVar("T")
ANCHOR = re.compile(r"e[0-9]+\Z")
EVENTS = frozenset({"laugh", "cough", "breath", "sigh", "cry", "sneeze", "gasp"})


@dataclass(frozen=True)
class Part:
    kind: str  # text, anchor, or event
    value: str


@dataclass(frozen=True)
class Turn:
    index: int
    channel: str
    start: str | None
    parts: tuple[Part, ...]


@dataclass(frozen=True)
class Edge:
    source: int
    target: int
    anchor: str | None  # None means the preceding same-channel turn's end.


@dataclass(frozen=True)
class TranscriptDAG:
    turns: tuple[Turn, ...]
    anchors: Mapping[str, int]  # Anchor name -> owner turn index.
    edges: tuple[Edge, ...]
    order: tuple[int, ...]
    root: int


@dataclass(frozen=True)
class Placement:
    turn: int
    channel: str
    requested_start: int
    start: int
    end: int  # Exclusive; all positions count A/B pairs, starting at zero.


@dataclass(frozen=True)
class Interleaving(Generic[T]):
    lanes: Mapping[str, tuple[T, ...]]
    occupied: Mapping[str, tuple[bool, ...]]
    placements: tuple[Placement, ...]
    anchors: Mapping[str, int]
    content_origin: int

    @property
    def tokens(self) -> tuple[T, ...]:
        return tuple(token for pair in zip(self.lanes["A"], self.lanes["B"]) for token in pair)

    @property
    def occupied_mask(self) -> tuple[bool, ...]:
        return tuple(value for pair in zip(self.occupied["A"], self.occupied["B"]) for value in pair)


def parse_xml(xml: str, *, max_turns: int = 10000) -> TranscriptDAG:
    """Accept a <duplex> document or an A/B fragment; validate before layout.

    Exactly one turn lacks start. Every other turn references a unique e<number>
    marker. Forward references and branching are allowed when the resulting
    anchor + same-channel ordering graph is acyclic. Turn attributes other than
    start (including type) are rejected instead of silently ignored.
    """
    if len(xml) > 2_000_000:
        raise ValueError("Transcript exceeds the XML size limit")
    if re.search(r"<!\s*(?:DOCTYPE|ENTITY)\b", xml, re.IGNORECASE):
        raise ValueError("DTD and entity declarations are not allowed")
    try:
        document = ET.fromstring(xml)
    except ET.ParseError:
        try:
            document = ET.fromstring("<duplex>" + xml + "</duplex>")
        except ET.ParseError as error:
            raise ValueError(f"Invalid transcript XML: {error}") from error
    if document.tag in {"A", "B"}:
        elements = [document]
    elif document.tag == "duplex" and not document.attrib:
        if (document.text or "").strip():
            raise ValueError("Text must be inside an A or B turn")
        elements = list(document)
    else:
        raise ValueError("Expected A/B turns or an attribute-free duplex wrapper")
    if not elements or len(elements) > max_turns:
        raise ValueError("Expected between 1 and max_turns turns")
    turns, anchors = [], {}
    for index, element in enumerate(elements):
        if element.tag not in {"A", "B"} or set(element.attrib) - {"start"}:
            raise ValueError(f"Turn {index}: only A/B and the start attribute are allowed")
        if (element.tail or "").strip():
            raise ValueError("Text must be inside an A or B turn")
        start = element.get("start")
        if start is not None and not ANCHOR.fullmatch(start):
            raise ValueError(f"Turn {index}: start must be e<number>")
        parts = [Part("text", element.text or "")]
        for child in element:
            if child.attrib or len(child) or (child.text or "").strip():
                raise ValueError(f"Turn {index}: anchors/events must be empty and attribute-free")
            if ANCHOR.fullmatch(child.tag):
                if child.tag in anchors:
                    raise ValueError(f"Duplicate anchor: {child.tag}")
                anchors[child.tag] = index
                parts.append(Part("anchor", child.tag))
            elif child.tag in EVENTS:
                parts.append(Part("event", child.tag))
            else:
                raise ValueError(f"Turn {index}: unsupported inline tag {child.tag!r}")
            parts.append(Part("text", child.tail or ""))
        # Ignore formatting outside the utterance, preserve all internal text.
        parts[0] = Part("text", parts[0].value.lstrip())
        parts[-1] = Part("text", parts[-1].value.rstrip())
        turns.append(Turn(index, element.tag, start, tuple(parts)))
    roots = [turn.index for turn in turns if turn.start is None]
    if len(roots) != 1:
        raise ValueError("Exactly one root turn must omit start; all others need start=e<number>")
    edges, previous = [], {}
    for turn in turns:
        if turn.start is not None:
            if turn.start not in anchors:
                raise ValueError(f"Unknown start anchor: {turn.start}")
            edges.append(Edge(anchors[turn.start], turn.index, turn.start))
        if turn.channel in previous:
            edges.append(Edge(previous[turn.channel], turn.index, None))
        previous[turn.channel] = turn.index
    children = [[] for _ in turns]
    indegree = [0] * len(turns)
    for edge in edges:
        children[edge.source].append(edge.target)
        indegree[edge.target] += 1
    ready = [i for i, degree in enumerate(indegree) if degree == 0]
    heapq.heapify(ready)
    order = []
    while ready:
        node = heapq.heappop(ready)
        order.append(node)
        for child in children[node]:
            indegree[child] -= 1
            if indegree[child] == 0:
                heapq.heappush(ready, child)
    if len(order) != len(turns):
        raise ValueError("Cycle in anchor dependencies or same-channel XML order")
    return TranscriptDAG(tuple(turns), anchors, tuple(edges), tuple(order), roots[0])


def expand_dag(
    dag: TranscriptDAG,
    encode_text: Callable[[str], Sequence[T]],
    pad_token: T,
    *,
    encode_event: Callable[[str], Sequence[T]] | None = None,
    prefixes: Mapping[str, Sequence[T]] | None = None,
    max_lane_tokens: int = 65536,
) -> Interleaving[T]:
    """Place whole turns at max(anchor position, previous same-lane end).

    Encode text spans separately to make every anchor an exact token boundary;
    the encoder must disable BOS/EOS insertion. No whitespace is inserted or
    removed between spans. Event tags require an explicit semantic encoder;
    they are never passed as XML or spoken tag names to the text encoder.

    Optional already-tokenized native/system prefixes are left-padded to the
    same length, preserving A-odd/B-even parity and one shared content origin.
    Structural PAD is distinguished from real tokens by occupied_mask.
    """
    prefixes = {} if prefixes is None else prefixes
    if set(prefixes) - {"A", "B"}:
        raise ValueError("Prefix channels must be A or B")
    prefix = {channel: tuple(prefixes.get(channel, ())) for channel in ("A", "B")}
    origin = max(map(len, prefix.values()))
    if max_lane_tokens < 1 or origin > max_lane_tokens:
        raise ValueError("Prefix exceeds the lane token limit")
    encoded, local_anchors = [], {}
    for turn in dag.turns:
        tokens = []
        for part in turn.parts:
            if part.kind == "anchor":
                local_anchors[part.value] = len(tokens)
            elif part.kind == "event":
                if encode_event is None:
                    raise ValueError(f"Event {part.value!r} requires an explicit event encoder")
                event_tokens = list(encode_event(part.value))
                if not event_tokens:
                    raise ValueError(f"Event {part.value!r} encoded to no tokens")
                tokens.extend(event_tokens)
            elif part.value:
                tokens.extend(encode_text(part.value))
            if len(tokens) + origin > max_lane_tokens:
                raise ValueError("Turn exceeds the lane token limit")
        if not tokens:
            raise ValueError(f"Turn {turn.index} contains no content tokens")
        encoded.append(tuple(tokens))
    incoming = [[] for _ in dag.turns]
    for edge in dag.edges:
        incoming[edge.target].append(edge)
    placements = {}
    anchors = {}
    for index in dag.order:
        turn = dag.turns[index]
        requested = origin if turn.start is None else anchors[turn.start]
        start = requested
        for edge in incoming[index]:
            if edge.anchor is None:
                start = max(start, placements[edge.source].end)
        end = start + len(encoded[index])
        if end > max_lane_tokens:
            raise ValueError("Expanded transcript exceeds the lane token limit")
        placements[index] = Placement(index, turn.channel, requested, start, end)
        for part in turn.parts:
            if part.kind == "anchor":
                anchors[part.value] = start + local_anchors[part.value]
    length = max(place.end for place in placements.values())
    lanes = {ch: [pad_token] * length for ch in ("A", "B")}
    occupied = {ch: [False] * length for ch in ("A", "B")}
    for ch in lanes:
        left = origin - len(prefix[ch])
        lanes[ch][left:origin] = prefix[ch]
        occupied[ch][left:origin] = [True] * len(prefix[ch])
    for index, place in placements.items():
        lanes[place.channel][place.start:place.end] = encoded[index]
        occupied[place.channel][place.start:place.end] = [True] * len(encoded[index])
    return Interleaving(
        {ch: tuple(tokens) for ch, tokens in lanes.items()},
        {ch: tuple(mask) for ch, mask in occupied.items()},
        tuple(placements[i] for i in range(len(dag.turns))), anchors, origin,
    )


def interleave_xml(xml: str, encode_text: Callable[[str], Sequence[T]], pad_token: T,
                   **kwargs) -> Interleaving[T]:
    return expand_dag(parse_xml(xml), encode_text, pad_token, **kwargs)