from pydub import AudioSegment
import io
import numpy as np

def convert_audio(audio_data, fmt: str, sample_rate: int = 24000):
    """
    Converts raw audio data (numpy array or list of floats) to the target format.
    Assumes mono audio.
    """
    # Ensure numpy array
    if not isinstance(audio_data, np.ndarray):
        audio_data = np.array(audio_data)

    # Check if float and normalize/convert to int16
    if audio_data.dtype.kind == 'f':
        # Clip to Avoid wrap-around
        audio_data = np.clip(audio_data, -1.0, 1.0)
        # Convert to 16-bit PCM
        audio_data = (audio_data * 32767).astype(np.int16)
    
    seg = AudioSegment(
        audio_data.tobytes(),
        frame_rate=sample_rate,
        sample_width=2, # 16-bit
        channels=1
    )
    
    buf = io.BytesIO()
    seg.export(buf, format=fmt)
    
    mime = {
        "wav": "audio/wav",
        "mp3": "audio/mpeg",
        "ogg": "audio/ogg",
        "flac": "audio/flac",
        "aac": "audio/aac"
    }.get(fmt, "audio/wav")
    
    return buf.getvalue(), mime