from __future__ import annotations

import contextlib
import gc
import hashlib
import io
import os
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import numpy as np
import soundfile as sf

from ..tts.tts_engine import _SilenceStderrFD


@dataclass(frozen=True)
class ChromaArtifacts:
    root: Path
    model_id: str
    revision: Optional[str]


class ChromaVoiceCloningEngine:
    """In-process Chroma voice cloning engine (optional; GPU-heavy).

    Design principles
    -----------------
    - Never download weights implicitly from speak()/speak_to_bytes().
    - Provide explicit prefetch hooks.
    - Keep imports lightweight until the engine is actually used.
    """

    DEFAULT_MODEL_ID = "FlashLabs/Chroma-4B"
    # Pin a known revision by default (safer than floating remote code).
    DEFAULT_REVISION = "864b4aea0c1359f91af62f1367df64657dc5e90f"

    def __init__(
        self,
        *,
        debug: bool = False,
        device: str = "auto",
        model_id: str = DEFAULT_MODEL_ID,
        revision: str | None = DEFAULT_REVISION,
        temperature: float = 0.7,
        top_k: int = 50,
        do_sample: bool = True,
        max_new_tokens_per_chunk: int = 512,
    ):
        self.debug = bool(debug)
        self._device_pref = str(device or "auto")
        self._model_id = str(model_id or self.DEFAULT_MODEL_ID)
        self._revision = revision if revision else None

        self._temperature = float(temperature)
        self._top_k = int(top_k)
        self._do_sample = bool(do_sample)
        self._max_new_tokens_per_chunk = int(max_new_tokens_per_chunk)

        self._model = None
        self._processor = None
        self._resolved_device = None

    def unload(self) -> None:
        """Best-effort release of loaded model/processor to free memory."""
        self._model = None
        self._processor = None
        self._resolved_device = None
        try:
            gc.collect()
        except Exception:
            pass
        try:
            import torch

            # Best-effort GPU/MPS cache release.
            try:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass
            try:
                if hasattr(torch, "mps") and torch.backends.mps.is_available():
                    torch.mps.empty_cache()
            except Exception:
                pass
        except Exception:
            pass

    def runtime_info(self) -> Dict[str, object]:
        info: Dict[str, object] = {
            "model_id": self._model_id,
            "revision": self._revision,
            "requested_device": self._device_pref,
            "resolved_device": self._resolved_device,
            "temperature": self._temperature,
            "top_k": self._top_k,
            "do_sample": self._do_sample,
            "max_new_tokens_per_chunk": self._max_new_tokens_per_chunk,
        }
        try:
            import torch

            info["torch_version"] = getattr(torch, "__version__", "?")
            info["cuda_available"] = bool(torch.cuda.is_available())
            try:
                info["mps_available"] = bool(torch.backends.mps.is_available())
            except Exception:
                info["mps_available"] = False
        except Exception:
            pass
        return info

    def set_quality_preset(self, preset: str) -> None:
        p = (preset or "").strip().lower()
        if p not in ("fast", "balanced", "high"):
            raise ValueError("preset must be one of: fast|balanced|high")
        if p == "fast":
            self._temperature = 0.6
            self._top_k = 30
            self._max_new_tokens_per_chunk = 384
        elif p == "balanced":
            self._temperature = 0.7
            self._top_k = 50
            self._max_new_tokens_per_chunk = 512
        else:
            self._temperature = 0.75
            self._top_k = 80
            self._max_new_tokens_per_chunk = 768

    def _artifact_root(self) -> Path:
        cache_dir = Path(os.path.expanduser("~/.cache/abstractvoice/chroma"))
        cache_dir.mkdir(parents=True, exist_ok=True)
        return cache_dir

    def _resolve_chroma_artifacts_local(self) -> ChromaArtifacts:
        root = self._artifact_root()
        required = [
            "config.json",
            "processor_config.json",
            "model.safetensors.index.json",
            "modeling_chroma.py",
            "processing_chroma.py",
            "configuration_chroma.py",
        ]
        missing = [name for name in required if not (root / name).exists()]
        if missing:
            raise RuntimeError(
                "Chroma artifacts are not present locally.\n"
                "Prefetch explicitly (outside the REPL):\n"
                "  abstractvoice-prefetch --chroma\n"
                "or:\n"
                "  python -m abstractvoice download --chroma\n"
                f"Looked under: {root}\n"
                f"Missing: {', '.join(missing)}"
            )
        return ChromaArtifacts(root=root, model_id=self._model_id, revision=self._revision)

    def ensure_chroma_artifacts_downloaded(self) -> ChromaArtifacts:
        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
        try:
            from huggingface_hub import snapshot_download
        except Exception as e:
            raise RuntimeError(
                "huggingface_hub is required to download Chroma artifacts.\n"
                "Install with: pip install huggingface_hub"
            ) from e

        import warnings

        root = self._artifact_root()
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                category=UserWarning,
                message=r".*local_dir_use_symlinks.*deprecated.*",
            )
            snapshot_download(
                repo_id=self._model_id,
                revision=self._revision,
                local_dir=str(root),
            )
        return self._resolve_chroma_artifacts_local()

    def are_chroma_artifacts_available(self) -> bool:
        root = self._artifact_root()
        return bool((root / "config.json").exists() and (root / "model.safetensors.index.json").exists())

    def _resolve_device(self) -> str:
        if self._device_pref and self._device_pref != "auto":
            return str(self._device_pref)
        try:
            from ..compute import best_torch_device

            return best_torch_device()
        except Exception:
            return "cpu"

    def _ensure_chroma_runtime(self) -> None:
        try:
            import importlib.util

            if importlib.util.find_spec("torch") is None:
                raise ImportError("torch not installed")
            if importlib.util.find_spec("transformers") is None:
                raise ImportError("transformers not installed")
        except Exception as e:
            raise RuntimeError(
                "Chroma requires the optional dependency group.\n"
                "Install with:\n"
                "  pip install \"abstractvoice[chroma]\"\n"
                f"Original error: {e}"
            ) from e

    def _ensure_model_loaded(self) -> None:
        if self._model is not None and self._processor is not None and self._resolved_device is not None:
            return

        self._ensure_chroma_runtime()
        artifacts = self._resolve_chroma_artifacts_local()

        # Keep interactive UX quiet by default (the REPL provides its own spinner).
        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
        os.environ.setdefault("TRANSFORMERS_NO_TQDM", "1")

        # Keep interactive UX quiet by default (the REPL provides its own spinner).
        try:
            from transformers.utils import logging as hf_logging

            hf_logging.disable_progress_bar()
            if not self.debug:
                hf_logging.set_verbosity_error()
        except Exception:
            pass

        import torch
        from transformers import AutoModelForCausalLM, AutoProcessor

        device = self._resolve_device()
        self._resolved_device = device

        torch_dtype = None
        if device == "cuda":
            torch_dtype = torch.bfloat16
        elif device == "mps":
            torch_dtype = torch.float16

        device_map = None
        # `device_map="auto"` is most reliable on CUDA; MPS/CPU should load normally then move.
        if self._device_pref == "auto" and device == "cuda":
            device_map = "auto"

        # transformers 5.0.0 deprecates `torch_dtype` in favor of `dtype`.
        with warnings.catch_warnings():
            # Torch/Torchaudio/Transformers often emit noisy warnings that are not actionable
            # for interactive users. Keep them hidden unless debug is enabled.
            if not self.debug:
                warnings.filterwarnings("ignore", category=UserWarning, module=r"torchaudio\..*")
                warnings.filterwarnings("ignore", category=UserWarning, message=r".*TorchCodec.*")
                warnings.filterwarnings("ignore", category=UserWarning, message=r".*output_attentions.*")
            stdout_ctx = contextlib.redirect_stdout(io.StringIO()) if not self.debug else contextlib.nullcontext()
            with stdout_ctx:
                with _SilenceStderrFD(enabled=not self.debug):
                    try:
                        model = AutoModelForCausalLM.from_pretrained(
                            str(artifacts.root),
                            trust_remote_code=True,
                            device_map=device_map,
                            dtype=torch_dtype,
                        )
                    except TypeError:
                        model = AutoModelForCausalLM.from_pretrained(
                            str(artifacts.root),
                            trust_remote_code=True,
                            device_map=device_map,
                            torch_dtype=torch_dtype,
                        )
        model.eval()

        stdout_ctx = contextlib.redirect_stdout(io.StringIO()) if not self.debug else contextlib.nullcontext()
        with stdout_ctx:
            with _SilenceStderrFD(enabled=not self.debug):
                processor = AutoProcessor.from_pretrained(str(artifacts.root), trust_remote_code=True)

        if device != "cuda" and device_map is None:
            try:
                model.to(device)
            except Exception:
                pass

        # Ensure tied weights are applied (Chroma relies on tied audio embeddings).
        try:
            if hasattr(model, "tie_weights"):
                model.tie_weights()
        except Exception:
            pass
        # Some releases log this key as MISSING even when it is meant to be tied.
        # Ensure backbone audio embeddings match decoder embeddings (best-effort).
        try:
            if hasattr(model, "backbone") and hasattr(model, "decoder"):
                b = getattr(model.backbone, "audio_embedding", None)
                d = getattr(model.decoder, "audio_embedding", None)
                if b is not None and d is not None:
                    bw = getattr(b, "embed_audio_tokens", None)
                    dw = getattr(d, "embed_audio_tokens", None)
                    if bw is not None and dw is not None and hasattr(bw, "weight") and hasattr(dw, "weight"):
                        if getattr(bw.weight, "shape", None) == getattr(dw.weight, "shape", None):
                            bw.weight.data.copy_(dw.weight.data)
                    # Prefer a hard tie (shared module) to avoid one side staying randomly
                    # initialized when the checkpoint omits a tied key.
                    try:
                        if bw is not None and dw is not None and bw is not dw:
                            b.embed_audio_tokens = dw
                    except Exception:
                        pass
        except Exception:
            pass

        self._patch_generation_compat(model)

        self._model = model
        self._processor = processor

    def _estimate_max_new_tokens(self, text: str, model) -> int:
        """Estimate a safe `max_new_tokens` (audio frames) for a text batch.

        Chroma audio frames decode at roughly:
          frames_per_second = sampling_rate / audio_frame_freq
        For the default config this is ~12.5 fps, so `512` frames ~= 41s.
        Without a good stopping signal, the model may generate until `max_new_tokens`,
        so we cap it based on expected speech duration to avoid long noisy tails.
        """
        s = " ".join(str(text or "").split()).strip()
        words = len(s.split()) if s else 0
        if words <= 0:
            # Fallback: roughly 5 chars/word in English.
            words = max(1, int(round(len(s) / 5.0))) if s else 1

        # Conservative speech-rate assumption (words/sec).
        wps = 3.2
        est_s = float(words) / float(wps)
        # Small pad to avoid truncation on short utterances.
        est_s += 0.4

        sr = 24000
        frame_freq = 1920
        try:
            cfg = getattr(model, "config", None)
            if cfg is not None:
                frame_freq = int(getattr(cfg, "audio_frame_freq", frame_freq) or frame_freq)
                codec_cfg = getattr(cfg, "codec_config", None)
                if codec_cfg is not None:
                    if isinstance(codec_cfg, dict):
                        sr = int(codec_cfg.get("sampling_rate", sr) or sr)
                    else:
                        sr = int(getattr(codec_cfg, "sampling_rate", sr) or sr)
        except Exception:
            pass

        fps = float(sr) / float(frame_freq) if frame_freq else 12.5
        # Add slack so we still have room for slower speech.
        frames = int(round(est_s * fps * 1.25))
        frames = max(64, frames)
        cap = int(self._max_new_tokens_per_chunk) if self._max_new_tokens_per_chunk else frames
        return int(min(frames, cap))

    def _patch_generation_compat(self, model) -> None:
        """Patch known transformers incompatibilities for Chroma remote code.

        Upstream Chroma code (pinned revision) expects a `use_model_defaults` positional arg
        in `GenerationMixin._prepare_generation_config`. transformers 5.0.0 removed it.
        """
        try:
            import inspect

            from transformers.generation import GenerationMode
            from transformers.generation.utils import GenerationMixin

            base_sig = inspect.signature(GenerationMixin._prepare_generation_config)
            if "use_model_defaults" in base_sig.parameters:
                return

            chroma_gen_cls = None
            for cls in model.__class__.mro():
                if cls.__name__ == "ChromaGenerationMixin":
                    chroma_gen_cls = cls
                    break
            if chroma_gen_cls is None:
                return

            current = getattr(chroma_gen_cls, "_prepare_generation_config", None)
            if current is None or getattr(current, "_abstractvoice_patched", False):
                return

            def _patched_prepare_generation_config(
                self, generation_config=None, use_model_defaults=None, **kwargs
            ):
                depth_decoder_kwargs = {k[len("decoder_") :]: v for k, v in kwargs.items() if k.startswith("decoder_")}
                kwargs = {k: v for k, v in kwargs.items() if not k.startswith("decoder_")}

                # transformers 5.0.0 removed the `use_model_defaults` positional arg; keep compatibility.
                try:
                    generation_config_out, model_kwargs = super(chroma_gen_cls, self)._prepare_generation_config(
                        generation_config, use_model_defaults, **kwargs
                    )
                except TypeError:
                    generation_config_out, model_kwargs = super(chroma_gen_cls, self)._prepare_generation_config(
                        generation_config, **kwargs
                    )

                try:
                    self.decoder.generation_config.update(**depth_decoder_kwargs)
                except Exception:
                    pass

                try:
                    decoder_min_new_tokens = getattr(self.decoder.generation_config, "min_new_tokens") or (
                        self.decoder.config.audio_num_codebooks - 1
                    )
                    decoder_max_new_tokens = getattr(self.decoder.generation_config, "max_new_tokens") or (
                        self.decoder.config.audio_num_codebooks - 1
                    )

                    if {decoder_min_new_tokens, decoder_max_new_tokens} != {self.decoder.config.audio_num_codebooks - 1}:
                        raise ValueError(
                            "decoder_generation_config's min_new_tokens "
                            f"({decoder_min_new_tokens}) and max_new_tokens ({decoder_max_new_tokens}) "
                            f"must be equal to self.config.num_codebooks - 1 ({self.decoder.config.audio_num_codebooks - 1})"
                        )
                    elif getattr(self.decoder.generation_config, "return_dict_in_generate", False):
                        self.decoder.generation_config.return_dict_in_generate = False
                except Exception:
                    pass

                # Monkey patch the get_generation_mode method to support Chroma model.
                try:
                    original_get_generation_mode = generation_config_out.get_generation_mode

                    def patched_get_generation_mode(assistant_model=None):
                        generation_mode = original_get_generation_mode(assistant_model)
                        if generation_mode not in (GenerationMode.GREEDY_SEARCH, GenerationMode.SAMPLE):
                            raise ValueError(
                                f"Generation mode {generation_mode} is not supported for Chroma model. "
                                "Please set generation parameters to use greedy or sampling generation."
                            )
                        return generation_mode

                    generation_config_out.get_generation_mode = patched_get_generation_mode
                except Exception:
                    pass

                return generation_config_out, model_kwargs

            _patched_prepare_generation_config._abstractvoice_patched = True  # type: ignore[attr-defined]
            chroma_gen_cls._prepare_generation_config = _patched_prepare_generation_config  # type: ignore[assignment]
        except Exception:
            # Best-effort only: do not break core flows if patching fails.
            return

    def _select_prompt_audio(self, reference_paths: Iterable[str | Path]) -> Path:
        paths = [Path(p) for p in reference_paths]
        if not paths:
            raise ValueError("reference_paths must contain at least one path")
        for p in paths:
            if not p.exists():
                raise FileNotFoundError(str(p))
        # Chroma prompt_audio currently expects a single audio file.
        return paths[0]

    def _prepare_prompt_audio_for_processor(self, prompt_audio: Path) -> Path:
        """Best-effort prompt-audio normalization for Chroma.

        Chroma's processor loads audio via torchaudio and resamples internally.
        We keep this preprocessing minimal (mono + 24kHz + PCM16 WAV) to reduce
        backend variability without altering voice characteristics.
        """
        try:
            if not prompt_audio.exists():
                return prompt_audio

            # Cache by (path, size, mtime) to avoid recomputing every turn.
            try:
                st = prompt_audio.stat()
                key = f"{prompt_audio.resolve()}|{st.st_size}|{st.st_mtime_ns}"
            except Exception:
                key = str(prompt_audio)
            h = hashlib.sha1(key.encode("utf-8", errors="ignore")).hexdigest()[:16]
            out_dir = self._artifact_root() / "prompt_cache"
            out_dir.mkdir(parents=True, exist_ok=True)
            out_path = out_dir / f"prompt_{h}.wav"
            if out_path.exists():
                return out_path

            audio, sr = sf.read(str(prompt_audio), always_2d=True, dtype="float32")
            mono = np.mean(audio, axis=1).astype(np.float32).reshape(-1)

            target_sr = 24000
            if int(sr) != int(target_sr):
                # Prefer torchaudio's sinc-based resampling when available (better fidelity than linear).
                try:
                    import torch
                    import torchaudio

                    t = torch.from_numpy(mono).unsqueeze(0)
                    t = torchaudio.functional.resample(t, orig_freq=int(sr), new_freq=int(target_sr))
                    mono = t.squeeze(0).detach().cpu().numpy().astype(np.float32)
                except Exception:
                    from ..audio.resample import linear_resample_mono

                    mono = linear_resample_mono(mono, int(sr), int(target_sr))

            # Cap extreme prompts for runtime stability; upstream examples go up to ~26s.
            max_seconds = 30.0
            max_samples = int(round(float(max_seconds) * float(target_sr)))
            if mono.size > max_samples:
                mono = mono[:max_samples]

            try:
                mono = np.nan_to_num(mono, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
                mono = np.clip(mono, -1.0, 1.0).astype(np.float32)
            except Exception:
                pass

            sf.write(str(out_path), mono, int(target_sr), format="WAV", subtype="PCM_16")
            return out_path
        except Exception:
            # If anything goes wrong, fall back to the original path.
            return prompt_audio

    def _build_conversation(self, text: str) -> List[List[dict]]:
        # Match upstream framing ("Chroma" virtual human) while keeping strict TTS behavior.
        system_prompt = (
            "You are Chroma, an advanced virtual human created by the FlashLabs. "
            "You possess the ability to understand auditory inputs and generate both text and speech. "
            "For this request, do not answer or paraphrase. Read the user's text aloud exactly as written."
        )
        return [[
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
            {"role": "user", "content": [{"type": "text", "text": str(text)}]},
        ]]

    def infer_to_wav_bytes(
        self,
        *,
        text: str,
        reference_paths: Iterable[str | Path],
        reference_text: Optional[str] = None,
        speed: Optional[float] = None,
    ) -> bytes:
        chunks = []
        sr_out = 24000
        for chunk, sr in self.infer_to_audio_chunks(
            text=text,
            reference_paths=reference_paths,
            reference_text=reference_text,
            speed=speed,
        ):
            chunks.append(np.asarray(chunk, dtype=np.float32).reshape(-1))
            sr_out = int(sr)
        audio = np.concatenate(chunks) if chunks else np.zeros((0,), dtype=np.float32)
        buf = io.BytesIO()
        sf.write(buf, audio, int(sr_out), format="WAV", subtype="PCM_16")
        return buf.getvalue()

    def infer_to_audio_chunks(
        self,
        *,
        text: str,
        reference_paths: Iterable[str | Path],
        reference_text: Optional[str] = None,
        speed: Optional[float] = None,
        max_chars: int = 240,
    ):
        self._ensure_model_loaded()

        if not reference_text or not str(reference_text).strip():
            raise RuntimeError(
                "Missing reference_text for Chroma cloning.\n"
                "If you're using VoiceCloner/VoiceManager, reference_text should be auto-generated and cached.\n"
                "If you're calling this engine directly, provide reference_text or set it via the voice store."
            )

        prompt_audio = self._select_prompt_audio(reference_paths)
        prompt_audio = self._prepare_prompt_audio_for_processor(prompt_audio)
        prompt_text = str(reference_text).strip()

        import re

        def _split_batches(s: str, limit: int) -> List[str]:
            s = " ".join(str(s).replace("\n", " ").split()).strip()
            if not s:
                return []
            sentences = re.split(r"(?<=[\\.!\\?\\。])\\s+", s)
            out: List[str] = []
            cur_s = ""
            for sent in sentences:
                sent = sent.strip()
                if not sent:
                    continue
                if len(sent) > limit:
                    words = sent.split(" ")
                    tmp = ""
                    for w in words:
                        cand = (tmp + " " + w).strip()
                        if len(cand) <= limit:
                            tmp = cand
                        else:
                            if tmp:
                                out.append(tmp)
                            tmp = w
                    if tmp:
                        out.append(tmp)
                    continue
                cand = (cur_s + " " + sent).strip()
                if len(cand) <= limit:
                    cur_s = cand
                else:
                    if cur_s:
                        out.append(cur_s)
                    cur_s = sent
            if cur_s:
                out.append(cur_s)
            return out

        batches = _split_batches(text, int(max_chars)) or [" "]

        model = self._model
        processor = self._processor
        if model is None or processor is None:
            raise RuntimeError("Chroma model not loaded")

        import torch

        for batch_text in batches:
            conversation = self._build_conversation(batch_text)
            with warnings.catch_warnings():
                if not self.debug:
                    warnings.filterwarnings("ignore", category=UserWarning, module=r"torchaudio\..*")
                    warnings.filterwarnings("ignore", category=UserWarning, message=r".*TorchCodec.*")
                stdout_ctx = contextlib.redirect_stdout(io.StringIO()) if not self.debug else contextlib.nullcontext()
                with stdout_ctx:
                    with _SilenceStderrFD(enabled=not self.debug):
                        inputs = processor(
                            conversation,
                            add_generation_prompt=True,
                            tokenize=False,
                            prompt_audio=[str(prompt_audio)],
                            prompt_text=[prompt_text],
                            return_tensors="pt",
                        )
            device = getattr(model, "device", None) or torch.device("cpu")
            # Match upstream usage: `processor(...).to(device)` (device move only).
            # We'll cast floating-point inputs to the model's dtype right before
            # generation to avoid fp16/bf16 dtype mismatches on GPU/MPS.
            try:
                inputs = inputs.to(device)
            except Exception:
                pass
            moved = dict(inputs)
            # Chroma weights often run in fp16/bf16 on GPU/MPS; ensure floating-point
            # inputs match model weight dtype to avoid runtime dtype mismatches like:
            # "Input type (float) and bias type (c10::Half) should be the same".
            try:
                target_dtype = getattr(model, "dtype", None)
                if target_dtype is None:
                    target_dtype = next(iter(model.parameters())).dtype  # type: ignore[assignment]
                if target_dtype is not None:
                    for k, v in list(moved.items()):
                        if isinstance(v, torch.Tensor) and v.is_floating_point() and v.dtype != target_dtype:
                            moved[k] = v.to(dtype=target_dtype)
            except Exception:
                pass

            max_new = self._estimate_max_new_tokens(batch_text, model)
            with warnings.catch_warnings():
                if not self.debug:
                    warnings.filterwarnings("ignore", category=UserWarning, message=r".*output_attentions.*")
                    warnings.filterwarnings("ignore", category=UserWarning, message=r".*sdp attention.*")
                stdout_ctx = contextlib.redirect_stdout(io.StringIO()) if not self.debug else contextlib.nullcontext()
                with stdout_ctx:
                    with _SilenceStderrFD(enabled=not self.debug):
                        with torch.inference_mode():
                            out = model.generate(
                                **moved,
                                max_new_tokens=int(max_new),
                                do_sample=bool(self._do_sample),
                                temperature=float(self._temperature),
                                top_k=int(self._top_k),
                                use_cache=True,
                                output_attentions=False,
                                output_audio=True,
                            )
            audio_list = out.audio if hasattr(out, "audio") else out
            if not audio_list:
                continue
            a = audio_list[0]
            if isinstance(a, (list, tuple)) and a:
                a = a[0]
            if isinstance(a, torch.Tensor):
                arr = a.detach().cpu().numpy()
            else:
                arr = np.asarray(a)

            mono = np.asarray(arr, dtype=np.float32).reshape(-1)
            sr = 24000

            if speed and speed != 1.0:
                try:
                    from ..tts.tts_engine import apply_speed_without_pitch_change

                    mono = apply_speed_without_pitch_change(mono, float(speed), sr=sr)
                except Exception:
                    pass

            # Sanitize before any normalization (NaNs/Infs produce loud artifacts).
            try:
                mono = np.nan_to_num(mono, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
            except Exception:
                pass

            # Chroma's codec output can exceed [-1, 1] slightly; normalize to avoid
            # hard clipping/distortion in playback and PCM encoders.
            try:
                peak = float(np.max(np.abs(mono))) if mono.size else 0.0
                if peak > 1.0:
                    mono = mono / peak
            except Exception:
                pass

            # Final clamp for safety (avoids stray overs on some backends).
            try:
                mono = np.clip(mono, -1.0, 1.0).astype(np.float32)
            except Exception:
                pass

            yield mono, sr
