"""Audio I/O utilities using miniaudio.

This module provides functions for reading and writing audio files.
- Reading: Uses miniaudio to support WAV, MP3, FLAC, and Vorbis formats.
           Uses ffmpeg for M4A/AAC, OGG, Opus, and WebM format support.
- Writing: Uses miniaudio for WAV and ffmpeg for MP3, FLAC, OGG, Opus, Vorbis, and WebM encoding.
"""

import io
import shutil
import subprocess
from pathlib import Path
from typing import Optional, Tuple, Union

import numpy as np

# Format mapping for miniaudio
_FORMAT_MAP = {
    "wav": "wav",
    "mp3": "mp3",
    "flac": "flac",
    "ogg": "vorbis",
    "vorbis": "vorbis",
    "m4a": "m4a",
    "aac": "m4a",
    "webm": "webm",
}

# Sample format mapping
_SAMPLE_FORMAT_MAP = {
    "int16": "SIGNED16",
    "int32": "SIGNED32",
    "float32": "FLOAT32",
}


def _detect_format_from_bytes(data: bytes) -> str:
    """Detect audio format from bytes data using magic bytes."""
    if data[:4] == b"RIFF" and data[8:12] == b"WAVE":
        return "wav"
    elif data[:3] == b"ID3" or (data[0:2] == b"\xff\xfb" or data[0:2] == b"\xff\xfa"):
        return "mp3"
    elif data[:4] == b"fLaC":
        return "flac"
    elif data[:4] == b"OggS":
        return "vorbis"
    elif data[4:8] == b"ftyp":
        # M4A/MP4/AAC container format
        return "m4a"
    elif data[:4] == b"\x1a\x45\xdf\xa3":
        # WebM/Matroska container (EBML header)
        return "webm"
    else:
        raise ValueError("Unable to detect audio format from bytes")


def _decode_ffmpeg(
    input_data: Union[str, Path, bytes],
) -> Tuple[np.ndarray, int, int]:
    """Decode audio using ffmpeg (for formats not supported by miniaudio like M4A).

    Args:
        input_data: Path to the audio file or raw bytes data.

    Returns:
        Tuple of (samples as int16 numpy array, sample_rate, nchannels).

    Raises:
        RuntimeError: If ffmpeg is not found or decoding fails.
    """
    ffmpeg_path = shutil.which("ffmpeg")
    if ffmpeg_path is None:
        raise RuntimeError(
            "\n"
            "========================================\n"
            "  ffmpeg not found!\n"
            "========================================\n"
            "\n"
            "ffmpeg is required for M4A/AAC/WebM audio decoding.\n"
            "\n"
            "Install ffmpeg:\n"
            "  macOS:  brew install ffmpeg\n"
            "  Ubuntu: sudo apt install ffmpeg\n"
        )

    # First, get audio info using ffprobe
    ffprobe_path = shutil.which("ffprobe")
    if ffprobe_path is None:
        raise RuntimeError("ffprobe not found (usually installed with ffmpeg)")

    if isinstance(input_data, bytes):
        # Use stdin for bytes input
        probe_cmd = [
            ffprobe_path,
            "-v",
            "quiet",
            "-print_format",
            "json",
            "-show_streams",
            "-select_streams",
            "a:0",
            "-i",
            "pipe:0",
        ]
        probe_result = subprocess.run(
            probe_cmd,
            input=input_data,
            capture_output=True,
        )
    else:
        probe_cmd = [
            ffprobe_path,
            "-v",
            "quiet",
            "-print_format",
            "json",
            "-show_streams",
            "-select_streams",
            "a:0",
            str(input_data),
        ]
        probe_result = subprocess.run(probe_cmd, capture_output=True)

    if probe_result.returncode != 0:
        raise RuntimeError(f"ffprobe failed: {probe_result.stderr.decode()}")

    import json

    probe_info = json.loads(probe_result.stdout.decode())
    if not probe_info.get("streams"):
        raise RuntimeError("No audio streams found in file")

    stream = probe_info["streams"][0]
    sample_rate = int(stream.get("sample_rate", 44100))
    nchannels = int(stream.get("channels", 2))

    # Decode to raw PCM using ffmpeg
    if isinstance(input_data, bytes):
        decode_cmd = [
            ffmpeg_path,
            "-i",
            "pipe:0",
            "-f",
            "s16le",  # Output: signed 16-bit little-endian PCM
            "-acodec",
            "pcm_s16le",
            "-ar",
            str(sample_rate),
            "-ac",
            str(nchannels),
            "pipe:1",
        ]
        decode_result = subprocess.run(
            decode_cmd,
            input=input_data,
            capture_output=True,
        )
    else:
        decode_cmd = [
            ffmpeg_path,
            "-i",
            str(input_data),
            "-f",
            "s16le",
            "-acodec",
            "pcm_s16le",
            "-ar",
            str(sample_rate),
            "-ac",
            str(nchannels),
            "pipe:1",
        ]
        decode_result = subprocess.run(decode_cmd, capture_output=True)

    if decode_result.returncode != 0:
        raise RuntimeError(f"ffmpeg decoding failed: {decode_result.stderr.decode()}")

    # Convert raw PCM bytes to numpy array
    samples = np.frombuffer(decode_result.stdout, dtype=np.int16)

    return samples, sample_rate, nchannels


def read(
    file: Union[str, Path, io.BytesIO],
    always_2d: bool = False,
    dtype: str = "float64",
) -> Tuple[np.ndarray, int]:
    """Read an audio file using miniaudio (or ffmpeg for M4A/AAC/OGG/Opus/WebM).

    Args:
        file: Path to the audio file or a BytesIO object.
        always_2d: If True, always return a 2D array (samples, channels).
        dtype: Data type for the output array. Supports 'float32', 'float64', 'int16'.

    Returns:
        Tuple of (audio_data, sample_rate).
        audio_data is a numpy array with shape (samples,) for mono or (samples, channels) for multi-channel.
    """
    # Check if this is a file that needs ffmpeg
    use_ffmpeg = False
    if isinstance(file, (str, Path)):
        ext = Path(file).suffix.lstrip(".").lower()
        if ext in ("m4a", "aac", "ogg", "opus", "webm"):
            use_ffmpeg = True
    elif isinstance(file, io.BytesIO):
        file.seek(0)
        header = file.read(12)
        file.seek(0)
        if (
            header[4:8] == b"ftyp"
            or header[:4] == b"OggS"
            or header[:4] == b"\x1a\x45\xdf\xa3"
        ):
            use_ffmpeg = True

    if use_ffmpeg:
        # Use ffmpeg for M4A/AAC decoding
        if isinstance(file, io.BytesIO):
            file.seek(0)
            input_data = file.read()
        else:
            input_data = file
        samples, sample_rate, nchannels = _decode_ffmpeg(input_data)
    else:
        # Use miniaudio for other formats
        import miniaudio

        if isinstance(file, (str, Path)):
            # Get file info to preserve original sample rate and channels
            info = miniaudio.get_file_info(str(file))
            decoded = miniaudio.decode_file(
                str(file),
                nchannels=info.nchannels,
                sample_rate=info.sample_rate,
            )
        elif isinstance(file, io.BytesIO):
            file.seek(0)
            data = file.read()
            # Detect format and get info to preserve original sample rate and channels
            fmt = _detect_format_from_bytes(data)
            if fmt == "wav":
                info = miniaudio.wav_get_info(data)
            elif fmt == "mp3":
                info = miniaudio.mp3_get_info(data)
            elif fmt == "flac":
                info = miniaudio.flac_get_info(data)
            elif fmt == "vorbis":
                info = miniaudio.vorbis_get_info(data)
            else:
                raise ValueError(f"Unsupported format: {fmt}")
            decoded = miniaudio.decode(
                data,
                nchannels=info.nchannels,
                sample_rate=info.sample_rate,
            )
        else:
            raise TypeError(f"Unsupported file type: {type(file)}")

        sample_rate = decoded.sample_rate
        nchannels = decoded.nchannels

        # Convert to numpy array
        # miniaudio returns samples as array of signed 16-bit integers interleaved
        samples = np.array(decoded.samples, dtype=np.int16)

    # Reshape to (samples, channels) if multi-channel
    if nchannels > 1:
        samples = samples.reshape(-1, nchannels)

    # Convert to requested dtype
    if dtype in ("float32", "float64"):
        samples = samples.astype(dtype) / 32768.0
    elif dtype == "int16":
        pass  # Already int16
    else:
        samples = samples.astype(dtype)

    # Handle always_2d
    if always_2d and samples.ndim == 1:
        samples = samples[:, np.newaxis]

    return samples, sample_rate


def _check_ffmpeg_available() -> bool:
    """Check if ffmpeg is available on the system."""
    return shutil.which("ffmpeg") is not None


def _get_ffmpeg_path() -> str:
    """Get the path to ffmpeg executable.

    Returns:
        Path to ffmpeg executable.

    Raises:
        RuntimeError: If ffmpeg is not found.
    """
    ffmpeg_path = shutil.which("ffmpeg")
    if ffmpeg_path is None:
        raise RuntimeError(
            "\n"
            "========================================\n"
            "  ffmpeg not found!\n"
            "========================================\n"
            "\n"
            "ffmpeg is required for MP3/FLAC/WebM encoding and M4A/AAC/WebM decoding.\n"
            "\n"
            "Install ffmpeg:\n"
            "  macOS:  brew install ffmpeg\n"
            "  Ubuntu: sudo apt install ffmpeg\n"
            "\n"
            "Alternatively, use WAV format which doesn't require ffmpeg.\n"
        )
    return ffmpeg_path


def _encode_ffmpeg(
    data: np.ndarray,
    samplerate: int,
    nchannels: int,
    output: Union[str, Path, io.BytesIO],
    format: str = "mp3",
    bitrate: str = "128k",
) -> None:
    """Encode audio using ffmpeg.

    Args:
        data: Audio data as int16 numpy array (samples,) or (samples, channels)
        samplerate: Sample rate in Hz
        nchannels: Number of channels
        output: Output file path or BytesIO object
        format: Output format (mp3, flac, ogg, opus, vorbis, etc.)
        bitrate: Audio bitrate for lossy formats (default: 128k)
    """
    ffmpeg_path = _get_ffmpeg_path()

    # Prepare raw PCM bytes (interleaved for stereo)
    if data.ndim == 1:
        pcm_bytes = data.tobytes()
    else:
        pcm_bytes = data.flatten().tobytes()

    # Build ffmpeg command
    cmd = [
        ffmpeg_path,
        "-y",  # Overwrite output
        "-f",
        "s16le",  # Input: signed 16-bit little-endian PCM
        "-ar",
        str(samplerate),  # Sample rate
        "-ac",
        str(nchannels),  # Channels
        "-i",
        "pipe:0",  # Read from stdin
    ]

    # Add format-specific options
    if format == "mp3":
        cmd.extend(["-b:a", bitrate])
    elif format == "opus":
        cmd.extend(["-c:a", "libopus", "-b:a", bitrate])
    elif format == "webm":
        cmd.extend(["-c:a", "libopus", "-b:a", bitrate])
    elif format in ("ogg", "vorbis"):
        # Use FLAC codec in OGG container for maximum compatibility
        # Native vorbis encoder has limitations (experimental, stereo-only)
        # FLAC in OGG provides lossless compression and works with any channels
        cmd.extend(["-c:a", "flac"])

    # Map format to output container format
    if format == "vorbis":
        format = "ogg"  # Vorbis uses OGG container

    cmd.extend(["-f", format])

    if isinstance(output, io.BytesIO):
        cmd.append("pipe:1")
        result = subprocess.run(
            cmd,
            input=pcm_bytes,
            capture_output=True,
        )
        if result.returncode != 0:
            raise RuntimeError(f"ffmpeg failed: {result.stderr.decode()}")
        output.write(result.stdout)
        output.seek(0)
    else:
        cmd.append(str(output))
        result = subprocess.run(
            cmd,
            input=pcm_bytes,
            capture_output=True,
        )
        if result.returncode != 0:
            raise RuntimeError(f"ffmpeg failed: {result.stderr.decode()}")


def write(
    file: Union[str, Path, io.BytesIO],
    data: np.ndarray,
    samplerate: int,
    format: Optional[str] = None,
) -> None:
    """Write audio data to a file.

    Args:
        file: Path to the output file or a BytesIO object.
        data: Audio data as numpy array. Shape can be (samples,) for mono
              or (samples, channels) for multi-channel.
        samplerate: Sample rate in Hz.
        format: Output format. Supports 'wav', 'flac', 'mp3', 'ogg', 'opus', 'vorbis', 'webm'.
                If None, inferred from file extension.

    Note:
        WAV uses miniaudio for encoding.
        MP3, FLAC, OGG, Opus, Vorbis, and WebM use ffmpeg (must be installed: brew install ffmpeg).
    """
    import miniaudio

    # Determine format
    if format is None:
        if isinstance(file, (str, Path)):
            format = Path(file).suffix.lstrip(".").lower()
        else:
            format = "wav"  # Default to WAV for BytesIO

    format = format.lower()

    # Ensure data is numpy array (handle MLX arrays and other array-like types)
    if not isinstance(data, np.ndarray):
        # Check for MLX array (has tolist but not __array__)
        if hasattr(data, "tolist") and not hasattr(data, "__array__"):
            data = np.array(data.tolist())
        elif hasattr(data, "__array__"):
            data = np.asarray(data)
        else:
            data = np.array(data)

    # Convert to int16 for encoding
    if data.dtype in (np.float32, np.float64):
        # Clip to [-1, 1] range and convert to int16
        data = np.clip(data, -1.0, 1.0)
        data = (data * 32767).astype(np.int16)
    elif data.dtype != np.int16:
        data = data.astype(np.int16)

    # Get number of channels
    if data.ndim == 1:
        nchannels = 1
        samples_flat = data
    else:
        nchannels = data.shape[1]
        samples_flat = data.flatten()

    if format == "pcm" or format == "raw":
        # Flatten for miniaudio (interleaved)
        pcm_bytes = samples_flat.tobytes()

        if isinstance(file, io.BytesIO):
            file.write(pcm_bytes)
            file.seek(0)
        else:
            with open(file, "wb") as f:
                f.write(pcm_bytes)

    elif format == "wav":
        import array

        # Convert to array.array for miniaudio
        samples_array = array.array("h", samples_flat.tolist())

        # Create DecodedSoundFile
        sound = miniaudio.DecodedSoundFile(
            name="output",
            nchannels=nchannels,
            sample_rate=samplerate,
            sample_format=miniaudio.SampleFormat.SIGNED16,
            samples=samples_array,
        )

        if isinstance(file, io.BytesIO):
            # Write to temp file, read back to BytesIO
            import tempfile

            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                tmp_path = tmp.name
            try:
                miniaudio.wav_write_file(tmp_path, sound)
                with open(tmp_path, "rb") as f:
                    file.write(f.read())
                file.seek(0)
            finally:
                import os

                os.unlink(tmp_path)
        else:
            miniaudio.wav_write_file(str(file), sound)

    elif format in ("flac", "mp3", "ogg", "opus", "vorbis", "webm"):
        # Check for ffmpeg early to provide a clear error message
        if not _check_ffmpeg_available():
            import warnings

            warnings.warn(
                f"ffmpeg is required for {format.upper()} encoding but was not found. "
                "Install with: brew install ffmpeg (macOS) or sudo apt install ffmpeg (Ubuntu). "
                "Falling back will fail - consider using WAV format instead.",
                RuntimeWarning,
                stacklevel=2,
            )
        _encode_ffmpeg(data, samplerate, nchannels, file, format=format)
    else:
        raise ValueError(f"Unsupported output format: {format}")


# Convenience aliases to match soundfile API
def sf_read(
    file: Union[str, Path, io.BytesIO],
    always_2d: bool = False,
) -> Tuple[np.ndarray, int]:
    """Read audio file (soundfile-compatible API).

    This is a drop-in replacement for soundfile.read().

    Args:
        file: Path to audio file or BytesIO object.
        always_2d: If True, always return 2D array.

    Returns:
        Tuple of (data, samplerate).
    """
    return read(file, always_2d=always_2d, dtype="float64")


def sf_write(
    file: Union[str, Path, io.BytesIO],
    data: np.ndarray,
    samplerate: int,
    format: Optional[str] = None,
) -> None:
    """Write audio file (soundfile-compatible API).

    This is a drop-in replacement for soundfile.write().

    Args:
        file: Path to output file or BytesIO object.
        data: Audio data as numpy array.
        samplerate: Sample rate in Hz.
        format: Output format (wav, flac, mp3).
    """
    write(file, data, samplerate, format=format)