"""VoxCPM2 AudioVAE V2 decoder: latent frames ``.npy`` → 48 kHz float WAV (listening samples).

The input of each request is a latent array ``[num_frames, latent_dim]`` (float32, *VAE*
scale, i.e. already de-normalised) whose path is passed as ``audio_path`` (the worker
protocol's input file). Output: ``params["out_path"]``, 32-bit float mono WAV at the
decoder's output rate, written atomically; with ``params["out_path_16k"]`` also a 16 kHz
16-bit copy (for ASR checks). Same pinned weights as :mod:`vae_encode`.
"""

from __future__ import annotations

import importlib
from collections.abc import Mapping, Sequence
from pathlib import Path

from nida.bench.protocol import WorkRequest
from nida.bench.workers.runner import main_for
from nida.bench.workers.vae_voxcpm import VoxCPMVAEAdapter

__all__ = ["ASR_RATE_HZ", "VAEDecodeAdapter", "wav_out_path"]

#: Sample rate of the optional ASR copy.
ASR_RATE_HZ = 16000


def wav_out_path(request: WorkRequest) -> Path:
    """The ``.wav`` output path named by ``request.params["out_path"]``.

    Raises:
        ValueError: if the parameter is missing or does not end in ``.wav``.
    """
    raw = request.params.get("out_path")
    if not raw:
        raise ValueError(f"{request.clip_id}: missing params.out_path")
    path = Path(raw)
    if path.suffix != ".wav":
        raise ValueError(f"{request.clip_id}: params.out_path must end in .wav, got {raw!r}")
    return path


class VAEDecodeAdapter(VoxCPMVAEAdapter):
    name = "vae_decode"

    def _decode(self, request: WorkRequest) -> Mapping[str, object]:
        out_path = wav_out_path(request)
        torch = self._torch
        numpy = importlib.import_module("numpy")
        soundfile = importlib.import_module("soundfile")
        latents = numpy.load(request.audio_path, allow_pickle=False)
        if latents.ndim != 2 or latents.shape[0] < 1:
            raise ValueError(f"{request.clip_id}: expected [frames, dim], got {latents.shape}")
        if not bool(numpy.isfinite(latents).all()):
            raise ValueError(f"{request.clip_id}: non-finite latents")
        z = torch.from_numpy(latents.astype("float32")).transpose(0, 1).unsqueeze(0)
        with torch.inference_mode():
            decoded = self._vae.decode(z.to(self._device))  # (1, 1, samples)
        out = decoded.reshape(-1).float().cpu().numpy()
        rate_hz = int(self._vae.out_sample_rate)
        target = Path(out_path)
        target.parent.mkdir(parents=True, exist_ok=True)
        tmp = target.with_name(f".{target.name}.tmp.wav")  # never leave a partial output
        soundfile.write(str(tmp), out, rate_hz, subtype="FLOAT")
        tmp.replace(target)
        payload: dict[str, object] = {}
        if request.params.get("out_path_16k"):
            functional = importlib.import_module("torchaudio.functional")
            low = functional.resample(torch.from_numpy(out), rate_hz, ASR_RATE_HZ).numpy()
            target_16k = Path(request.params["out_path_16k"])
            tmp_16k = target_16k.with_name(f".{target_16k.name}.tmp.wav")
            soundfile.write(str(tmp_16k), numpy.clip(low, -1.0, 1.0), ASR_RATE_HZ, "PCM_16")
            tmp_16k.replace(target_16k)
            payload["out_path_16k"] = str(target_16k)
        return {
            **payload,
            "out_path": str(out_path),
            "out_sample_rate_hz": rate_hz,
            "num_samples": int(out.shape[0]),
            "num_latent_frames": int(latents.shape[0]),
        }

    def run_batch(self, batch: Sequence[WorkRequest]) -> list[Mapping[str, object]]:
        return [self._decode(r) for r in batch]


if __name__ == "__main__":
    raise SystemExit(main_for(VAEDecodeAdapter()))