app/services/audio_utils.py

from __future__ import annotations

import contextlib
import math
import shutil
import struct
import subprocess
import tempfile
import wave
from pathlib import Path
from typing import Any


DEFAULT_SAMPLE_RATE = 24_000
DEFAULT_SAMPLE_WIDTH = 2
DEFAULT_CHANNELS = 1


class AudioUtilityError(RuntimeError):
    """Raised when a local audio utility cannot complete safely."""


def get_audio_duration_seconds(file_path: str | Path) -> float:
    path = Path(file_path)
    if not path.exists():
        raise AudioUtilityError(f"Audio file does not exist: {path}")

    ffprobe_duration = _ffprobe_duration_seconds(path)
    if ffprobe_duration is not None:
        return ffprobe_duration

    if path.suffix.lower() == ".wav":
        with contextlib.closing(wave.open(str(path), "rb")) as wav_file:
            frame_count = wav_file.getnframes()
            frame_rate = wav_file.getframerate()
            if frame_rate <= 0:
                raise AudioUtilityError(f"Invalid WAV frame rate for {path}")
            return round(frame_count / float(frame_rate), 3)

    try:
        from mutagen import File as MutagenFile
    except ImportError as exc:
        raise AudioUtilityError("Install mutagen to detect non-WAV audio durations.") from exc

    audio = MutagenFile(path)
    if audio is None or audio.info is None:
        raise AudioUtilityError(f"Could not read audio metadata for {path}")
    length = getattr(audio.info, "length", None)
    if length is None:
        raise AudioUtilityError(f"Audio duration is missing for {path}")
    return round(float(length), 3)


def media_has_audio_stream(file_path: str | Path) -> bool:
    return _media_has_stream(file_path, "a:0", "audio")


def media_has_video_stream(file_path: str | Path) -> bool:
    return _media_has_stream(file_path, "v:0", "video")


def _media_has_stream(file_path: str | Path, stream_selector: str, expected: str) -> bool:
    path = Path(file_path)
    if not path.exists() or not _ffprobe_available():
        return False
    command = [
        "ffprobe",
        "-v",
        "error",
        "-select_streams",
        stream_selector,
        "-show_entries",
        "stream=codec_type",
        "-of",
        "csv=p=0",
        str(path),
    ]
    result = subprocess.run(command, capture_output=True, text=True, timeout=30, check=False)
    return result.returncode == 0 and expected in result.stdout.lower()


def normalize_audio_volume(file_path: str | Path, *, enabled: bool = True) -> Path:
    path = Path(file_path)
    if not enabled or not _ffmpeg_available() or not path.exists():
        return path

    with tempfile.NamedTemporaryFile(
        suffix=path.suffix,
        delete=False,
        dir=str(path.parent),
    ) as temp_file:
        temp_path = Path(temp_file.name)

    command = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(path),
        "-af",
        "loudnorm=I=-18:TP=-2:LRA=11",
        str(temp_path),
    ]
    result = subprocess.run(command, capture_output=True, text=True, timeout=60, check=False)
    if result.returncode == 0 and temp_path.exists() and temp_path.stat().st_size > 0:
        temp_path.replace(path)
    else:
        temp_path.unlink(missing_ok=True)

    return path


def trim_silence_and_pad_audio(
    file_path: str | Path,
    *,
    padding_seconds: float = 0.4,
    enabled: bool = True,
) -> Path:
    path = Path(file_path)
    if not enabled or not _ffmpeg_available() or not path.exists():
        return path

    with tempfile.NamedTemporaryFile(
        suffix=path.suffix,
        delete=False,
        dir=str(path.parent),
    ) as temp_file:
        temp_path = Path(temp_file.name)

    command = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        str(path),
        "-af",
        (
            "silenceremove=start_periods=1:start_duration=0.15:start_threshold=-45dB:"
            "stop_periods=1:stop_duration=0.3:stop_threshold=-45dB,"
            f"apad=pad_dur={max(padding_seconds, 0):.2f}"
        ),
        str(temp_path),
    ]
    result = subprocess.run(command, capture_output=True, text=True, timeout=60, check=False)
    if result.returncode == 0 and temp_path.exists() and temp_path.stat().st_size > 0:
        temp_path.replace(path)
    else:
        temp_path.unlink(missing_ok=True)

    return path


def merge_wav_files(input_files: list[str | Path], output_file: str | Path) -> Path:
    paths = [Path(file) for file in input_files]
    if not paths:
        raise AudioUtilityError("No WAV files were provided for merging.")

    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    params: Any | None = None
    expected_format: tuple[int, int, int, str, str] | None = None
    frames: list[bytes] = []
    for path in paths:
        with contextlib.closing(wave.open(str(path), "rb")) as wav_file:
            current_params = wav_file.getparams()
            comparable = (
                current_params.nchannels,
                current_params.sampwidth,
                current_params.framerate,
                current_params.comptype,
                current_params.compname,
            )
            if params is None:
                params = current_params
                expected_format = comparable
            elif comparable != expected_format:
                raise AudioUtilityError("WAV files use different audio formats.")
            frames.append(wav_file.readframes(current_params.nframes))

    if params is None:
        raise AudioUtilityError("Could not read WAV parameters.")

    with contextlib.closing(wave.open(str(output_path), "wb")) as output_wav:
        output_wav.setparams(params)
        for frame_bytes in frames:
            output_wav.writeframes(frame_bytes)

    return output_path


def _ffmpeg_available() -> bool:
    return shutil.which("ffmpeg") is not None


def _ffprobe_available() -> bool:
    return shutil.which("ffprobe") is not None


def _ffprobe_duration_seconds(path: Path) -> float | None:
    if not _ffprobe_available():
        return None
    command = [
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "format=duration",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        str(path),
    ]
    result = subprocess.run(command, capture_output=True, text=True, timeout=30, check=False)
    if result.returncode != 0:
        return None
    try:
        duration = float(result.stdout.strip())
    except ValueError:
        return None
    if duration <= 0:
        return None
    return round(duration, 3)


def create_mock_audio(
    output_file: str | Path,
    duration_seconds: float,
    *,
    tone: bool = False,
) -> Path:
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    safe_duration = max(float(duration_seconds), 0.5)
    frame_count = int(DEFAULT_SAMPLE_RATE * safe_duration)
    amplitude = 420 if tone else 0

    with contextlib.closing(wave.open(str(output_path), "wb")) as wav_file:
        wav_file.setnchannels(DEFAULT_CHANNELS)
        wav_file.setsampwidth(DEFAULT_SAMPLE_WIDTH)
        wav_file.setframerate(DEFAULT_SAMPLE_RATE)
        for index in range(frame_count):
            sample = 0
            if amplitude:
                sample = int(amplitude * math.sin(2 * math.pi * 440 * index / DEFAULT_SAMPLE_RATE))
            wav_file.writeframes(struct.pack("<h", sample))

    return output_path