esp32-server/app/audio/opus.py

"""Opus encode/decode — lazy (CON-002).

Device-facing codec for the Xiaozhi bridge:

* :class:`OpusCodec` — stateless packet codec. ``decode`` turns one device
  Opus packet into PCM; ``encode`` is NOT used for the reply path (use
  :class:`Opus16kEncoder`, which owns frame alignment).
* :class:`Opus16kEncoder` — streaming encoder for reply audio: feed arbitrary
  16 kHz mono PCM, get back exact 20 ms (320-sample) Opus frames; ``flush``
  emits one final padded frame (smallest valid Opus frame size) so short
  utterances still produce a playable packet.
* :class:`PassthroughOpus` — PCM in, PCM out. Lets the whole pipeline run
  without a native codec (tests / synthetic clients).

``RealOpus`` is constructed only when selected (kind ``"real"``); the native
``opus`` library is located via ``OPUS_DLL_DIR`` (env), then a per-venv
``.venv/Scripts/opus`` directory (this repo keeps ``opus.dll`` there).
"""
from __future__ import annotations

import os

from abc import ABC, abstractmethod


def _ensure_native() -> None:
    """Make the native ``opus`` library discoverable by opuslib.

    opuslib loads it through ``ctypes.util.find_library('opus')`` (PATH on
    Windows). We do NOT assume any particular install layout: first an
    explicit ``OPUS_DLL_DIR`` env override, then the venv-local dir next to
    this repo (``.venv/Scripts/opus``), which this project ships the DLL in.
    """
    import ctypes.util

    if ctypes.util.find_library("opus"):
        return
    repo_root = os.path.dirname(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    )
    candidates = [os.environ.get("OPUS_DLL_DIR", "")]
    candidates.append(os.path.join(os.getcwd(), ".venv", "Scripts", "opus"))
    candidates.append(os.path.join(repo_root, ".venv", "Scripts", "opus"))
    # Windows os.add_dll_directory requires an ABSOLUTE path (a relative one
    # raises WinError 87); normalize every candidate before use.
    for cand in (os.path.abspath(c) for c in candidates if c):
        if os.path.isfile(os.path.join(cand, "opus.dll")):
            os.environ["PATH"] = cand + os.pathsep + os.environ.get("PATH", "")
            if hasattr(os, "add_dll_directory"):
                os.add_dll_directory(cand)
            break


class OpusCodec(ABC):
    @abstractmethod
    def decode(self, data: bytes) -> bytes:
        """One Opus packet -> PCM (16-bit LE mono)."""


class PassthroughOpus(OpusCodec):
    """PCM in, PCM out. Lets the whole pipeline run without a native codec."""

    def decode(self, data: bytes) -> bytes:
        return data


class RealOpus(OpusCodec):
    """Real Opus codec via ``opuslib``.

    Frame math (verified): 20 ms @ 16 kHz mono = 320 samples = 640 B PCM.
    ``Decoder.decode(packet, 320)``; an undecodable packet (loss) returns
    silence of the requested frame size — we raise so the caller can skip it
    rather than inject silent PCM into the VAD buffer.
    """

    FRAME_SAMPLES = 320  # 20 ms @ 16 kHz

    def __init__(self, sample_rate: int = 16000, channels: int = 1) -> None:
        _ensure_native()
        from opuslib.classes import Decoder  # lazy (CON-002)

        self._dec = Decoder(sample_rate, channels)

    def decode(self, data: bytes) -> bytes:
        pcm = self._dec.decode(data, self.FRAME_SAMPLES)
        if pcm is None:
            raise ValueError("undecodable opus packet")
        return pcm


class OpusEncoder:
    """Streaming mono PCM -> exact 20 ms Opus frames.

    Feed PCM in any size; the encoder buffers a partial frame and emits
    complete 20 ms frames as they become available. ``flush`` pads the
    remainder up to the smallest valid Opus frame (2.5 ms) so the device
    always gets a decodable final packet.

    Frame math is rate-generic: at 16 kHz a frame is 320 samples (640 B);
    at 24 kHz (the TTS output rate this bridge encodes replies at) it is
    480 samples (960 B). The valid Opus frame sizes (2.5/5/10/20/40/60 ms)
    scale with the rate.
    """

    def __init__(self, sample_rate: int = 16000, channels: int = 1) -> None:
        _ensure_native()
        from opuslib.classes import Encoder

        self._sr = sample_rate
        self._enc = Encoder(sample_rate, channels, "voip")
        self._carry = b""
        self.FRAME_SAMPLES = sample_rate // 50  # 20 ms
        self.FRAME_BYTES = self.FRAME_SAMPLES * 2
        # Valid Opus frame sizes (2.5/5/10/20/40/60 ms) at this rate.
        self._VALID = tuple(
            int(round(sample_rate * ms / 1000)) for ms in (2.5, 5, 10, 20, 40, 60)
        )

    def feed(self, pcm: bytes) -> list[bytes]:
        """Add PCM; return any complete 20 ms frames as Opus packets."""
        self._carry += pcm
        out: list[bytes] = []
        while len(self._carry) >= self.FRAME_BYTES:
            frame = self._carry[: self.FRAME_BYTES]
            self._carry = self._carry[self.FRAME_BYTES :]
            out.append(self._enc.encode(frame, self.FRAME_SAMPLES))
        return out

    def flush(self) -> list[bytes]:
        """Emit the trailing partial frame (padded to a valid frame size)."""
        n = (len(self._carry) + 1) // 2  # keep an odd trailing sample
        if n == 0:
            return []
        for valid in self._VALID:
            if n <= valid:
                target = valid
                break
        else:  # pragma: no cover — carry is always < one 20 ms frame
            target = self._VALID[-1]
        pad = (target - n) * 2
        pcm = self._carry + b"\x00" * pad
        self._carry = b""
        return [self._enc.encode(pcm, target)]


# Backward-compat alias (this class was 16 kHz-specific).
Opus16kEncoder = OpusEncoder


class _PassthroughEncoder:
    """Encoder shim with the same interface as :class:`Opus16kEncoder`."""

    def feed(self, pcm: bytes) -> list[bytes]:
        return [pcm] if pcm else []

    def flush(self) -> list[bytes]:
        return []


def make_opus(kind: str = "passthrough", **kw) -> OpusCodec:
    if kind == "passthrough":
        return PassthroughOpus()
    if kind == "real":
        return RealOpus(**kw)
    raise ValueError(f"unknown opus kind: {kind}")


def make_encoder(kind: str = "passthrough", **kw):
    if kind == "passthrough":
        return _PassthroughEncoder()
    if kind == "real":
        return Opus16kEncoder(**kw)
    raise ValueError(f"unknown opus kind: {kind}")