call-center/services/ai_voice_runtime_service/ack_bank.py

from __future__ import annotations

import contextlib
import hashlib
import json
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from threading import Lock

from services.ai_voice_runtime_service.audiosocket import resample_pcm16le
from services.ai_voice_runtime_service.providers.tts import TTSProvider


def _ack_bank_dir() -> Path:
    explicit = str(os.getenv("AI_VOICE_V2_PREBAKED_ACK_DIR", "") or "").strip()
    if explicit:
        return Path(explicit).expanduser()
    data_dir = str(os.getenv("CC_DATA_DIR", "") or "").strip()
    if data_dir:
        return Path(data_dir).expanduser() / "voice_ack_bank"
    local_data_dir = Path(".data_local")
    if local_data_dir.exists():
        return local_data_dir / "voice_ack_bank"
    return Path(".data") / "voice_ack_bank"


@dataclass(slots=True)
class AckClip:
    text: str
    pcm_8k_bytes: bytes
    source: str
    sample_rate_hz: int = 8000


class PrebakedAckBank:
    def __init__(self, *, tts_provider: TTSProvider, ack_dir: Path | None = None) -> None:
        self._tts_provider = tts_provider
        self._ack_dir = ack_dir or _ack_bank_dir()
        self._lock = Lock()

    def _cache_key(
        self,
        *,
        text: str,
        language: str | None,
        style_hints: dict[str, object] | None,
    ) -> str:
        payload = {
            "provider": getattr(self._tts_provider, "name", "tts"),
            "voice_fingerprint": self._tts_provider.cache_fingerprint(language, style_hints=style_hints),
            "language": str(language or "").strip() or None,
            "style_hints": style_hints or {},
            "text": text,
        }
        encoded = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
        return hashlib.sha256(encoded).hexdigest()

    def _paths(self, cache_key: str) -> tuple[Path, Path]:
        prefix = self._ack_dir / cache_key[:2] / cache_key[2:4]
        return prefix / f"{cache_key}.pcm", prefix / f"{cache_key}.json"

    def _load(self, cache_key: str, *, text: str) -> AckClip | None:
        pcm_path, meta_path = self._paths(cache_key)
        if not pcm_path.exists():
            return None
        try:
            pcm_bytes = pcm_path.read_bytes()
            source = "prebaked_cache"
            if meta_path.exists():
                metadata = json.loads(meta_path.read_text(encoding="utf-8"))
                source = str(metadata.get("source") or source)
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return None
        if not pcm_bytes:
            return None
        return AckClip(text=text, pcm_8k_bytes=pcm_bytes, source=source)

    def _store(self, cache_key: str, *, clip: AckClip, language: str | None, style_hints: dict[str, object] | None) -> None:
        pcm_path, meta_path = self._paths(cache_key)
        pcm_path.parent.mkdir(parents=True, exist_ok=True)
        metadata = {
            "provider": getattr(self._tts_provider, "name", "tts"),
            "language": str(language or "").strip() or None,
            "style_hints": style_hints or {},
            "text": clip.text,
            "sample_rate_hz": clip.sample_rate_hz,
            "source": clip.source,
        }
        pcm_tmp: str | None = None
        meta_tmp: str | None = None
        try:
            with tempfile.NamedTemporaryFile(dir=pcm_path.parent, delete=False, suffix=".pcm.tmp") as handle:
                handle.write(clip.pcm_8k_bytes)
                pcm_tmp = handle.name
            with tempfile.NamedTemporaryFile(dir=meta_path.parent, delete=False, suffix=".json.tmp", mode="w", encoding="utf-8") as handle:
                json.dump(metadata, handle, ensure_ascii=False, sort_keys=True)
                meta_tmp = handle.name
            os.replace(pcm_tmp, pcm_path)
            os.replace(meta_tmp, meta_path)
        finally:
            for path in (pcm_tmp, meta_tmp):
                if path and os.path.exists(path):
                    with contextlib.suppress(OSError):
                        os.unlink(path)

    def get_clip(
        self,
        *,
        text: str,
        language: str | None,
        style_hints: dict[str, object] | None = None,
    ) -> AckClip:
        normalized_text = str(text or "").strip()
        if not normalized_text:
            return AckClip(text="", pcm_8k_bytes=b"", source="empty")
        cache_key = self._cache_key(text=normalized_text, language=language, style_hints=style_hints)
        with self._lock:
            cached = self._load(cache_key, text=normalized_text)
            if cached is not None:
                return cached
            synthesis = self._tts_provider.synthesize(
                normalized_text,
                language=language,
                style_hints=style_hints,
            )
            pcm_8k = resample_pcm16le(
                synthesis.audio_bytes,
                input_rate_hz=synthesis.sample_rate_hz,
                output_rate_hz=8000,
            )
            clip = AckClip(
                text=normalized_text,
                pcm_8k_bytes=pcm_8k,
                source="prebaked_materialized",
            )
            self._store(cache_key, clip=clip, language=language, style_hints=style_hints)
            return clip