"""Robust audio preprocessing.

Libsyn MP3s confuse faster-whisper's decoder on this box (bogus multi-channel
allocation). We therefore demux+decode ourselves with PyAV in bounded chunks,
downmix to mono, resample to 16 kHz s16, and write a plain WAV that Whisper
ingests instantly.

Full-length prep of a 3 h episode takes ~1 min and ~340 MB; the file is
deleted by transcribe.py afterwards unless DS_PREP_KEEP=1.
"""
import os
import wave

import av
import numpy as np

from config import AUDIO_DIR
from download import slug as ep_slug

PREP_DIR = os.path.join(AUDIO_DIR, "_prep")
TARGET_RATE = 16000


def prepared_path(ep: dict) -> str:
    return os.path.join(PREP_DIR, ep_slug(ep) + ".16k.wav")


def prepare_audio(src_mp3: str, dst_wav: str,
                  max_seconds: float | None = None) -> str:
    os.makedirs(PREP_DIR, exist_ok=True)
    if os.path.exists(dst_wav):
        return dst_wav

    resampler = av.AudioResampler(format="s16", layout="mono",
                                  rate=TARGET_RATE)
    total_written = 0
    limit = int(max_seconds * TARGET_RATE) if max_seconds else None

    tmp = dst_wav + ".tmp"
    with av.open(src_mp3, metadata_errors="ignore") as container:
        # pick the first real audio stream explicitly
        streams = [s for s in container.streams.audio]
        if not streams:
            raise ValueError(f"no audio stream in {src_mp3}")
        stream = streams[0]

        import wave
        wf = wave.open(tmp, "wb")
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(TARGET_RATE)
        try:
            for frame in container.decode(stream):
                resampled = resampler.resample(frame)
                for rf in resampled:
                    arr = rf.to_ndarray()
                    if arr.ndim == 2:      # (1, samples) planar mono
                        arr = arr[0]
                    pcm = arr.astype(np.int16).tobytes()
                    n = len(pcm) // 2
                    if limit and total_written + n > limit:
                        n = limit - total_written
                        pcm = pcm[: n * 2]
                    wf.writeframes(pcm)
                    total_written += n
                    if limit and total_written >= limit:
                        raise StopIteration
        except StopIteration:
            pass
        finally:
            wf.close()

    os.replace(tmp, dst_wav)
    return dst_wav


def load_pcm(dst_wav: str) -> np.ndarray:
    """Load a prepared 16 kHz s16 mono wav as float32 in [-1, 1]."""
    with wave.open(dst_wav, "rb") as wf:
        assert wf.getnchannels() == 1 and wf.getframerate() == TARGET_RATE
        frames = wf.readframes(wf.getnframes())
    return np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0