"""
Vocence TTS engine — Qwen3-TTS sequential best-of-up-to-2 (single process).

Architectural note:
  An earlier version spawned subprocess workers for true GPU-parallel best-of-N.
  In the chute container, child Python processes segfault during `<frozen site>`
  initialization (sandboxed FS / restricted privileges → child can't even start).
  Chute permits exactly one Python process running miner.py — so best-of-N must
  happen sequentially in the main process.

Tradeoffs:
  - Sequential generation only (max two samples per request when time allows).
  - With sampling enabled (do_sample=True, temp=0.9, ...), each call is typically
    ~15-45s on the chute GPU.
  - If the first sample finishes in under 45s, a second sample is generated and
    UTMOS + Whisper pick the better of the two. If the first sample takes ≥45s,
    it is returned immediately (no second sample, no scoring) to save latency.

Vocence contract (do not change):
    Miner(path_hf_repo: Path)
    warmup() -> None
    generate_wav(instruction: str, text: str) -> tuple[np.ndarray, int]
"""
from __future__ import annotations

import json
import re
import threading
import time
from pathlib import Path
from typing import Any, Mapping, Optional

import numpy as np


# ---------------------------------------------------------------------------
# Generation kwargs. Sampling matches qwen-tts example defaults; max_new_tokens
# is taken from the snapshot's generation_config.json (see Miner._max_new_tokens).
# ---------------------------------------------------------------------------
_GEN_TEMPERATURE = 0.65
_GEN_TOP_P = 1.0
_GEN_TOP_K = 50
_GEN_REPETITION_PENALTY = 1.05
# Fallback only if generation_config.json has no max_new_tokens (snapshot default is 8192).
_DEFAULT_MAX_NEW_TOKENS = 2048
_GEN_DO_SAMPLE = True

# At most two sequential candidates when the first finishes quickly enough.
_DEFAULT_NUM_CANDIDATES = 2

# If the first generation completes in under this many seconds, optionally
# generate a second candidate and score both; otherwise return the first only.
_FIRST_GEN_FAST_THRESHOLD_SEC = 45.0

# Validity thresholds for the candidate filter.
_MIN_DURATION_SEC = 2.0
_MAX_DURATION_SEC = 29.5
_MIN_RMS = 1e-3
_MAX_PEAK = 0.99


# ---------------------------------------------------------------------------
# Repo / config helpers
# ---------------------------------------------------------------------------

_CONFIG_NAME = "config.json"
_VOCENCE_YAML = "vocence_config.yaml"
_GENERATION_CONFIG_NAME = "generation_config.json"


# qwen-tts ships Qwen3TTSSpeakerEncoderConfig.__init__ without **kwargs, but
# transformers' to_dict() serializes nested PretrainedConfigs with model_type/
# transformers_version/torch_dtype embedded. On load, Qwen3TTSConfig.__init__
# splat-passes that dict into the sub-config and TypeError fires on the very
# first unknown key. Wrap __init__ to drop those metadata keys before forwarding.
from qwen_tts.core.models import configuration_qwen3_tts as _qcfg

_orig_se_init = _qcfg.Qwen3TTSSpeakerEncoderConfig.__init__


def _se_init(self, *args, **kwargs):
    for k in ("model_type", "transformers_version", "torch_dtype", "dtype"):
        kwargs.pop(k, None)
    return _orig_se_init(self, *args, **kwargs)


_qcfg.Qwen3TTSSpeakerEncoderConfig.__init__ = _se_init


# Qwen3TTSTalkerModel reads config.text_vocab_size but Qwen3TTSTalkerConfig
# never sets it (and the saved config.json omits it). Default to Qwen3's
# tokenizer vocab (151936 — matches the saved text_embedding.weight rows).
_TEXT_VOCAB_SIZE_DEFAULT = 151936
_orig_talker_init = _qcfg.Qwen3TTSTalkerConfig.__init__


def _talker_init(self, *args, **kwargs):
    text_vocab_size = kwargs.pop("text_vocab_size", _TEXT_VOCAB_SIZE_DEFAULT)
    _orig_talker_init(self, *args, **kwargs)
    self.text_vocab_size = text_vocab_size


_qcfg.Qwen3TTSTalkerConfig.__init__ = _talker_init


def _read_vocence_yaml(repo: Path) -> dict[str, Any]:
    path = repo / _VOCENCE_YAML
    if not path.is_file():
        return {}
    from yaml import safe_load
    with path.open("r", encoding="utf-8") as fh:
        data = safe_load(fh)
    return data if isinstance(data, Mapping) else {}


def _ensure_snapshot(repo: Path) -> Path:
    repo = repo.resolve()
    marker = repo / _CONFIG_NAME
    if not marker.is_file():
        raise FileNotFoundError(f"Model snapshot incomplete: {marker} missing.")
    return repo


def _read_max_new_tokens(repo: Path, fallback: int) -> int:
    """Match HF snapshot generation_config.json so decode budget isn't silently clipped."""
    path = repo / _GENERATION_CONFIG_NAME
    if not path.is_file():
        return fallback
    try:
        with path.open("r", encoding="utf-8") as fh:
            data = json.load(fh)
        raw = data.get("max_new_tokens", fallback)
        n = int(raw)
        return n if n > 0 else fallback
    except (OSError, ValueError, TypeError, json.JSONDecodeError):
        return fallback


# ---------------------------------------------------------------------------
# Quality scoring
# ---------------------------------------------------------------------------

_WORD_RE = re.compile(r"\w+", re.UNICODE)


def _word_error_rate(reference: str, hypothesis: str) -> float:
    """Levenshtein-on-tokens WER, identical to validator's evaluation.py."""
    ref = _WORD_RE.findall(reference.lower())
    hyp = _WORD_RE.findall(hypothesis.lower())
    if not ref:
        return 1.0 if hyp else 0.0
    n, m = len(ref), len(hyp)
    prev = list(range(m + 1))
    for i in range(1, n + 1):
        curr = [i] + [0] * m
        for j in range(1, m + 1):
            cost = 0 if ref[i - 1] == hyp[j - 1] else 1
            curr[j] = min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + cost)
        prev = curr
    return min(1.0, prev[m] / n)


def _is_valid(waveform: np.ndarray, sample_rate: int) -> bool:
    """Reject obviously broken samples before scoring."""
    n = waveform.shape[0]
    if n == 0:
        print("[error] validity: empty waveform", flush=True)
        return False
    duration = n / float(sample_rate)
    if duration < _MIN_DURATION_SEC or duration > _MAX_DURATION_SEC:
        print("[error] validity: duration out of bounds", flush=True)
        return False
    rms = float(np.sqrt(np.mean(np.square(waveform))))
    if not np.isfinite(rms) or rms < _MIN_RMS:
        print("[error] validity: RMS out of bounds", flush=True)
        return False
    peak = float(np.max(np.abs(waveform)))
    if peak >= _MAX_PEAK:
        print("[error] validity: peak out of bounds", flush=True)
        return False
    return True


def _utmosv2_cache_root() -> Path:
    """UTMOSv2 cache dir; allow override via env, fall back to ~/.cache/utmosv2."""
    import os
    return Path(os.environ.get("UTMOSV2_CACHE_DIR")
                or (Path.home() / ".cache" / "utmosv2"))


def _dump_utmosv2_cache_state(prefix: str = "  ") -> None:
    """Print what currently exists under the UTMOSv2 cache so we can see
    whether the package downloaded ANYTHING vs nothing. Useful diagnostic
    when 'Done.' is printed but checkpoint files are missing."""
    root = _utmosv2_cache_root()
    if not root.exists():
        print(f"{prefix}cache root does not exist: {root}", flush=True)
        return
    found = []
    for p in root.rglob("*"):
        if p.is_file():
            try:
                size = p.stat().st_size
            except OSError:
                size = -1
            rel = p.relative_to(root)
            found.append((str(rel), size))
    if not found:
        print(f"{prefix}cache root empty: {root}", flush=True)
        return
    print(f"{prefix}{root} contents:", flush=True)
    for rel, size in sorted(found):
        print(f"{prefix}  {rel}  ({size} bytes)", flush=True)


def _ensure_utmosv2_weights() -> bool:
    """Force-download UTMOSv2 fold checkpoints when the package's auto-download
    leaves them missing. Tries the most common HF repo IDs for sarulab-speech's
    UTMOSv2 release. Returns True if at least the fusion_stage3 fold0 file
    ends up on disk at the path UTMOSv2 looks for.
    """
    target_root = _utmosv2_cache_root()
    target_models = target_root / "models"
    target_models.mkdir(parents=True, exist_ok=True)
    sentinel = target_models / "fusion_stage3" / "fold0_s42_best_model.pth"
    if sentinel.is_file() and sentinel.stat().st_size > 1024:
        print(f"  UTMOSv2 weights already present at {sentinel}", flush=True)
        return True

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("  huggingface_hub not installed; cannot manually fetch", flush=True)
        return False

    candidates = (
        "sarulab-speech/UTMOSv2",
        "sarulab-speech/utmosv2",
        "Tetsuya-S/UTMOSv2",
    )
    for repo_id in candidates:
        try:
            path = snapshot_download(
                repo_id=repo_id,
                local_dir=str(target_root),
                allow_patterns=["models/**", "*.json", "*.yaml"],
            )
            if sentinel.is_file():
                print(f"  manual download from {repo_id} -> {path} OK", flush=True)
                return True
            else:
                print(f"  {repo_id} downloaded but sentinel still missing: {sentinel}",
                      flush=True)
        except Exception as e:
            print(f"  hf_hub_download from {repo_id} failed: "
                  f"{type(e).__name__}: {e}", flush=True)
            continue
    return False


class _CompositeScorer:
    """UTMOSv2 (naturalness proxy) + faster-whisper WER (script proxy).

    Final score = 0.5 * (utmos / 5.0) + 0.5 * (1 - WER), each in [0, 1].
    Falls back gracefully if either component is unavailable.
    """

    def __init__(self) -> None:
        # UTMOSv2 (sarulab-speech). Has known issue where create_model's
        # auto-download prints 'Done.' but only fetches the manifest, leaving
        # per-fold .pth files missing. Detect → manual hf_hub fallback → retry.
        self._utmos = None
        try:
            import torch
            import utmosv2
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

            try:
                self._utmos = utmosv2.create_model(pretrained=True, device=device)
            except FileNotFoundError as e:
                print(f"scorer: UTMOSv2 first load failed: {e}", flush=True)
                print("scorer: cache state at point of failure:", flush=True)
                _dump_utmosv2_cache_state()
                print("scorer: attempting manual hf_hub download...", flush=True)
                if _ensure_utmosv2_weights():
                    self._utmos = utmosv2.create_model(pretrained=True, device=device)
                else:
                    raise

            print(f"scorer: UTMOSv2 loaded on {device}", flush=True)
        except Exception as e:
            print(f"scorer: UTMOSv2 unavailable ({type(e).__name__}: {e}) "
                  f"- continuing without naturalness signal", flush=True)
            _dump_utmosv2_cache_state()

        # faster-whisper for script alignment
        self._whisper = None
        try:
            from faster_whisper import WhisperModel
            import torch
            device = "cuda" if torch.cuda.is_available() else "cpu"
            compute = "float16" if device == "cuda" else "int8"
            self._whisper = WhisperModel("base", device=device, compute_type=compute)
            print(f"scorer: faster-whisper-base on {device}", flush=True)
        except Exception as e:
            print(f"scorer: Whisper unavailable ({type(e).__name__}: {e}) "
                  f"- continuing without script signal", flush=True)

    def _utmos_score(self, wav: np.ndarray, sr: int) -> float:
        """UTMOSv2 README API: model.predict(data=ndarray, sr=int). Keyword-only."""
        if self._utmos is None:
            return 0.5
        try:
            arr = np.ascontiguousarray(wav, dtype=np.float32)
            mos = self._utmos.predict(data=arr, sr=int(sr))
            if hasattr(mos, "item"):
                mos = float(mos.item() if mos.ndim == 0 else mos.flatten()[0].item())
            elif isinstance(mos, np.ndarray):
                mos = float(mos.flatten()[0])
            else:
                mos = float(mos)
            return max(0.0, min(1.0, mos / 5.0))
        except Exception as e:
            print(f"UTMOSv2 predict error ({type(e).__name__}): {e}", flush=True)
            return 0.5

    def _whisper_wer(self, wav: np.ndarray, sr: int, target_text: str) -> float:
        """Returns 1 - WER ∈ [0, 1]; higher is better."""
        if self._whisper is None or not target_text.strip():
            return 0.5
        try:
            if sr != 16000:
                import torch
                import torchaudio.functional as AF
                t = torch.from_numpy(wav).unsqueeze(0)
                wav = AF.resample(t, sr, 16000).squeeze(0).numpy()
            segments, _ = self._whisper.transcribe(wav, language="en", beam_size=1)
            hyp = " ".join(seg.text for seg in segments).strip()
            return max(0.0, 1.0 - _word_error_rate(target_text, hyp))
        except Exception as e:
            print(f"Whisper transcribe error ({type(e).__name__}): {e}", flush=True)
            return 0.5

    def score(self, wav: np.ndarray, sr: int, target_text: str) -> float:
        u = self._utmos_score(wav, sr)
        w = self._whisper_wer(wav, sr, target_text)
        return 0.3 * u + 0.7 * w


# ---------------------------------------------------------------------------
# Instruction format conversion
# ---------------------------------------------------------------------------
#
# The validator sends the trait segment as a pipe-separated key:value string
# (built by vocence/pipeline/evaluation.py::format_task_prompt_for_tts):
#     "gender: X | pitch: Y | speed: Z | age_group: ... | emotion: ... | tone: ... | accent: ..."
# We rewrite that into a natural-language sentence the TTS model parses better:
#     "A {age} {gender} speaker with a {tone} tone speaks {speed_adv} and
#      {emotion_adv} at a {pitch} pitch, with a {accent} accent."
#
# Closed enums (mirrored from evaluation.py VOICE_TRAIT_ENUMS):
#   gender:    male, female, neutral
#   pitch:     low, mid, high
#   speed:     slow, normal, fast
#   age_group: child, young_adult, adult, senior
#   emotion:   neutral, happy, sad, angry, calm, excited, serious, fearful
#   tone:      warm, cold, friendly, formal, casual, authoritative
#   accent:    us, uk, au, in, neutral, other

_SPEED_ADVERBS = {
    "slow": "slowly",
    "normal": "at a normal pace",
    "fast": "quickly",
}

_EMOTION_ADVERBS = {
    "neutral": "in a neutral manner",
    "happy": "happily",
    "sad": "sadly",
    "angry": "angrily",
    "calm": "calmly",
    "excited": "excitedly",
    "serious": "seriously",
    "fearful": "fearfully",
}

_ACCENT_NAMES = {
    "us": "American",
    "uk": "British",
    "au": "Australian",
    "in": "Indian",
    "neutral": "neutral",
    "other": "neutral",
}

_AGE_PHRASES = {
    "child": "child",
    "young_adult": "young adult",
    "adult": "adult",
    "senior": "senior",
}


def _structured_to_natural(instruction: str) -> str:
    """Rewrite 'gender: X | pitch: Y | ...' as a natural-language sentence.

    Pass-through for any input that is not a key:value pipe-separated string
    (so warmup() and ad-hoc callers can still hand in plain prose)."""
    if "|" not in instruction or ":" not in instruction:
        return instruction

    parts: dict[str, str] = {}
    for chunk in instruction.split("|"):
        if ":" not in chunk:
            continue
        k, v = chunk.split(":", 1)
        parts[k.strip().lower()] = v.strip().lower()

    if not any(k in parts for k in ("gender", "pitch", "speed", "age_group",
                                    "emotion", "tone", "accent")):
        return instruction

    age = _AGE_PHRASES.get(parts.get("age_group", "adult"),
                           parts.get("age_group", "adult").replace("_", " "))
    gender = parts.get("gender", "neutral")
    tone = parts.get("tone", "casual")
    pitch = parts.get("pitch", "mid")
    speed_raw = parts.get("speed", "normal")
    emotion_raw = parts.get("emotion", "neutral")
    accent_raw = parts.get("accent", "neutral")

    speed_adv = _SPEED_ADVERBS.get(speed_raw, f"at a {speed_raw} pace")
    emotion_adv = _EMOTION_ADVERBS.get(emotion_raw, f"in a {emotion_raw} manner")
    accent = _ACCENT_NAMES.get(accent_raw, accent_raw)

    def _a(word: str) -> str:
        return "an" if word and word[0].lower() in "aeiou" else "a"

    if gender == "neutral":
        speaker = f"{_a(age).capitalize()} {age} speaker"
    else:
        speaker = f"{_a(age).capitalize()} {age} {gender} speaker"

    return (
        f"{speaker} with {_a(tone)} {tone} tone speaks {speed_adv} "
        f"and {emotion_adv} at {_a(pitch)} {pitch} pitch, "
        f"with {_a(accent)} {accent} accent."
    )


# ---------------------------------------------------------------------------
# Qwen3-TTS loader
# ---------------------------------------------------------------------------

def _resolve_device(prefer_cuda: bool) -> str:
    import torch
    if prefer_cuda and torch.cuda.is_available():
        return "cuda:0"
    return "cpu"


def _resolve_dtype(torch, prefer_bf16: bool):
    if prefer_bf16 and torch.cuda.is_available():
        return torch.bfloat16
    return torch.float32


def _probe_flash_attn() -> tuple[bool, str]:
    """Try to import flash_attn so we know whether it's actually available.
    Returns (available, version-or-error-message).
    """
    try:
        import flash_attn  # type: ignore
        ver = getattr(flash_attn, "__version__", "unknown")
        return True, ver
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def _force_dtype_on_config_tree(config: Any, dtype: Any) -> None:
    """Recursively set `.dtype` on a PretrainedConfig and every sub-config.

    Qwen3TTSConfig contains nested sub_configs (talker_config →
    code_predictor_config, speaker_encoder_config). transformers' from_pretrained
    propagates the `dtype=` kwarg only to the outer config; each inner sub-model
    then calls _flash_attn_2_can_dispatch and reads its own config.dtype, which
    is None — triggering 'Flash Attention 2 without specifying a torch dtype'
    even though the model is correctly loaded in bf16. Pre-populating dtype on
    every node of the config tree makes that check pass cleanly.
    """
    try:
        config.dtype = dtype
    except Exception:
        pass
    sub_specs = getattr(type(config), "sub_configs", None) or {}
    for attr in sub_specs:
        sub = getattr(config, attr, None)
        if sub is not None:
            _force_dtype_on_config_tree(sub, dtype)


def _load_qwen(checkpoint_dir: str, device_map: str, torch_dtype, use_flash2: bool):
    """Load Qwen3-TTS with explicit logging of attention impl + fallback reason.

    Critically: do NOT swallow flash-attn errors silently. If use_flash2=True
    but flash-attn is unavailable or rejected, print the exception so the
    miner.py user knows the deploy isn't using flash-attn (and can fix the
    chute_config.yml wheel pin). Falls back to sdpa as last resort so the
    chute still serves traffic.
    """
    from qwen_tts import Qwen3TTSModel

    if use_flash2:
        avail, info = _probe_flash_attn()
        if avail:
            print(f"[load] flash_attn package present (version={info})", flush=True)
        else:
            print(f"[load] flash_attn import FAILED ({info}); "
                  f"forcing attn_implementation=sdpa", flush=True)
            use_flash2 = False

    requested = "flash_attention_2" if use_flash2 else "sdpa"

    # Pre-build the config and force dtype onto every nested sub-config.
    # Without this, Qwen3TTS's talker / code_predictor / speaker_encoder
    # sub-models each emit a spurious 'FA2 without specifying a torch dtype'
    # warning even though the loaded model is bf16.
    prebuilt_config = None
    try:
        from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig
        prebuilt_config = Qwen3TTSConfig.from_pretrained(checkpoint_dir)
        _force_dtype_on_config_tree(prebuilt_config, torch_dtype)
        print("[load] pre-populated dtype on Qwen3TTSConfig sub-config tree", flush=True)
    except Exception as e:
        print(f"[load] could not pre-build config ({type(e).__name__}: {e}); "
              f"falling back to dtype kwarg only", flush=True)

    common: dict[str, Any] = dict(
        pretrained_model_name_or_path=checkpoint_dir,
        device_map=device_map,
        dtype=torch_dtype,
        attn_implementation=requested,
    )
    if prebuilt_config is not None:
        common["config"] = prebuilt_config
    print(f"[load] Qwen3TTSModel.from_pretrained(attn_implementation={requested!r})",
          flush=True)
    try:
        model = Qwen3TTSModel.from_pretrained(**common)
        print(f"[load] OK: attn_implementation={requested!r}", flush=True)
        # Probe the loaded model to confirm what it actually uses
        try:
            inner = getattr(model, "model", None)
            if inner is not None and hasattr(inner, "config"):
                actual = getattr(inner.config, "_attn_implementation", None) or \
                         getattr(inner.config, "attn_implementation", None)
                if actual:
                    print(f"[load] model reports config._attn_implementation={actual!r}",
                          flush=True)
        except Exception:
            pass
        return model
    except Exception as e:
        if requested == "sdpa":
            raise  # already on the safe path; propagate
        print(f"[load] FAILED with attn={requested!r}: {type(e).__name__}: {e}",
              flush=True)
        print(f"[load] retrying with attn_implementation='sdpa'", flush=True)
        common["attn_implementation"] = "sdpa"
        return Qwen3TTSModel.from_pretrained(**common)


def _to_mono_f32(segment: np.ndarray) -> np.ndarray:
    arr = np.asarray(segment, dtype=np.float32)
    if arr.ndim > 1:
        arr = arr.mean(axis=1)
    return arr


def _call_qwen_with_kwarg_dropping(tts: Any, kwargs: dict) -> tuple[Any, Any]:
    """qwen-tts versions vary in which kwargs they accept; drop optional ones if rejected."""
    try:
        return tts.generate_voice_design(**kwargs)
    except TypeError:
        for drop in (
            ("max_new_tokens",),
            ("max_new_tokens", "top_k"),
            ("max_new_tokens", "top_k", "repetition_penalty"),
            ("max_new_tokens", "top_k", "repetition_penalty", "top_p"),
        ):
            trimmed = {k: v for k, v in kwargs.items() if k not in drop}
            try:
                return tts.generate_voice_design(**trimmed)
            except TypeError:
                continue
        raise


# ---------------------------------------------------------------------------
# Miner
# ---------------------------------------------------------------------------

class Miner:
    """Sequential best-of-up-to-2 Qwen3-TTS engine with optional UTMOS+Whisper selection."""

    def __init__(self, path_hf_repo: Path) -> None:
        self._root = _ensure_snapshot(Path(path_hf_repo))
        cfg = _read_vocence_yaml(self._root)
        runtime = cfg.get("runtime") or {}
        generation = cfg.get("generation") or {}
        limits = cfg.get("limits") or {}

        self._language = str(
            limits.get("default_language") or runtime.get("default_language", "English")
        )
        self._cap_instruction = int(limits.get("max_instruction_chars", 600))
        self._cap_text = int(limits.get("max_text_chars", 2000))

        self._num_candidates = int(
            generation.get("num_candidates")
            or runtime.get("num_candidates")
            or _DEFAULT_NUM_CANDIDATES
        )
        if self._num_candidates < 1:
            self._num_candidates = 1
        # Strategy uses at most two sequential samples; larger YAML values are clamped.
        self._num_candidates = min(self._num_candidates, 2)

        prefer_cuda = str(runtime.get("device_preference", "cuda")).lower() == "cuda"
        want_bf16 = str(runtime.get("dtype", "bfloat16")).lower() == "bfloat16"
        flash = bool(runtime.get("use_flash_attention_2", False))

        import torch
        device_map = _resolve_device(prefer_cuda)
        torch_dtype = _resolve_dtype(torch, want_bf16)
        self._tts = _load_qwen(str(self._root), device_map, torch_dtype, flash)
        self._max_new_tokens = _read_max_new_tokens(self._root, _DEFAULT_MAX_NEW_TOKENS)
        print(
            f"Qwen3-TTS ready (sequential best-of-up-to-{self._num_candidates}, "
            f"fast-first threshold={_FIRST_GEN_FAST_THRESHOLD_SEC}s, "
            f"max_new_tokens={self._max_new_tokens}, "
            f"device={device_map}, dtype={torch_dtype}).",
            flush=True,
        )

        self._scorer = _CompositeScorer()

    def __repr__(self) -> str:
        return f"Miner(qwen3-tts, sequential_best_of_up_to_{self._num_candidates})"

    # -- Vocence contract -----------------------------------------------------

    def warmup(self) -> None:
        """One full pass — TTS synth + scorer — so the first /speak pays no
        cold-start cost on any of the three models (Qwen3-TTS, UTMOSv2,
        faster-whisper). Mirrors the real generate_wav path: same instruction
        rewrite, same `_generate_single`, then one `_scorer.score` call on the
        produced audio."""
        status: dict[str, object] = {"done": False, "error": None}

        def _once() -> None:
            try:
                warm_text = "This is a warmup utterance for the voice engine."
                instruction = _structured_to_natural(
                    "gender: female | age_group: adult | pitch: mid | "
                    "speed: normal | emotion: neutral | tone: casual | accent: us"
                )
                t0 = time.monotonic()
                wav, sr = self._generate_single(
                    instruction=instruction,
                    text=warm_text,
                    sampling=False,
                )
                t_tts = time.monotonic() - t0
                # Scorer warmup: one UTMOSv2 + one faster-whisper pass on the
                # generated audio so their lazy weight loading / CUDA graph
                # capture happens here, not on the first real /speak.
                t1 = time.monotonic()
                s = self._scorer.score(wav, sr, warm_text)
                t_score = time.monotonic() - t1
                print(f"[warmup] warm up completed successfully", flush=True)

                status["done"] = True
            except Exception as exc:
                status["error"] = str(exc)

        worker = threading.Thread(target=_once, daemon=True)
        worker.start()
        worker.join(timeout=240.0)
        if not status["done"]:
            raise RuntimeError(status["error"] or "warmup exceeded 240s")

    def generate_wav(self, instruction: str, text: str) -> tuple[np.ndarray, int]:
        """Sequential best-of up to two candidates.

        Time the first generation. If it takes ``>= _FIRST_GEN_FAST_THRESHOLD_SEC``
        seconds, return that first valid waveform immediately (no second sample,
        no scorer). If it finishes faster and ``num_candidates`` allows, generate
        one more sample, score both with _CompositeScorer (UTMOS + Whisper), and
        return the higher-scoring valid candidate.
        """
        instruction = _structured_to_natural(instruction)
        if self._cap_instruction > 0:
            instruction = instruction[: self._cap_instruction]
        if self._cap_text > 0:
            text = text[: self._cap_text]

        candidates: list[tuple[np.ndarray, int]] = []
        first_error: Optional[Exception] = None

        # ---- 1st candidate (timed) ------------------------------------------
        t0 = time.monotonic()
        first_wav: Optional[np.ndarray] = None
        first_sr: Optional[int] = None
        try:
            first_wav, first_sr = self._generate_single(instruction, text, sampling=True)
        except Exception as e:
            first_error = e
        first_elapsed = time.monotonic() - t0

        if first_wav is not None and first_sr is not None and _is_valid(first_wav, first_sr):
            candidates.append((first_wav, first_sr))

        want_second = (
            first_elapsed < _FIRST_GEN_FAST_THRESHOLD_SEC
            and self._num_candidates >= 2
        )

        # ---- first generation slow: return first valid only, no scoring -----
        if not want_second and candidates:
            return candidates[0]

        extra = 1 if want_second else 0

        # ---- additional candidates -------------------------------------------
        for i in range(1, 1 + extra):
            try:
                wav, sr = self._generate_single(instruction, text, sampling=True)
            except Exception as e:
                if first_error is None:
                    first_error = e
                print(f"[gen] sample {i} failed: {type(e).__name__}: {e}", flush=True)
                continue
            if _is_valid(wav, sr):
                candidates.append((wav, sr))
            else:
                print(f"[gen] sample {i} rejected by validity filter", flush=True)
        candidate_elapsed = time.monotonic() - t0 - first_elapsed
        print(f"[gen] generated {len(candidates)} candidates in {candidate_elapsed:.1f}s", flush=True)
        # ---- score + pick best ---------------------------------------------
        if candidates:
            if len(candidates) == 1:
                return candidates[0]
            scores = [self._scorer.score(wav, sr, text) for wav, sr in candidates]
            score_time = time.monotonic() - t0 - first_elapsed - candidate_elapsed
            print(f"[gen] scored {len(candidates)} candidates in {score_time:.1f}s", flush=True)
            best_idx = int(np.argmax(scores))
            print(
                f"[gen] best-of-{len(candidates)}/{self._num_candidates}: "
                f"picked={best_idx}",
                flush=True,
            )
            total_time = time.monotonic() - t0
            return candidates[best_idx]

        # ---- fallback: every candidate failed validity ---------------------
        print("[gen] all candidates invalid; falling back to single greedy call", flush=True)
        try:
            return self._generate_single(instruction, text, sampling=False)
        except Exception:
            if first_error is not None:
                raise first_error
            raise
 
    # -- internal -------------------------------------------------------------

    def _generate_single(
        self,
        instruction: str,
        text: str,
        sampling: bool = True,
    ) -> tuple[np.ndarray, int]:
        kwargs: dict[str, Any] = dict(
            text=text,
            instruct=instruction,
            language=self._language,
            max_new_tokens=self._max_new_tokens,
        )
        if sampling:
            kwargs.update(
                do_sample=_GEN_DO_SAMPLE,
                temperature=_GEN_TEMPERATURE,
                top_p=_GEN_TOP_P,
                top_k=_GEN_TOP_K,
                repetition_penalty=_GEN_REPETITION_PENALTY,
            )

        waves, sr = _call_qwen_with_kwarg_dropping(self._tts, kwargs)

        if isinstance(waves, (list, tuple)):
            if not waves:
                raise ValueError("TTS generation returned no audio")
            first = waves[0]
        else:
            first = waves
        if first is None:
            raise ValueError("TTS generation returned empty channel")

        return _to_mono_f32(first), int(sr)