"""
audio_io.py — Lectura y escritura de audio sin sorpresas.

- Lee WAV/FLAC/AIFF/OGG/MP3 con soundfile; si falla (M4A, AAC…), usa ffmpeg.
- Devuelve siempre float32 con forma (muestras, canales).
- Escribe en el formato pedido; al bajar a 16 bits aplica dither TPDF.
"""

from __future__ import annotations

import io
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import soundfile as sf

AUDIO_EXTENSIONS = [".wav", ".flac", ".aif", ".aiff", ".ogg", ".opus", ".mp3", ".m4a", ".aac", ".wma"]

# Subtipos que se pueden conservar tal cual al escribir el resultado.
_LOSSLESS = {
    "WAV": {"PCM_16", "PCM_24", "PCM_32", "FLOAT", "DOUBLE"},
    "FLAC": {"PCM_16", "PCM_24"},
    "AIFF": {"PCM_16", "PCM_24", "PCM_32", "FLOAT"},
}
_EXT = {"WAV": ".wav", "FLAC": ".flac", "AIFF": ".aiff"}

SUBTYPE_LABEL = {
    "PCM_16": "PCM entero 16 bits",
    "PCM_24": "PCM entero 24 bits",
    "PCM_32": "PCM entero 32 bits",
    "FLOAT": "coma flotante 32 bits",
    "DOUBLE": "coma flotante 64 bits",
}
SUBTYPE_BITS = {"PCM_U8": 8, "PCM_S8": 8, "PCM_16": 16, "PCM_24": 24, "PCM_32": 32, "FLOAT": 32, "DOUBLE": 64}


@dataclass
class AudioMeta:
    name: str
    container: str      # WAV, FLAC, MP3…
    subtype: str        # PCM_24, FLOAT, MPEG_LAYER_III…
    samplerate: int
    channels: int
    frames: int

    @property
    def duration(self) -> float:
        return self.frames / self.samplerate if self.samplerate else 0.0

    @property
    def bits(self) -> int | None:
        return SUBTYPE_BITS.get(self.subtype)

    @property
    def is_integer_pcm(self) -> bool:
        return self.subtype.startswith("PCM_")

    def describe(self) -> str:
        sub = SUBTYPE_LABEL.get(self.subtype, self.subtype)
        ch = "mono" if self.channels == 1 else f"{self.channels} canales"
        return f"{self.container} · {sub} · {self.samplerate} Hz · {ch} · {fmt_duration(self.duration)}"


def fmt_duration(sec: float) -> str:
    sec = max(0.0, sec)
    h, rem = divmod(int(round(sec)), 3600)
    m, s = divmod(rem, 60)
    return f"{h}:{m:02d}:{s:02d}" if h else f"{m}:{s:02d}"


def _read_ffmpeg(path: str) -> tuple[np.ndarray, int]:
    if not shutil.which("ffmpeg"):
        raise RuntimeError("Formato no soportado y ffmpeg no está instalado.")
    proc = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", path, "-vn", "-f", "wav", "-acodec", "pcm_f32le", "-"],
        capture_output=True, check=False,
    )
    if proc.returncode != 0 or not proc.stdout:
        raise RuntimeError(f"No se pudo leer el audio: {proc.stderr.decode(errors='ignore')[:300]}")
    data, sr = sf.read(io.BytesIO(proc.stdout), dtype="float32", always_2d=True)
    return data, sr


def read_audio(path: str | Path, start: float | None = None, stop: float | None = None):
    """Devuelve (datos float32 (n, ch), sr, AudioMeta). start/stop en segundos (opcional)."""
    path = str(path)
    name = Path(path).name
    try:
        info = sf.info(path)
        meta = AudioMeta(name, info.format, info.subtype, info.samplerate, info.channels, info.frames)
        a = int(start * info.samplerate) if start else 0
        b = int(stop * info.samplerate) if stop else None
        data, sr = sf.read(path, start=a, stop=b, dtype="float32", always_2d=True)
    except Exception:
        data, sr = _read_ffmpeg(path)
        ext = Path(path).suffix.lstrip(".").upper() or "?"
        meta = AudioMeta(name, ext, "COMPRIMIDO", sr, data.shape[1], data.shape[0])
        if start or stop:
            a = int((start or 0) * sr)
            b = int(stop * sr) if stop else None
            data = data[a:b]
    return data, sr, meta


def read_meta(path: str | Path) -> AudioMeta:
    try:
        info = sf.info(str(path))
        return AudioMeta(Path(path).name, info.format, info.subtype, info.samplerate, info.channels, info.frames)
    except Exception:
        _, _, meta = read_audio(path)
        return meta


def output_format(meta: AudioMeta, container: str | None, subtype: str | None) -> tuple[str, str, str]:
    """
    Decide (contenedor, subtipo, extensión) de salida.
    container/subtype None = "como el original"; si el original es comprimido → WAV 24 bits.
    """
    if container is None:
        container = meta.container if meta.container in _LOSSLESS else "WAV"
    container = container.upper()
    if subtype is None:
        subtype = meta.subtype if meta.subtype in _LOSSLESS.get(container, set()) else "PCM_24"
    if subtype not in _LOSSLESS[container]:
        subtype = "PCM_24"
    return container, subtype, _EXT[container]


def write_audio(path: str | Path, data: np.ndarray, sr: int, container: str = "WAV", subtype: str = "PCM_24",
                seed: int = 0) -> None:
    """Escribe (n, ch) float. En formatos enteros recorta a [-1, 1) y en 16 bits aplica dither TPDF."""
    data = np.asarray(data, dtype=np.float64)
    if subtype.startswith("PCM_"):
        if subtype in ("PCM_16", "PCM_U8", "PCM_S8"):
            lsb = 1.0 / (2 ** (SUBTYPE_BITS[subtype] - 1))
            rng = np.random.default_rng(seed)
            data = data + (rng.random(data.shape) - rng.random(data.shape)) * lsb
        bits = SUBTYPE_BITS[subtype]
        data = np.clip(data, -1.0, 1.0 - 1.0 / 2 ** (bits - 1))
    sf.write(str(path), data, sr, format=container, subtype=subtype)