"""Assertion and response validation helpers for tests."""

import io
import json
import math
import tempfile
import threading
import wave
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, TypeVar

if TYPE_CHECKING:
    from tests.helpers.runtime import DiffusionResponse

import av
import numpy as np
import soundfile as sf
from PIL import Image

from tests.helpers.media import (
    convert_audio_bytes_to_text,
    cosine_similarity_text,
    decode_b64_image,
    preprocess_text,
)

_ResponseT = TypeVar("_ResponseT")

_GENDER_PIPELINE = None
_GENDER_PIPELINE_LOCK = threading.Lock()

# Assertions that sample a random variable rather than testing whether serving works:
# whether TTS audio is intelligible enough to transcribe back, and whether a preset voice's
# timbre estimates as the expected gender. Both are decided by the model and the estimator,
# so a caller measuring a success rate has to tell them apart from a serving failure. The
# messages are constants because a caller matching the wording by hand would silently stop
# matching if it were reworded, reclassifying every quality failure as a hard failure.
AUDIO_MISMATCH_MESSAGE = "The audio content is not same as the text"
GENDER_MISMATCH_MESSAGE = "estimated gender is"
_QUALITY_FAILURE_MESSAGES = (AUDIO_MISMATCH_MESSAGE, GENDER_MISMATCH_MESSAGE)


@dataclass(frozen=True)
class SuccessRateGate:
    """Gate a sampled quality check on how many requests pass, not on each one passing.

    Built by the send helpers from their own arguments, so a test states a policy and never
    touches this class. See ``collect_at_success_rate``.

    Attributes:
        min_successes: Fail below this many successes out of ``request_num``. Predeclare it
            from a measured baseline; tuning it until CI passes defeats the gate.
        concurrency: In-flight requests, in batches until the sample count is reached.
            Defaults to sending all of them at once.
    """

    min_successes: int
    concurrency: int | None = None


def is_quality_failure(exc: BaseException) -> bool:
    """True when ``exc`` is one of the sampled quality assertions above.

    A caller absorbing a few of these into a success-rate budget must not absorb an HTTP
    failure, missing audio or a text-keyword miss as well, or a broken server passes.
    """
    if not isinstance(exc, AssertionError):
        return False
    return any(message in str(exc) for message in _QUALITY_FAILURE_MESSAGES)


def collect_at_success_rate(
    sample: Callable[[], _ResponseT],
    gate: SuccessRateGate,
    *,
    request_num: int,
    report: Callable[[str], None] = print,
) -> list[_ResponseT]:
    """Send ``request_num`` samples and gate on how many passed, returning the successes.

    ``sample`` sends one request and asserts it. Only the sampled quality assertions are
    counted; anything else propagates immediately. Does not retry and does not exit early,
    since both would bias the rate being measured.
    """
    if gate.min_successes > request_num:
        raise ValueError(f"min_successes={gate.min_successes} exceeds request_num={request_num}, gate can never pass")
    concurrency = min(gate.concurrency or request_num, request_num)
    responses: list[_ResponseT] = []
    failures: list[str] = []
    lock = threading.Lock()

    def _one_sample(_: int) -> None:
        try:
            response = sample()
        except AssertionError as exc:
            if not is_quality_failure(exc):
                raise
            with lock:
                failures.append(str(exc))
            return
        with lock:
            responses.append(response)

    for start in range(0, request_num, concurrency):
        size = min(concurrency, request_num - start)
        with ThreadPoolExecutor(max_workers=size) as executor:
            for _ in executor.map(_one_sample, range(size)):
                pass

    successes = request_num - len(failures)
    report(f"success rate: {successes}/{request_num} succeeded")
    assert successes >= gate.min_successes, (
        f"success rate {successes}/{request_num} is below the "
        f"{gate.min_successes}/{request_num} gate "
        f"(95% CI {wilson_interval(successes, request_num)}). Failures:\n"
        + "\n".join(f"  - {failure}" for failure in failures)
    )
    return responses


def wilson_interval(successes: int, total: int, z: float = 1.96) -> str:
    """Format a 95% Wilson score interval, so a rate failure shows how marginal it is."""
    p = successes / total
    denom = 1 + z * z / total
    centre = (p + z * z / (2 * total)) / denom
    half = z / denom * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total))
    return f"{max(0.0, centre - half):.1%}-{min(1.0, centre + half):.1%}"


# Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
# short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
# leading SFX token), which flakes the deterministic similarity gate. Short
# clips first get a conservative containment fallback for minor ASR repeats/noise.
# A test can also opt in to ASR escalation by setting
# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``)
# in its request_config: on a failed fast pass the clip is re-transcribed with
# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while
# a genuine model artifact still fails (the strong ASR mismatches too).
_PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
_MIN_PCM_SPEECH_HNR_DB = 1.0
_PRESET_VOICE_GENDER_MAP: dict[str, str] = {
    "serena": "female",
    "uncle_fu": "male",
    "chelsie": "female",
    "clone": "female",
    "ethan": "male",
}


def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
    """Allow minor ASR repeats/noise for very short speech clips."""
    transcript_clean = preprocess_text(transcript)
    expected_clean = preprocess_text(expected)
    if not transcript_clean or not expected_clean:
        return False

    transcript_words = transcript_clean.split()
    expected_words = expected_clean.split()
    if not transcript_words or not expected_words:
        return False

    short_text = min(len(transcript_clean), len(expected_clean)) <= 15
    small_word_delta = len(transcript_words) <= len(expected_words) + 2
    return short_text and small_word_delta and expected_clean in transcript_clean


def assert_image_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
    run_level: str = None,
) -> None:
    """
    Validate image diffusion response.

    Expected request_config schema:
        {
            "request_type": "image",
            "extra_body": {
                "num_outputs_per_prompt": 1,
                "width": ...,
                "height": ...,
                ...
            }
        }
    """
    assert response.images is not None, "Image response is None"
    assert len(response.images) > 0, "No images in response"

    extra_body = request_config.get("extra_body") or {}

    num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt")
    if num_outputs_per_prompt is not None:
        assert len(response.images) == num_outputs_per_prompt, (
            f"Expected {num_outputs_per_prompt} images, got {len(response.images)}"
        )

    if run_level in {"advanced_model", "full_model"}:
        width = extra_body.get("width")
        height = extra_body.get("height")

        if width is not None or height is not None:
            if isinstance(width, (list, tuple)) and isinstance(height, (list, tuple)):
                assert len(response.images) == len(width) == len(height), (
                    f"Per-output size lists require one image per entry; got {len(response.images)} images, "
                    f"len(width)={len(width)}, len(height)={len(height)}"
                )
                for img, w, h in zip(response.images, width, height, strict=True):
                    assert_image_valid(img, width=int(w), height=int(h))
            else:
                for img in response.images:
                    assert_image_valid(
                        img,
                        width=_maybe_int(width) if width is not None else None,
                        height=_maybe_int(height) if height is not None else None,
                    )


def assert_images_generations_response(
    response: dict[str, Any],
    request_config: dict[str, Any],
    run_level: str | None = None,
) -> None:
    """Validate a successful ``/v1/images/generations`` JSON response."""
    del run_level
    request_body = request_config.get("json") or {}
    data = response.get("data")
    assert isinstance(data, list), "Image generation response is missing data[]"

    expected_count = int(request_body.get("n", 1))
    assert len(data) == expected_count, f"Expected {expected_count} images, got {len(data)}"

    width = height = None
    size = request_body.get("size")
    if isinstance(size, str) and "x" in size:
        width_value, height_value = size.lower().split("x", 1)
        width, height = int(width_value), int(height_value)

    for item in data:
        assert isinstance(item, dict), "Image generation data entries must be objects"
        b64_json = item.get("b64_json")
        assert isinstance(b64_json, str) and b64_json, "Image generation response is missing b64_json"
        image = decode_b64_image(b64_json)
        assert_image_valid(image, width=width, height=height)


def assert_video_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
    run_level: str = None,
) -> None:
    """
    Validate video diffusion response.

    Expected request_config schema:
        {
            "request_type": "video",
            "form_data": {
                "prompt": "...",
                "num_frames": ...,
                "width": ...,
                "height": ...,
                "fps": ...,
                ...
            }
        }
    """
    form_data = request_config.get("form_data", {})

    assert response.videos is not None, "Video response is None"
    assert len(response.videos) > 0, "No videos in response"

    expected_frames = _maybe_int(form_data.get("num_frames"))
    expected_width = _maybe_int(form_data.get("width"))
    expected_height = _maybe_int(form_data.get("height"))
    expected_fps = _maybe_int(form_data.get("fps"))

    # Skip num_frames assertion for Helios models because they round up frames
    model = request_config.get("model", "")
    if "Helios" in model:
        expected_frames = None

    for vid_bytes in response.videos:
        assert_video_valid(
            vid_bytes,
            num_frames=expected_frames,
            width=expected_width,
            height=expected_height,
            fps=expected_fps,
        )


def assert_video_first_frame_matches(
    video: Path | bytes | BytesIO,
    expected: np.ndarray,
    *,
    max_mean_absolute_error: float,
) -> None:
    """Assert that an encoded video's first frame matches a reference image."""
    if isinstance(video, Path):
        source: str | BytesIO = str(video)
    else:
        video_bytes = video if isinstance(video, bytes) else video.getvalue()
        source = BytesIO(video_bytes)

    with av.open(source) as container:
        first_frame = next(container.decode(video=0)).to_ndarray(format="rgb24")

    assert first_frame.shape == expected.shape
    mean_absolute_error = float(np.abs(first_frame.astype(np.float32) - expected).mean() / 255.0)
    assert mean_absolute_error < max_mean_absolute_error, (
        f"Expected first-frame MAE < {max_mean_absolute_error}, got {mean_absolute_error:.6f}."
    )


def assert_audio_diffusion_response(
    response: "DiffusionResponse",
    request_config: dict[str, Any],
    run_level: str = None,
) -> None:
    """
    Validate audio diffusion response.

    `response.audios` carries one entry per choice, each a `dict` with raw WAV
    bytes (`wav_bytes`) and the OpenAI audio metadata (`id`, `expires_at`).
    """
    assert response.audios, "Audio response is empty"
    for audio in response.audios:
        wav_bytes = audio.get("wav_bytes")
        assert wav_bytes, "Audio entry missing decoded WAV bytes"
        with wave.open(io.BytesIO(wav_bytes), "rb") as wav_file:
            assert wav_file.getnframes() > 0, "Decoded WAV has zero frames"
            assert wav_file.getframerate() > 0, "Decoded WAV has invalid sample rate"


def _maybe_int(value: Any) -> int | None:
    if value is None:
        return None
    return int(value)


def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None):
    """Assert the file is a loadable image with optional exact dimensions."""
    if isinstance(image, Path):
        assert image.exists(), f"Image not found: {image}"
        image = Image.open(image)
        image.load()
    assert image.width > 0 and image.height > 0
    if width is not None:
        assert image.width == width, f"Expected width={width}, got {image.width}"
    if height is not None:
        assert image.height == height, f"Expected height={height}, got {image.height}"
    return image


def assert_video_valid(
    video: Path | bytes | BytesIO,
    *,
    num_frames: int | None = None,
    width: int | None = None,
    height: int | None = None,
    fps: float | None = None,
) -> dict[str, int | float]:
    """Assert the MP4 has the expected resolution and frame count.

    For several diffusion backends, encoded MP4 frame count follows a codec-aligned
    convention (e.g. request `num_frames=8` can produce 9 encoded frames). Keep
    this compatibility behavior to avoid false negatives in online-serving tests.
    """
    temp_path = None
    cap = None
    try:
        import cv2

        if isinstance(video, Path):
            if not video.exists():
                raise AssertionError(f"Video file not found: {video}")
            video_path = str(video)
        else:
            suffix = ".mp4"
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, mode="wb") as tmp:
                if isinstance(video, bytes):
                    tmp.write(video)
                elif isinstance(video, BytesIO):
                    tmp.write(video.getvalue())
                else:
                    raise TypeError(f"Unsupported video type: {type(video)}")
                temp_path = Path(tmp.name)
                video_path = str(temp_path)

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise AssertionError("Failed to open video capture")

        actual_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        actual_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        actual_fps = float(cap.get(cv2.CAP_PROP_FPS))
        actual_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        if width is not None:
            assert actual_width == width, f"Expected width={width}, got {actual_width}"
        if height is not None:
            assert actual_height == height, f"Expected height={height}, got {actual_height}"
        if fps is not None and actual_fps:
            assert abs(actual_fps - float(fps)) < 1.0, f"Expected fps~={fps}, got {actual_fps}"
        if num_frames is not None:
            expected_frames = (int(num_frames) // 4) * 4 + 1
            assert actual_frames == expected_frames, f"Expected frames={expected_frames}, got {actual_frames}"

        return {
            "width": actual_width,
            "height": actual_height,
            "fps": actual_fps,
            "num_frames": actual_frames,
        }
    except Exception as e:
        print(f"ERROR: {type(e).__name__}: {e}", flush=True)
        raise
    finally:
        if cap is not None:
            cap.release()
        if temp_path and temp_path.exists():
            try:
                temp_path.unlink()
            except OSError:
                pass


def assert_audio_valid(
    audio_or_path: Path | np.ndarray,
    *,
    sample_rate: int,
    channels: int,
    duration_s: float,
) -> None:
    """Assert WAV file or (batch, channels, samples) ndarray matches expected audio format."""
    expected_samples = int(duration_s * sample_rate)
    if isinstance(audio_or_path, np.ndarray):
        audio = audio_or_path
        assert audio.ndim == 3, f"Expected audio ndim=3 (batch, channels, samples), got shape {audio.shape}"
        assert audio.shape[0] == 1, f"Expected batch size 1, got {audio.shape[0]}"
        assert audio.shape[1] == channels, f"Expected {channels} channels, got {audio.shape[1]}"
        assert audio.shape[2] == expected_samples, (
            f"Expected {expected_samples} samples ({duration_s}s @ {sample_rate} Hz), got {audio.shape[2]}"
        )
        return

    path = audio_or_path
    assert path.exists(), f"Audio not found: {path}"
    info = sf.info(str(path))
    assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}"
    assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}"
    assert info.frames == expected_samples, (
        f"Expected {expected_samples} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}"
    )


def _load_gender_pipeline():
    global _GENDER_PIPELINE
    if _GENDER_PIPELINE is not None:
        return _GENDER_PIPELINE
    model_name = "7wolf/wav2vec2-base-gender-classification"
    try:
        from transformers import pipeline

        _GENDER_PIPELINE = pipeline(task="audio-classification", model=model_name, device=-1)
        return _GENDER_PIPELINE
    except Exception as exc:  # pragma: no cover
        print(f"Warning: failed to create gender pipeline '{model_name}': {exc}")
        _GENDER_PIPELINE = None
        return None


def _median_pitch_hz_from_autocorr(mono: np.ndarray, sr: int) -> float | None:
    x = np.asarray(mono, dtype=np.float64)
    x = x - np.mean(x)
    if x.size < int(0.15 * sr):
        return None
    frame_len = int(0.04 * sr)
    hop = max(frame_len // 2, 1)
    f0_min_hz, f0_max_hz = 70.0, 400.0
    lag_min = max(1, int(sr / f0_max_hz))
    lag_max = min(frame_len - 2, int(sr / f0_min_hz))
    if lag_max <= lag_min:
        return None
    win = np.hamming(frame_len)
    pitches: list[float] = []
    for start in range(0, int(x.shape[0]) - frame_len, hop):
        frame = x[start : start + frame_len] * win
        frame = frame - np.mean(frame)
        if float(np.sqrt(np.mean(frame**2))) < 1e-4:
            continue
        ac = np.correlate(frame, frame, mode="full")[frame_len - 1 :]
        ac = ac / (float(ac[0]) + 1e-12)
        region = ac[lag_min : lag_max + 1]
        peak_rel = int(np.argmax(region))
        peak_lag = peak_rel + lag_min
        if peak_lag <= 0:
            continue
        f0 = float(sr) / float(peak_lag)
        if f0_min_hz <= f0 <= f0_max_hz:
            pitches.append(f0)
    if len(pitches) < 4:
        return None
    return float(np.median(np.asarray(pitches, dtype=np.float64)))


def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
    data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True)
    if data.size == 0:
        raise ValueError("Empty audio")
    mono = np.mean(data, axis=1)
    try:
        target_sr = 16000
        if int(sr) != target_sr and mono.size > 1:
            src_len = int(mono.shape[0])
            dst_len = max(1, int(round(src_len * float(target_sr) / float(sr))))
            src_idx = np.arange(src_len, dtype=np.float32)
            dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
            mono = np.interp(dst_idx, src_idx, mono.astype(np.float32, copy=False)).astype(np.float32)
            sr = target_sr

        median_f0 = _median_pitch_hz_from_autocorr(mono, sr)
        clf = _load_gender_pipeline()
        if clf is None:
            print("gender model not available, returning 'unknown'")
            return "unknown"
        with _GENDER_PIPELINE_LOCK:
            outputs = clf(mono, sampling_rate=sr)
        if not outputs:
            return "unknown"
        top = outputs[0]
        label = str(top.get("label", "")).lower()
        conf = float(top.get("score", 0.0))
        if conf < 0.6:
            gender = "unknown"
        elif ("female" in label) or ("жен" in label):
            gender = "female"
        elif ("male" in label) or ("муж" in label):
            gender = "male"
        else:
            gender = "unknown"

        if gender == "female" and median_f0 is not None and median_f0 < 165.0 and conf < 0.88:
            print(f"gender pitch assist: reclassifying female->male (median_f0={median_f0:.1f} Hz, conf={conf:.3f})")
            gender = "male"
        elif gender == "male" and median_f0 is not None and median_f0 > 230.0 and conf < 0.88:
            print(f"gender pitch assist: reclassifying male->female (median_f0={median_f0:.1f} Hz, conf={conf:.3f})")
            gender = "female"
        print(
            f"gender classifier: label={label}, conf={conf:.3f}, gender={gender}"
            + (f", median_f0={median_f0:.1f}Hz" if median_f0 is not None else "")
        )
        return gender
    except Exception as exc:  # pragma: no cover
        print(f"Warning: gender classification failed, returning 'unknown': {exc}")
        return "unknown"


def _assert_preset_voice_gender_from_audio(
    audio_bytes: bytes | None,
    voice_name: str | None,
    *,
    response_format: str | None = None,
) -> None:
    """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown)."""
    if response_format == "pcm":
        return
    if not voice_name or not audio_bytes:
        return
    key = str(voice_name).lower()
    expected_gender = _PRESET_VOICE_GENDER_MAP.get(key)
    if expected_gender is None:
        return
    estimated_gender = _estimate_voice_gender_from_audio(audio_bytes)
    print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
    if estimated_gender != "unknown":
        assert estimated_gender == expected_gender, (
            f"{voice_name!r} is expected {expected_gender}, but {GENDER_MISMATCH_MESSAGE} {estimated_gender!r}"
        )


def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ) -> float:
    frame_len = int(0.03 * sr)
    hop = frame_len // 2
    hnr_values: list[float] = []
    for start in range(0, len(pcm_samples) - frame_len, hop):
        frame = pcm_samples[start : start + frame_len].astype(np.float32, copy=False)
        frame = frame - np.mean(frame)
        if np.max(np.abs(frame)) < 0.01:
            continue
        ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :]
        ac = ac / (ac[0] + 1e-10)
        min_lag = int(sr / 400)
        max_lag = min(int(sr / 80), len(ac))
        if min_lag >= max_lag:
            continue
        peak = float(np.max(ac[min_lag:max_lag]))
        if 0 < peak < 1:
            hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10)))
    return float(np.mean(hnr_values)) if hnr_values else 0.0


def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB) -> None:
    """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech.

    min_hnr_db defaults to the global _MIN_PCM_SPEECH_HNR_DB (1.0 dB),
    which matches the cleaner TTS models the helper was originally calibrated
    for. Quieter codecs (e.g. MOSS-TTS-Nano, whose voice_clone output is
    intrinsically around -2 dB) can pass a lower per-test threshold via
    request_config["min_hnr_db"] to keep the catastrophic-failure check
    while not gating CI on a model-intrinsic property.
    """
    assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes"
    assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16"
    pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
    hnr = _compute_pcm_hnr_db(pcm_samples)
    print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB)")
    assert hnr >= min_hnr_db, (
        f"Audio distortion detected: HNR={hnr:.2f} dB < {min_hnr_db} dB. "
        "Voice clone decoder may be losing ref_code speaker context on later chunks."
    )


def _response_has_audio_output(response: Any) -> bool:
    if response.audio_bytes:
        return len(response.audio_bytes) > 0
    if isinstance(getattr(response, "audio_content", None), str) and response.audio_content.strip():
        return True
    audio_data = getattr(response, "audio_data", None)
    return bool(audio_data)


def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool:
    if run_level not in {"advanced_model", "full_model"}:
        return False
    modalities = request_config.get("modalities", ["text", "audio"])
    if "audio" not in modalities:
        return False
    keywords_dict = request_config.get("key_words", {}) or {}
    # When text is not an output modality, the keyword loop validates keywords
    # against the audio transcript -- for keywords under ANY word_type
    # (text/image/audio/video), not just "audio". Mirror that here so the
    # transcript is actually computed; otherwise the loop hits
    # `assert transcript is not None` with transcript=None (e.g. an audio-only
    # request carrying key_words={"text": [...]}).
    if "text" not in modalities and any(
        keywords_dict.get(word_type) for word_type in ("text", "image", "audio", "video")
    ):
        return True
    if request_config.get("audio_ref_text"):
        return True
    return "text" in modalities


def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool:
    if run_level not in {"advanced_model", "full_model"}:
        return False
    if request_config.get("response_format") == "pcm":
        return False
    return bool(request_config.get("input"))


def _resolve_audio_transcript(
    response: Any,
    request_config: dict[str, Any],
    run_level: str,
    *,
    speech_api: bool,
) -> str | None:
    """Run Whisper only when this run_level / request_config needs a transcript for assertions."""
    needs = (
        _speech_assertion_needs_audio_transcript(request_config, run_level)
        if speech_api
        else _omni_assertion_needs_audio_transcript(request_config, run_level)
    )
    if not needs:
        return None
    existing = getattr(response, "audio_content", None)
    if isinstance(existing, str) and existing.strip():
        return existing
    audio_bytes = getattr(response, "audio_bytes", None)
    if not audio_bytes:
        return None
    return convert_audio_bytes_to_text(audio_bytes, language=request_config.get("transcript_language"))


def assert_omni_response(response: Any, request_config: dict[str, Any], run_level):
    """
    Validate response results.

    Args:
        response: OmniResponse object

    Raises:
        AssertionError: When the response does not meet validation criteria
    """
    assert response.success, "The request failed."

    modalities = request_config.get("modalities", ["text", "audio"])

    if run_level in {"advanced_model", "full_model"}:
        transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=False)
        # Verify output success
        if "audio" in modalities:
            assert _response_has_audio_output(response), "No audio output is generated"
            if transcript is not None:
                print(f"audio content is: {transcript}")
            speaker = request_config.get("speaker")
            if speaker:
                _assert_preset_voice_gender_from_audio(
                    response.audio_bytes,
                    speaker,
                    response_format=request_config.get("response_format"),
                )
        if "text" in modalities:
            assert response.text_content is not None, "No text output is generated"
            print(f"text content is: {response.text_content}")

        # Verify keywords in output
        word_types = ["text", "image", "audio", "video"]
        keywords_dict = request_config.get("key_words", {})
        for word_type in word_types:
            keywords = keywords_dict.get(word_type)
            if "text" in modalities:
                if keywords:
                    text_lower = response.text_content.lower()
                    assert any(str(kw).lower() in text_lower for kw in keywords), (
                        "The output does not contain any of the keywords."
                    )
            else:
                if keywords:
                    assert transcript is not None, "No audio transcript for keyword validation"
                    audio_lower = transcript.lower()
                    assert any(str(kw).lower() in audio_lower for kw in keywords), (
                        "The output does not contain any of the keywords."
                    )

        # Verify similarity (Whisper transcript vs streamed/detokenized text)
        if "audio" in modalities:
            audio_ref_text = request_config.get("audio_ref_text")
            similarity_threshold = request_config.get("similarity_threshold", 0.8)
            if "text" in modalities:
                assert transcript is not None, "No audio transcript for similarity validation"
                text_output = (response.text_content or "").strip()
                # For very short outputs (e.g. one-word answers), n-gram cosine
                # similarity with length penalty is unreliable because Whisper
                # may hallucinate extra context around the short utterance.  Use
                # a containment check instead: the shorter text must appear in
                # the longer one (after preprocessing removes punctuation).
                _SHORT_TEXT_THRESHOLD = 15
                if len(text_output) <= _SHORT_TEXT_THRESHOLD or len(transcript) <= _SHORT_TEXT_THRESHOLD:
                    shorter = text_output.lower() if len(text_output) <= len(transcript) else transcript.lower()
                    longer = transcript.lower() if len(text_output) <= len(transcript) else text_output.lower()
                    import re as _re

                    shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
                    longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
                    assert shorter_clean and (shorter_clean in longer_clean), (
                        f"{AUDIO_MISMATCH_MESSAGE} "
                        f"(short-text containment check failed: "
                        f"text={text_output!r}, transcript={transcript!r})"
                    )
                    print(f"short-text containment check passed: {shorter_clean!r} in {longer_clean!r}")
                else:
                    similarity = cosine_similarity_text(
                        transcript.lower(),
                        text_output.lower(),
                    )
                    print(f"similarity is: {similarity}")
                    assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE
            if audio_ref_text:
                assert transcript is not None, "No audio transcript for reference-text validation"
                audio_similarity = cosine_similarity_text(
                    transcript.strip().lower(),
                    str(audio_ref_text).lower(),
                )
                assert audio_similarity > similarity_threshold, (
                    f"The audio content does not match reference text: similarity={audio_similarity:.3f}"
                )


def _assert_transcript_matches(
    transcript: str,
    audio_bytes: bytes | None,
    expected_text: Any,
    *,
    threshold: float,
    escalation_model: str | None = None,
    language: str | None = None,
) -> None:
    """Assert spoken audio matches ``expected_text``.

    ``transcript`` is the fast whisper-``small`` result. If it clears
    ``threshold`` the check passes immediately.

    If the cosine check fails, very short clips get a conservative containment
    fallback that accepts minor ASR repeats/noise only when the expected text is
    still present and the transcript has few extra words.

    When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
    request_config key) and the fast check plus containment fallback fail, the
    clip is re-transcribed with that stronger ASR and the assertion is decided on
    its verdict -- so a weak whisper-``small`` mishear on a short clip does not
    flake the gate, while a genuine model artifact still fails (the strong ASR
    mismatches too). ``language`` applies to that escalated pass too, so a request
    that pins a language does not silently fall back to auto-detection on retry.
    """
    expected = str(expected_text).strip().lower()
    similarity = cosine_similarity_text(transcript.strip().lower(), expected)
    print(f"Cosine similarity: {similarity:.3f}")
    if similarity > threshold:
        return

    if _short_transcript_contains_expected(transcript, expected):
        print("short speech containment check passed")
        return

    if escalation_model and audio_bytes:
        print(
            f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "
            f"escalating to whisper-{escalation_model} to rule out an ASR mishear"
        )
        strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model, language=language)
        strong_similarity = cosine_similarity_text(strong_transcript.strip().lower(), expected)
        print(
            f"audio content (whisper-{escalation_model}): {strong_transcript}\n"
            f"Cosine similarity (whisper-{escalation_model}): {strong_similarity:.3f}"
        )
        assert strong_similarity > threshold, (
            f"Transcript doesn't match input after ASR escalation: "
            f"input={expected_text!r}; whisper-small='{transcript}' (sim={similarity:.2f}); "
            f"whisper-{escalation_model}='{strong_transcript}' (sim={strong_similarity:.2f})"
        )
        return

    assert similarity > threshold, (
        f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
    )


def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
    """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.

    When ``request_config`` carries ``status_code`` and/or ``err_message``, the
    request is expected to be rejected: assert it failed and that the HTTP status
    / error text match. Otherwise the normal success-path checks run.
    """
    expected_status = request_config.get("status_code")
    expected_err = request_config.get("err_message")
    if expected_status is not None or expected_err is not None:
        assert not response.success, "Expected an error response, but the request succeeded."
        if expected_status is not None:
            allowed = expected_status if isinstance(expected_status, (list, tuple)) else (expected_status,)
            assert response.status_code in allowed, f"Expected HTTP status in {allowed}, got {response.status_code}"
        if expected_err is not None:
            alternatives = expected_err if isinstance(expected_err, (list, tuple)) else (expected_err,)
            error_text = response.error_message or ""
            assert any(alt in error_text for alt in alternatives), (
                f"Expected one of {alternatives} in error text, got: {error_text!r}"
            )
        return

    assert response.success, "The request failed."

    # Optional floor on decoded audio size (models with very short clips may use a lower value).
    min_audio = request_config.get("min_audio_bytes")
    if min_audio is not None:
        n = int(min_audio)
        if n > 0:
            ab = response.audio_bytes
            assert ab is not None, "Expected audio bytes when min_audio_bytes is set"
            assert len(ab) > n, f"Audio payload too small: {len(ab)} bytes, expected more than {n} (min_audio_bytes)"

    req_fmt = request_config.get("response_format")
    if req_fmt == "pcm" and response.audio_bytes:
        if response.audio_format:
            assert "pcm" in response.audio_format.lower(), (
                f"Expected audio/pcm content-type, got {response.audio_format!r}"
            )
    elif req_fmt == "wav" and response.audio_format:
        assert req_fmt in response.audio_format

    if run_level in {"advanced_model", "full_model"}:
        if req_fmt == "pcm" and response.audio_bytes:
            min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB))
            _assert_pcm_int16_speech_hnr(response.audio_bytes, min_hnr_db=min_hnr_db)

        transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True)
        if transcript is not None:
            expected_text = request_config.get("input")
            if expected_text:
                print(f"audio content is: {transcript}")
                print(f"input text is: {expected_text}")
                _assert_transcript_matches(
                    transcript,
                    getattr(response, "audio_bytes", None),
                    expected_text,
                    threshold=0.9,
                    escalation_model=request_config.get("transcript_escalation_model"),
                    language=request_config.get("transcript_language"),
                )
        _assert_preset_voice_gender_from_audio(
            response.audio_bytes,
            request_config.get("voice"),
            response_format=request_config.get("response_format"),
        )


def assert_diffusion_response(response: "DiffusionResponse", request_config: dict[str, Any], run_level: str = None):
    assert response.success, "The request failed."
    has_any_content = any(content is not None for content in (response.images, response.videos, response.audios))
    assert has_any_content, "Response contains no images, videos, or audios"
    if response.images is not None:
        assert_image_diffusion_response(response=response, request_config=request_config, run_level=run_level)
    if response.videos is not None:
        assert_video_diffusion_response(response=response, request_config=request_config, run_level=run_level)
    if response.audios is not None:
        assert_audio_diffusion_response(response=response, request_config=request_config, run_level=run_level)


def _http_response_body_materialize(resp: Any) -> tuple[bytes, dict[str, Any] | None]:
    """Serialize ``HttpResponse``-like body to UTF-8 bytes and parse a JSON object when possible."""
    jb = getattr(resp, "json_body", None)
    if jb is not None:
        raw = json.dumps(jb, ensure_ascii=False).encode("utf-8")
        if isinstance(jb, dict):
            return raw, jb
        try:
            parsed = json.loads(raw.decode("utf-8", errors="replace"))
        except json.JSONDecodeError:
            return raw, None
        return raw, parsed if isinstance(parsed, dict) else None
    err = getattr(resp, "error_message", None)
    raw = (err or "").encode("utf-8", errors="replace")
    try:
        parsed = json.loads(raw.decode("utf-8", errors="replace"))
    except json.JSONDecodeError:
        return raw, None
    return raw, parsed if isinstance(parsed, dict) else None


def assert_err_message_in_text(
    haystack: str,
    err_message: str | tuple[str, ...] | list[str] | set[str] | frozenset[str],
    *,
    sequence_match: Literal["all", "any"] = "all",
) -> None:
    """Assert ``err_message`` appears in ``haystack`` (case-insensitive).

    For non-string sequences: ``sequence_match='all'`` requires every substring (HTTP-style);
    ``sequence_match='any'`` requires at least one (WebSocket first-frame JSON helpers).
    """
    hl = haystack.lower()
    if isinstance(err_message, (list, tuple, set, frozenset)):
        if sequence_match == "all":
            missing = [s for s in err_message if str(s).lower() not in hl]
            assert not missing, (
                f"Expected error text to contain all of {err_message!r}; missing {missing!r}. haystack={haystack!r}"
            )
        else:
            assert any(str(s).lower() in hl for s in err_message), (
                f"Expected error text to contain one of {err_message!r}. haystack={haystack!r}"
            )
    else:
        assert str(err_message).lower() in hl, f"Expected error text to contain {str(err_message)!r}, got: {haystack!r}"


def assert_http_error(
    resp: Any,
    *,
    err_code: int | tuple[int, ...] | list[int] | None = None,
    err_message: str | tuple[str, ...] | list[str] | None = None,
    websocket_json_message: bool = False,
) -> dict[str, Any] | None:
    """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object.

    Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when
    ``request_config`` contains optional ``err_code`` and/or ``err_message``.

    When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame).
    Tuple/list ``err_message`` then uses **any** substring match; HTTP mode still requires **all** pieces.

    - ``err_code``: exact HTTP ``int``, or membership if a non-string sequence (e.g. ``(400, 422)``).
      When ``err_code`` is set and the actual status is a client error in ``400..499`` other than ``404``,
      the JSON body must include FastAPI ``detail`` and/or OpenAI-style ``error`` (2xx skips this).
    - ``err_message``: substring match (case-insensitive) against serialized ``json_body`` and ``error_message``.
      If ``err_message`` is a non-string sequence (``list`` / ``tuple`` / ``set`` / ``frozenset``), **every**
      element must appear as a substring; a plain ``str`` still requires that single substring.
    """
    if websocket_json_message:
        if err_code is not None:
            raise ValueError("assert_http_error: err_code is incompatible with websocket_json_message=True")
        jb_ws = getattr(resp, "json_body", None)
        assert jb_ws is not None, resp
        if err_message is None:
            return jb_ws if isinstance(jb_ws, dict) else None
        assert_err_message_in_text(
            json.dumps(jb_ws, ensure_ascii=False),
            err_message,
            sequence_match="any",
        )
        return jb_ws if isinstance(jb_ws, dict) else None

    if err_code is None and err_message is None:
        return None

    actual = getattr(resp, "status_code", None)
    assert actual is not None, "response missing status_code"

    if err_code is not None:
        if isinstance(err_code, int):
            assert actual == err_code, (resp, err_code)
        else:
            allowed = tuple(err_code)
            assert actual in allowed, (resp, allowed)

    body_bytes, payload = _http_response_body_materialize(resp)

    if err_code is not None and actual is not None and 400 <= actual < 500 and actual != 404:
        assert payload is not None, getattr(resp, "error_message", resp)
        assert "detail" in payload or "error" in payload, payload

    if err_message is not None:
        pieces: list[str] = []
        jb = getattr(resp, "json_body", None)
        if jb is not None:
            pieces.append(json.dumps(jb, ensure_ascii=False))
        em = getattr(resp, "error_message", None)
        if em:
            pieces.append(str(em))
        haystack = "\n".join(pieces) if pieces else body_bytes.decode("utf-8", errors="replace")
        assert_err_message_in_text(haystack, err_message, sequence_match="all")

    return payload


__all__ = [
    "assert_audio_diffusion_response",
    "assert_audio_speech_response",
    "assert_diffusion_response",
    "assert_err_message_in_text",
    "assert_http_error",
    "assert_image_diffusion_response",
    "assert_image_valid",
    "assert_omni_response",
    "assert_video_diffusion_response",
    "assert_video_valid",
    "assert_audio_valid",
]