src/turn_detection/data/audio.py

"""Audio payload inspection and exact hashing without mandatory DSP packages."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import io
import json
from pathlib import Path
import struct
from typing import Any, Mapping, Sequence
import wave


class AudioPayloadError(ValueError):
    """Raised when an audio value cannot be resolved to bytes or samples."""


@dataclass(frozen=True)
class AudioInfo:
    """Content fingerprint and inexpensive container metadata."""

    sha256: str
    num_bytes: int
    format: str
    sample_rate: int | None = None
    num_channels: int | None = None
    num_frames: int | None = None
    bits_per_sample: int | None = None
    duration_seconds: float | None = None
    path: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "audio_sha256": self.sha256,
            "audio_num_bytes": self.num_bytes,
            "audio_format": self.format,
            "sample_rate": self.sample_rate,
            "num_channels": self.num_channels,
            "num_frames": self.num_frames,
            "bits_per_sample": self.bits_per_sample,
            "duration_seconds": self.duration_seconds,
            "audio_path": self.path,
        }


def _resolve_path(path_value: str | Path, base_dir: str | Path | None) -> Path:
    path = Path(path_value).expanduser()
    if not path.is_absolute() and base_dir is not None:
        path = Path(base_dir) / path
    return path


def _read_path(path: Path) -> bytes:
    try:
        return path.read_bytes()
    except OSError as exc:
        raise AudioPayloadError(f"cannot read audio path {path}: {exc}") from exc


def extract_audio_bytes(
    audio: Any,
    *,
    base_dir: str | Path | None = None,
) -> tuple[bytes, str | None]:
    """Resolve common Hugging Face/local audio representations to exact bytes.

    Hugging Face Parquet rows normally contain ``{"bytes": ..., "path": ...}``.
    Byte content is preferred over the often-virtual path.  Decoded arrays are
    represented canonically when no encoded payload is available.
    """

    if isinstance(audio, bytes):
        return audio, None
    if isinstance(audio, (bytearray, memoryview)):
        return bytes(audio), None
    if isinstance(audio, (str, Path)):
        path = _resolve_path(audio, base_dir)
        return _read_path(path), str(audio)
    if isinstance(audio, Mapping):
        path_value = audio.get("path")
        encoded = audio.get("bytes")
        if encoded is not None:
            if not isinstance(encoded, (bytes, bytearray, memoryview)):
                raise AudioPayloadError("audio['bytes'] is not bytes-like")
            return bytes(encoded), str(path_value) if path_value is not None else None
        if path_value:
            path = _resolve_path(str(path_value), base_dir)
            return _read_path(path), str(path_value)
        if audio.get("array") is not None:
            return _canonical_array_bytes(audio["array"], audio.get("sampling_rate")), None
    if hasattr(audio, "read"):
        stream = audio
        try:
            position = stream.tell()
        except (AttributeError, OSError):
            position = None
        payload = stream.read()
        if position is not None:
            try:
                stream.seek(position)
            except (AttributeError, OSError):
                pass
        if isinstance(payload, str):
            payload = payload.encode("utf-8")
        if isinstance(payload, (bytes, bytearray, memoryview)):
            return bytes(payload), None
    raise AudioPayloadError(
        "unsupported audio value; expected bytes, a path, an Audio mapping, or a decoded array"
    )


def _canonical_array_bytes(array: Any, sampling_rate: Any) -> bytes:
    header: dict[str, Any] = {"sampling_rate": sampling_rate}
    if hasattr(array, "dtype"):
        header["dtype"] = str(array.dtype)
    if hasattr(array, "shape"):
        header["shape"] = tuple(int(value) for value in array.shape)
    if hasattr(array, "tobytes"):
        body = array.tobytes(order="C")
    elif isinstance(array, Sequence) and not isinstance(array, (str, bytes, bytearray)):
        # JSON is slower than ndarray.tobytes but deterministic and dependency-free.
        try:
            body = json.dumps(array, separators=(",", ":"), allow_nan=False).encode("utf-8")
        except (TypeError, ValueError) as exc:
            raise AudioPayloadError(f"cannot serialize decoded audio array: {exc}") from exc
    else:
        raise AudioPayloadError("audio['array'] cannot be serialized")
    return b"TURN_ARRAY\0" + json.dumps(header, sort_keys=True).encode("utf-8") + b"\0" + body


def _infer_format(payload: bytes, path: str | None) -> str:
    if payload.startswith(b"RIFF") and payload[8:12] == b"WAVE":
        return "wav"
    if payload.startswith(b"fLaC"):
        return "flac"
    if payload.startswith(b"OggS"):
        return "ogg"
    if payload.startswith(b"ID3") or payload[:2] in (b"\xff\xfb", b"\xff\xf3", b"\xff\xf2"):
        return "mp3"
    if payload.startswith(b"TURN_ARRAY\0"):
        return "decoded_array"
    if path:
        suffix = Path(path).suffix.lower().lstrip(".")
        if suffix:
            return suffix
    return "unknown"


def _wav_metadata(payload: bytes) -> tuple[int, int, int, int, float] | None:
    try:
        with wave.open(io.BytesIO(payload), "rb") as wav_file:
            sample_rate = wav_file.getframerate()
            channels = wav_file.getnchannels()
            frames = wav_file.getnframes()
            bits = wav_file.getsampwidth() * 8
    except (wave.Error, EOFError, OSError, struct.error):
        return None
    duration = frames / sample_rate if sample_rate else None
    return sample_rate, channels, frames, bits, duration


def _flac_metadata(payload: bytes) -> tuple[int, int, int, int, float] | None:
    # FLAC STREAMINFO is always the first metadata block and has a 34-byte body.
    if len(payload) < 42 or payload[:4] != b"fLaC":
        return None
    block_type = payload[4] & 0x7F
    block_length = int.from_bytes(payload[5:8], "big")
    if block_type != 0 or block_length < 34 or len(payload) < 8 + block_length:
        return None
    streaminfo = payload[8 : 8 + block_length]
    packed = int.from_bytes(streaminfo[10:18], "big")
    sample_rate = packed >> 44
    channels = ((packed >> 41) & 0x7) + 1
    bits = ((packed >> 36) & 0x1F) + 1
    frames = packed & ((1 << 36) - 1)
    duration = frames / sample_rate if sample_rate else None
    return sample_rate, channels, frames, bits, duration


def inspect_audio(audio: Any, *, base_dir: str | Path | None = None) -> AudioInfo:
    """Hash exact audio bytes and extract WAV/FLAC STREAMINFO metadata."""

    payload, path = extract_audio_bytes(audio, base_dir=base_dir)
    if not payload:
        raise AudioPayloadError("audio payload is empty")
    digest = hashlib.sha256(payload).hexdigest()
    audio_format = _infer_format(payload, path)
    metadata = _wav_metadata(payload) if audio_format == "wav" else None
    if audio_format == "flac":
        metadata = _flac_metadata(payload)
    if metadata is None:
        return AudioInfo(digest, len(payload), audio_format, path=path)
    sample_rate, channels, frames, bits, duration = metadata
    return AudioInfo(
        sha256=digest,
        num_bytes=len(payload),
        format=audio_format,
        sample_rate=sample_rate,
        num_channels=channels,
        num_frames=frames,
        bits_per_sample=bits,
        duration_seconds=duration,
        path=path,
    )


def exact_audio_sha256(audio: Any, *, base_dir: str | Path | None = None) -> str:
    """Return the SHA-256 digest of the exact encoded payload/canonical array."""

    return inspect_audio(audio, base_dir=base_dir).sha256