"""Audio I/O utilities for mlx-music.

Handles reading/writing WAV files, resampling, and stereo normalization.
All audio in mlx-music is 48kHz stereo float32 by convention.
"""

import numpy as np
from pathlib import Path
from typing import Optional, Tuple

TARGET_SR = 48000
TARGET_CHANNELS = 2  # stereo


def load_wav(path: str, target_sr: int = TARGET_SR) -> Tuple[np.ndarray, int]:
    """Load audio file and normalize to stereo at target sample rate.

    Args:
        path: Path to audio file (WAV, FLAC, OGG, etc.)
        target_sr: Target sample rate (default 48000)

    Returns:
        Tuple of (audio, sample_rate) where audio is shape [samples, 2] float32.
    """
    import soundfile as sf

    audio, sr = sf.read(str(path), dtype="float32")

    # Ensure 2D: [samples, channels]
    if audio.ndim == 1:
        audio = np.stack([audio, audio], axis=-1)  # mono -> stereo
    elif audio.shape[1] == 1:
        audio = np.concatenate([audio, audio], axis=-1)
    elif audio.shape[1] > 2:
        audio = audio[:, :2]  # take first 2 channels

    # Resample if needed
    if sr != target_sr:
        audio = _resample(audio, sr, target_sr)
        sr = target_sr

    return audio, sr


def save_wav(
    audio: np.ndarray,
    path: str,
    sample_rate: int = TARGET_SR,
    normalize: bool = True,
) -> None:
    """Save audio array as WAV file.

    Args:
        audio: Audio array, shape [samples] or [samples, channels], float32.
        path: Output path (.wav)
        sample_rate: Sample rate (default 48000)
        normalize: If True, normalize peak to 0.95 to avoid clipping.
    """
    import soundfile as sf

    if normalize:
        peak = np.abs(audio).max()
        if peak > 0:
            audio = audio * (0.95 / peak)

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    sf.write(str(path), audio, sample_rate, subtype="FLOAT")


def normalize_stereo_48k(audio: np.ndarray, sr: int) -> Tuple[np.ndarray, int]:
    """Normalize audio to stereo 48kHz float32.

    Args:
        audio: Input audio array
        sr: Input sample rate

    Returns:
        (audio, 48000) normalized tuple
    """
    # Ensure 2D
    if audio.ndim == 1:
        audio = np.stack([audio, audio], axis=-1)
    elif audio.shape[1] == 1:
        audio = np.concatenate([audio, audio], axis=-1)
    elif audio.shape[1] > 2:
        audio = audio[:, :2]

    # Resample
    if sr != TARGET_SR:
        audio = _resample(audio, sr, TARGET_SR)

    return audio.astype(np.float32), TARGET_SR


def duration_to_samples(duration_seconds: float, sample_rate: int = TARGET_SR) -> int:
    """Convert duration in seconds to sample count."""
    return int(duration_seconds * sample_rate)


def samples_to_duration(num_samples: int, sample_rate: int = TARGET_SR) -> float:
    """Convert sample count to duration in seconds."""
    return num_samples / sample_rate


def _resample(audio: np.ndarray, src_sr: int, target_sr: int) -> np.ndarray:
    """Simple linear interpolation resampling.

    For production quality, consider using librosa or scipy.signal.resample.
    This is sufficient for basic usage.
    """
    if src_sr == target_sr:
        return audio

    ratio = target_sr / src_sr
    target_len = int(len(audio) * ratio)

    # Use numpy interp for each channel
    x_old = np.linspace(0, 1, len(audio))
    x_new = np.linspace(0, 1, target_len)

    if audio.ndim == 1:
        return np.interp(x_new, x_old, audio).astype(np.float32)

    channels = []
    for ch in range(audio.shape[1]):
        channels.append(np.interp(x_new, x_old, audio[:, ch]))
    return np.stack(channels, axis=-1).astype(np.float32)
