banban/talkingq-url/utils/audio_duration.py

import asyncio
import io
import math
import os
import shutil
import subprocess
import tempfile
import wave

try:
    import soundfile as sf
except Exception:  # pragma: no cover - depends on runtime native libs
    sf = None

from utils.audio_format import (
    DEFAULT_CHANNELS,
    DEFAULT_SAMPLE_RATE,
    DEFAULT_SAMPLE_WIDTH,
    detect_audio_format,
)
from utils.logger import session_logger


def _duration_from_wave(audio_data: bytes) -> int | None:
    try:
        with wave.open(io.BytesIO(audio_data), "rb") as wav_file:
            frame_rate = wav_file.getframerate()
            frame_count = wav_file.getnframes()
            if frame_rate <= 0 or frame_count <= 0:
                return None
            return max(1, int(round(frame_count * 1000 / frame_rate)))
    except Exception:
        return None


def _duration_from_soundfile(audio_data: bytes) -> int | None:
    if sf is None:
        return None
    try:
        with sf.SoundFile(io.BytesIO(audio_data)) as audio_file:
            sample_rate = int(audio_file.samplerate or 0)
            frames = int(audio_file.frames or 0)
            if sample_rate <= 0 or frames <= 0:
                return None
            return max(1, int(round(frames * 1000 / sample_rate)))
    except Exception:
        return None


def _duration_from_pcm_16k_mono(audio_data: bytes) -> int | None:
    frame_size = DEFAULT_CHANNELS * DEFAULT_SAMPLE_WIDTH
    if DEFAULT_SAMPLE_RATE <= 0 or frame_size <= 0 or not audio_data:
        return None
    frames = len(audio_data) / frame_size
    if frames <= 0:
        return None
    return max(1, int(round(frames * 1000 / DEFAULT_SAMPLE_RATE)))


def _duration_from_ffprobe(audio_data: bytes, suffix: str) -> int | None:
    ffprobe_path = shutil.which("ffprobe")
    if not ffprobe_path:
        return None

    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
            tmp_file.write(audio_data)
            tmp_path = tmp_file.name

        process = subprocess.run(
            [
                ffprobe_path,
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                tmp_path,
            ],
            check=False,
            capture_output=True,
            text=True,
            timeout=5,
        )
        if process.returncode != 0:
            return None
        duration_seconds = float((process.stdout or "").strip())
        if not math.isfinite(duration_seconds) or duration_seconds <= 0:
            return None
        return max(1, int(round(duration_seconds * 1000)))
    except Exception:
        return None
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
            except OSError:
                pass


def _suffix_for_audio(*, audio_format: str, mime_type: str | None) -> str:
    normalized_mime = (mime_type or "").strip().lower()
    if audio_format == "mp3" or normalized_mime in {"audio/mpeg", "audio/mp3"}:
        return ".mp3"
    if audio_format == "wav" or normalized_mime in {"audio/wav", "audio/x-wav"}:
        return ".wav"
    if normalized_mime in {"audio/aac", "audio/x-aac"}:
        return ".aac"
    if normalized_mime in {"audio/mp4", "audio/x-m4a", "audio/m4a"}:
        return ".m4a"
    if normalized_mime == "audio/webm":
        return ".webm"
    return ".audio"


def _can_treat_as_pcm(*, audio_format: str, mime_type: str | None) -> bool:
    normalized_mime = (mime_type or "").strip().lower()
    return audio_format == "pcm_s16le_16k_mono" and normalized_mime in {
        "",
        "application/octet-stream",
        "audio/pcm",
        "audio/raw",
        "audio/s16le",
    }


def parse_audio_duration_ms_sync(
    audio_data: bytes,
    *,
    mime_type: str | None = None,
    source: str = "audio",
) -> int | None:
    if not audio_data:
        return None

    audio_format = detect_audio_format(audio_data)

    duration_ms = _duration_from_wave(audio_data)
    if duration_ms is not None:
        return duration_ms

    duration_ms = _duration_from_soundfile(audio_data)
    if duration_ms is not None:
        return duration_ms

    duration_ms = _duration_from_ffprobe(
        audio_data,
        suffix=_suffix_for_audio(audio_format=audio_format, mime_type=mime_type),
    )
    if duration_ms is not None:
        return duration_ms

    if _can_treat_as_pcm(audio_format=audio_format, mime_type=mime_type):
        duration_ms = _duration_from_pcm_16k_mono(audio_data)
        if duration_ms is not None:
            return duration_ms

    session_logger.warning(
        "system",
        "audio_duration",
        (
            "failed to parse audio duration: "
            f"source={source}, detected_format={audio_format}, "
            f"mime_type={mime_type or 'unknown'}, size_bytes={len(audio_data)}"
        ),
    )
    return None


async def parse_audio_duration_ms(
    audio_data: bytes,
    *,
    mime_type: str | None = None,
    source: str = "audio",
) -> int | None:
    return await asyncio.to_thread(
        parse_audio_duration_ms_sync,
        audio_data,
        mime_type=mime_type,
        source=source,
    )