app/services/audio.py

import io
import struct

import numpy as np
import soundfile as sf

from app.logging_config import get_logger

logger = get_logger('audio')

VALID_FORMATS = {'mp3', 'wav', 'opus', 'aac', 'flac', 'pcm'}


def validate_format(fmt: str) -> str:
    fmt = fmt.lower()
    if fmt == 'mpeg':
        return 'mp3'
    if fmt not in VALID_FORMATS:
        logger.warning(f"Unknown format '{fmt}', falling back to wav")
        return 'wav'
    return fmt


def numpy_to_pcm_bytes(audio: np.ndarray) -> bytes:
    if audio.ndim > 1:
        audio = audio[0]
    pcm = (audio * 32767).clip(-32768, 32767).astype(np.int16)
    return pcm.tobytes()


def convert_audio(
    audio: np.ndarray, sample_rate: int, target_format: str = 'wav'
) -> io.BytesIO:
    buffer = io.BytesIO()

    if audio.ndim == 1:
        audio = audio[np.newaxis, :]

    if target_format == 'pcm':
        pcm_bytes = numpy_to_pcm_bytes(audio)
        buffer.write(pcm_bytes)
        buffer.seek(0)
        return buffer

    if target_format in ('opus', 'aac', 'mp3', 'flac'):
        sf.write(buffer, audio.T, sample_rate, format=target_format.upper())
        buffer.seek(0)
        return buffer

    sf.write(buffer, audio.T, sample_rate, format='WAV')
    buffer.seek(0)
    return buffer


def write_wav_header(
    sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16, num_frames: int = 0
) -> bytes:
    byte_rate = sample_rate * num_channels * bits_per_sample // 8
    block_align = num_channels * bits_per_sample // 8

    data_size = num_frames * block_align
    if num_frames == 0:
        data_size = 0xFFFFFFFF - 36

    chunk_size = 36 + data_size

    header = io.BytesIO()
    header.write(b'RIFF')
    header.write(struct.pack('<I', chunk_size))
    header.write(b'WAVE')
    header.write(b'fmt ')
    header.write(struct.pack('<I', 16))
    header.write(struct.pack('<H', 1))
    header.write(struct.pack('<H', num_channels))
    header.write(struct.pack('<I', sample_rate))
    header.write(struct.pack('<I', byte_rate))
    header.write(struct.pack('<H', block_align))
    header.write(struct.pack('<H', bits_per_sample))
    header.write(b'data')
    header.write(struct.pack('<I', data_size))

    return header.getvalue()


def get_mime_type(fmt: str) -> str:
    mime_types = {
        'wav': 'audio/wav',
        'mp3': 'audio/mpeg',
        'pcm': 'audio/L16',
        'opus': 'audio/opus',
        'aac': 'audio/aac',
        'flac': 'audio/flac',
    }
    return mime_types.get(fmt, f'audio/{fmt}')