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 Opus16kEncoder:
    """Streaming 16 kHz mono PCM -> exact 20 ms Opus frames.

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

    FRAME_SAMPLES = 320  # 20 ms @ 16 kHz
    FRAME_BYTES = 640
    # Valid Opus frame sizes (2.5/5/10/20/40/60 ms) in samples @ 16 kHz.
    _VALID = (40, 80, 160, 320, 640, 960)

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

        self._enc = Encoder(sample_rate, channels, "voip")
        self._carry = b""

    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) // 2
        if n == 0:
            return []
        for valid in self._VALID:
            if n <= valid:
                target = valid
                break
        else:  # pragma: no cover — carry is always < 320 samples
            target = 960
        pad = (target - n) * 2
        pcm = self._carry + b"\x00" * pad
        self._carry = b""
        return [self._enc.encode(pcm, target)]


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}")