# SPDX-License-Identifier: Apache-2.0
"""Shared utilities for audio engines (STT, TTS, STS)."""

import io
import struct
import wave

import numpy as np

# Default sample rate used when the model does not report one.
DEFAULT_SAMPLE_RATE = 24000
_MAX_WAV_CHUNK_SIZE = 0xFFFFFFFF


def wav_header(sample_rate: int, channels: int = 1, sample_width: int = 2) -> bytes:
    """Create a PCM WAV header suitable for streamed/unknown-length audio."""
    block_align = channels * sample_width
    byte_rate = sample_rate * block_align
    bits_per_sample = sample_width * 8
    return struct.pack(
        "<4sI4s4sIHHIIHH4sI",
        b"RIFF",
        _MAX_WAV_CHUNK_SIZE,
        b"WAVE",
        b"fmt ",
        16,
        1,  # PCM
        channels,
        sample_rate,
        byte_rate,
        block_align,
        bits_per_sample,
        b"data",
        _MAX_WAV_CHUNK_SIZE,
    )



def wav_bytes_to_pcm_frames(wav_bytes: bytes) -> tuple[int, int, int, bytes]:
    """Extract WAV metadata and raw PCM frames from WAV bytes."""
    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
        sample_rate = wf.getframerate()
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        pcm_bytes = wf.readframes(wf.getnframes())
    return sample_rate, channels, sample_width, pcm_bytes



def audio_to_wav_bytes(audio_array, sample_rate: int) -> bytes:
    """Convert a float32 audio array to 16-bit PCM WAV bytes.

    Args:
        audio_array: numpy or mlx array of float32 samples in [-1, 1]
        sample_rate: audio sample rate in Hz

    Returns:
        WAV-encoded bytes (RIFF header + PCM data)
    """
    # Ensure we have a numpy array for the wave module
    if not isinstance(audio_array, np.ndarray):
        # NumPy doesn't support bfloat16 — cast to float32 first
        if hasattr(audio_array, "dtype"):
            import mlx.core as mx

            if audio_array.dtype == mx.bfloat16:
                audio_array = audio_array.astype(mx.float32)
        audio_array = np.array(audio_array)

    # Flatten to 1-D (mono)
    audio_array = audio_array.flatten()

    # Clip to [-1, 1] then convert to int16
    audio_array = np.clip(audio_array, -1.0, 1.0)
    audio_int16 = (audio_array * 32767).astype(np.int16)

    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)  # 16-bit
        wf.setframerate(sample_rate)
        wf.writeframes(audio_int16.tobytes())
    return buf.getvalue()