UmbraAudio/umbraaudio/generation/staleness.py

"""Utterance-audio staleness guard.

Per-utterance audio is keyed by index (utt_XXXX.wav). Indexes stay stable
across re-segmentation while the underlying text changes, so a "file exists"
check silently reuses audio that no longer matches the utterance text. Each
wav gets a sidecar marker (utt_XXXX.wav.txt) recording the exact spoken text
AND the assigned speaker; a file is only considered current when both match.

Markers are JSON ``{"speaker": ..., "text": ..., "voice_fp": ..., "instruct": ...,
"sha256": ...}``. Legacy plain-text markers carry no speaker information, so they are
always treated as stale.  The ``voice_fp`` field is a truncated SHA-256 of the
voice-reference wav used during generation; when a speaker's ref changes
(e.g. gender re-assignment) the fingerprint no longer matches and the file
is considered stale, forcing regeneration.

``sha256`` is the FULL content hash of the rendered wav, recorded at generation
(specs/21 §9.1 M4-04 / controller §5.2: "every blob gets its hash at generation, not at
export"). It is recorded but **not** verified by default: hashing ~10k files on every
validation pass would re-read gigabytes. Pass ``verify_sha256=True`` when an integrity
check is actually wanted.
"""

from __future__ import annotations

import hashlib
import json
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Sequence
    from pathlib import Path

# Spec 31 section 10. A marker written before the language tag existed records no
# ``language``; that audio was English (the only language the pipeline could
# make), so an absent value means this. Treating absence as "unknown" instead
# would invalidate every existing clip - ~4,300 files and a full regeneration -
# for no gain.
DEFAULT_LANGUAGE = "en"


def marker_path(wav_path: Path) -> Path:
    """Sidecar marker path for a wav file (utt_0000.wav -> utt_0000.wav.txt)."""
    return wav_path.with_suffix(".wav.txt")


def _voice_fingerprint(voice_ref: Path) -> str:
    """Truncated SHA-256 fingerprint of a voice-reference wav file."""
    h = hashlib.sha256(voice_ref.read_bytes()).hexdigest()
    return h[:16]


def file_sha256(path: Path) -> str | None:
    """Full SHA-256 of a file, or None if it cannot be read."""
    try:
        return hashlib.sha256(path.read_bytes()).hexdigest()
    except OSError:
        return None


def is_current(
    text: str,
    wav_path: Path,
    speaker: str | None = None,
    voice_ref: Path | None = None,
    voice_mode: str | None = None,
    accepted_voice_modes: set[str] | None = None,
    instruct: str | None = None,
    fx_version: int | None = None,
    verify_sha256: bool = False,
    language: str | None = None,
) -> bool:
    """True when the wav exists, its marker is JSON, and text (+ speaker if
    given) match exactly.  If *voice_ref* is provided the marker's
    ``voice_fp`` must also match the current ref's fingerprint; when no
    ref is given the fingerprint check is skipped (backward compatible).
    If *instruct* is provided the marker's ``instruct`` must match too
    (a changed emotion/variation therefore invalidates old audio).

    Mode handling: if *voice_mode* is given (e.g. ``"icl"``) the marker must
    record that exact mode; otherwise if *accepted_voice_modes* is given the
    marker's mode just needs to be one of the accepted set (e.g.
    ``{"embedding", "icl"}``) so a historical mixed-mode chapter is treated as
    current. When *neither* is given the mode check is skipped.

    This lets a generation-technology change (e.g. embedding -> a brand-new
    mode) invalidate previously generated audio, while tolerating equivalent
    voice-identity modes (embedding/icl) that share the same ref fingerprint."""
    if not wav_path.exists():
        return False
    marker = marker_path(wav_path)
    if not marker.exists():
        return False
    try:
        payload = json.loads(marker.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return False
    if not isinstance(payload, dict):
        return False
    if payload.get("text") != text:
        return False
    if speaker is not None and payload.get("speaker") != speaker:
        return False
    if instruct is not None and payload.get("instruct") != instruct:
        return False
    if language is not None:
        # Absence means English (DEFAULT_LANGUAGE), so a legacy English clip
        # stays current while a line retagged as another language - or an
        # English line whose clip was made in another language - goes stale.
        recorded_language = payload.get("language") or DEFAULT_LANGUAGE
        if recorded_language != language:
            return False
    if voice_ref is not None:
        if not voice_ref.exists():
            return False
        expected_fp = _voice_fingerprint(voice_ref)
        if payload.get("voice_fp") != expected_fp:
            return False
    recorded = payload.get("voice_mode")
    if voice_mode is not None and recorded != voice_mode:
        return False
    if (
        voice_mode is None
        and accepted_voice_modes is not None
        and recorded not in accepted_voice_modes
    ):
        return False
    if fx_version is not None and payload.get("fx_version") != fx_version:
        return False
    if verify_sha256:
        recorded = payload.get("sha256")
        if not recorded or recorded != file_sha256(wav_path):
            return False
    return not (fx_version is None and payload.get("fx_version") is not None)


def write_marker(
    text: str,
    wav_path: Path,
    speaker: str | None = None,
    voice_ref: Path | None = None,
    voice_mode: str | None = None,
    instruct: str | None = None,
    fx_version: int | None = None,
    source_unit_key: str | None = None,
    effects: Sequence[str] | None = None,
    language: str | None = None,
) -> None:
    """Record *text*, *speaker*, and (optionally) the voice-reference
    fingerprint as the spoken content of *wav_path*. When *instruct* is given
    (the exact TTS instruction used) it is recorded too, so a later run with a
    different emotion/instruct regenerates instead of reusing the file.
    *fx_version* records which post-processing emotion recipe (if any) was
    stamped onto the audio; a bump forces regeneration of fx'd files.

    The rendered wav's own full SHA-256 is always recorded when the file exists,
    so every generated blob carries its content hash from the moment it is born
    rather than at export time."""
    payload: dict[str, str] = {"speaker": speaker or "", "text": text}
    if wav_path.exists():
        digest = file_sha256(wav_path)
        if digest:
            payload["sha256"] = digest
    if voice_ref is not None and voice_ref.exists():
        payload["voice_fp"] = _voice_fingerprint(voice_ref)
    if voice_mode is not None:
        payload["voice_mode"] = voice_mode
    if instruct is not None:
        payload["instruct"] = instruct
    if fx_version is not None:
        payload["fx_version"] = fx_version
    if source_unit_key is not None:
        payload["source_unit_key"] = source_unit_key
    if language is not None:
        # Recorded at generation so a language tag applied *later* (spec 31
        # section 9) invalidates the clip rather than leaving it with English
        # phonology behind an unchanged text.
        payload["language"] = language
    if effects:
        # The recorded sha256 identifies the FINAL bytes. Without this, nothing records
        # which post-processing turned the raw TTS output into those bytes, so a changed
        # trait list would change the hash with no explanation (spec 21 §12 gap G-A).
        payload["effects"] = list(effects)
    marker_path(wav_path).write_text(
        json.dumps(payload, ensure_ascii=False), encoding="utf-8"
    )