openmusic_analysis/audio/decoder.py

from __future__ import annotations

import json
import subprocess
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from openmusic_analysis.errors import AudioDecodeError


@dataclass(frozen=True)
class AudioMetadata:
    duration_ms: int
    source_sample_rate: int | None
    source_channels: int | None
    source_format: str | None


@dataclass(frozen=True)
class DecodedAudio:
    waveform: np.ndarray
    sample_rate: int
    metadata: AudioMetadata


class AudioDecoder:
    """The only source-file decoder; analyzers receive canonical in-memory PCM."""

    def __init__(
        self,
        *,
        ffmpeg_binary: str = "ffmpeg",
        ffprobe_binary: str = "ffprobe",
        canonical_sample_rate: int = 48_000,
        max_audio_seconds: float = 1800.0,
        timeout_seconds: float = 120.0,
    ) -> None:
        self.ffmpeg_binary = ffmpeg_binary
        self.ffprobe_binary = ffprobe_binary
        self.canonical_sample_rate = canonical_sample_rate
        self.max_audio_seconds = max_audio_seconds
        self.timeout_seconds = timeout_seconds

    def decode(self, source_path: str | Path) -> DecodedAudio:
        path = str(source_path)
        probe = self._probe(path)
        if probe["duration"] is not None and probe["duration"] > self.max_audio_seconds:
            raise AudioDecodeError(
                f"Audio duration exceeds the {self.max_audio_seconds:g} second limit."
            )

        command = [
            self.ffmpeg_binary,
            "-v",
            "error",
            "-nostdin",
            "-threads",
            "1",
            "-i",
            path,
            "-map",
            "0:a:0",
            "-vn",
            "-ac",
            "1",
            "-ar",
            str(self.canonical_sample_rate),
            "-acodec",
            "pcm_f32le",
            "-f",
            "f32le",
            "pipe:1",
        ]
        try:
            completed = subprocess.run(
                command,
                check=False,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=self.timeout_seconds,
            )
        except (OSError, subprocess.TimeoutExpired) as exc:
            raise AudioDecodeError() from exc
        if completed.returncode != 0 or not completed.stdout:
            raise AudioDecodeError()

        waveform = np.frombuffer(completed.stdout, dtype="<f4").copy()
        if waveform.size == 0 or not np.isfinite(waveform).all():
            raise AudioDecodeError()
        actual_duration = waveform.size / self.canonical_sample_rate
        if actual_duration > self.max_audio_seconds:
            raise AudioDecodeError(
                f"Audio duration exceeds the {self.max_audio_seconds:g} second limit."
            )
        metadata = AudioMetadata(
            duration_ms=int(round(actual_duration * 1000)),
            source_sample_rate=probe["sample_rate"],
            source_channels=probe["channels"],
            source_format=probe["format"],
        )
        return DecodedAudio(
            waveform=np.ascontiguousarray(waveform, dtype=np.float32),
            sample_rate=self.canonical_sample_rate,
            metadata=metadata,
        )

    def _probe(self, path: str) -> dict[str, int | float | str | None]:
        command = [
            self.ffprobe_binary,
            "-v",
            "error",
            "-select_streams",
            "a:0",
            "-show_entries",
            "stream=sample_rate,channels,duration:format=format_name,duration",
            "-of",
            "json",
            path,
        ]
        try:
            completed = subprocess.run(
                command,
                check=False,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=min(30.0, self.timeout_seconds),
                text=True,
            )
            payload = json.loads(completed.stdout) if completed.returncode == 0 else {}
        except (OSError, subprocess.TimeoutExpired, json.JSONDecodeError) as exc:
            raise AudioDecodeError() from exc
        streams = payload.get("streams") or []
        if not streams:
            raise AudioDecodeError()
        stream = streams[0]
        container = payload.get("format") or {}
        duration_value = stream.get("duration") or container.get("duration")
        try:
            duration = float(duration_value) if duration_value is not None else None
        except (TypeError, ValueError):
            duration = None
        return {
            "duration": duration,
            "sample_rate": _optional_int(stream.get("sample_rate")),
            "channels": _optional_int(stream.get("channels")),
            "format": container.get("format_name"),
        }


def _optional_int(value: object) -> int | None:
    try:
        return int(value) if value is not None else None
    except (TypeError, ValueError):
        return None