"""
Server management and performance validation for diffusion tests.
"""

from __future__ import annotations

import asyncio
import base64
import os
import shlex
import subprocess
import sys
import tempfile
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Sequence
from urllib.request import urlopen

import pytest
from openai import Client

from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    globally_suppress_loggers,
    init_logger,
)
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server.common.slack import upload_file_to_slack
from sglang.multimodal_gen.test.server.realtime_consistency import (
    build_realtime_init_payload,
    collect_realtime_output,
    encode_realtime_frames_to_mp4,
    prepare_realtime_first_frame,
    realtime_ws_url,
    record_realtime_key_frames,
    record_realtime_perf_stats,
)
from sglang.multimodal_gen.test.server.testcase_configs import (
    DiffusionSamplingParams,
    PerformanceSummary,
    ScenarioConfig,
    ToleranceConfig,
)
from sglang.multimodal_gen.test.test_utils import (
    get_expected_image_format,
    get_video_frame_count,
    is_image_url,
    prepare_perf_log,
    validate_image,
    validate_image_file,
    validate_openai_video,
    validate_video_file,
)

logger = init_logger(__name__)

globally_suppress_loggers()

FIRST_DENOISE_STEP_TOLERANCE = 4.0
FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 80.0
DECODING_STAGE_MIN_ABS_TOLERANCE_MS = 450.0
VIDEO_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 160.0

# Tracks mesh output file paths from generate_mesh for later correctness validation.
# Keyed by case_id, cleaned up after use.
MESH_OUTPUT_PATHS: dict[str, str] = {}


def _urlopen_with_retry(url: str, timeout: int = 30, max_retries: int = 3) -> bytes:
    """Download content from a URL with retry on transient failures."""
    for attempt in range(max_retries + 1):
        try:
            with urlopen(url, timeout=timeout) as response:
                return response.read()
        except (TimeoutError, OSError) as e:
            if attempt < max_retries:
                wait = 2**attempt
                logger.warning(
                    f"Download attempt {attempt + 1}/{max_retries + 1} failed "
                    f"for {url}: {e}. Retrying in {wait}s..."
                )
                time.sleep(wait)
            else:
                logger.error(
                    f"Failed to download from {url} after "
                    f"{max_retries + 1} attempts: {e}"
                )
                raise


def download_image_from_url(url: str) -> Path:
    """Download an image from a URL to a temporary file.

    Args:
        url: The URL of the image to download

    Returns:
        Path to the downloaded temporary file
    """
    logger.info(f"Downloading image from URL: {url}")

    # Determine file extension from URL
    ext = ".jpg"  # default
    if url.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif")):
        ext = url[url.rfind(".") :]

    # Create temporary file
    temp_file = (
        Path(tempfile.gettempdir()) / f"diffusion_test_image_{int(time.time())}{ext}"
    )

    data = _urlopen_with_retry(url)
    temp_file.write_bytes(data)
    logger.info(f"Downloaded image to: {temp_file}")
    return temp_file


def parse_dimensions(size_string: str | None) -> tuple[int | None, int | None]:
    """Parse a size string in "widthxheight" format to (width, height) tuple.

    Args:
        size_string: Size string in "widthxheight" format (e.g., "1024x1024") or None.
                    Spaces are automatically stripped.

    Returns:
        Tuple of (width, height) as integers if parsing succeeds, (None, None) otherwise.
    """
    if not size_string:
        return (None, None)

    # Strip spaces from the entire string
    size_string = size_string.strip()
    if not size_string:
        return (None, None)

    # Split by "x"
    parts = size_string.split("x")
    if len(parts) != 2:
        return (None, None)

    # Strip spaces from each part and try to convert to int
    try:
        width_str = parts[0].strip()
        height_str = parts[1].strip()

        if not width_str or not height_str:
            return (None, None)

        width = int(width_str)
        height = int(height_str)

        # Validate that both are positive
        if width <= 0 or height <= 0:
            return (None, None)

        return (width, height)
    except ValueError:
        return (None, None)


@dataclass
class ServerContext:
    """Context for a running diffusion server."""

    port: int
    process: subprocess.Popen
    model: str
    stdout_file: Path
    perf_log_path: Path
    log_dir: Path
    _stdout_fh: Any = field(repr=False)
    _log_thread: threading.Thread | None = field(default=None, repr=False)

    def log_tail(self, lines: int = 200) -> str:
        """Return recent server output for failure diagnostics."""
        try:
            content = self.stdout_file.read_text(encoding="utf-8", errors="ignore")
            return "\n".join(content.splitlines()[-lines:])
        except Exception:
            return ""

    def cleanup(self) -> None:
        """Clean up server resources."""
        try:
            kill_process_tree(self.process.pid)
        except Exception:
            pass
        try:
            self._stdout_fh.flush()
            self._stdout_fh.close()
        except Exception:
            pass

        # ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests
        # This is needed because ROCm memory release can be slower than CUDA
        if current_platform.is_hip():
            self._cleanup_rocm_gpu_memory()
            # Clean up downloaded models if HF cache is not persistent
            # This prevents disk exhaustion in CI when cache is not mounted
            self._cleanup_hf_cache_if_not_persistent()
        else:
            # Give the runtime a brief cooldown after server shutdown.
            time.sleep(2)

    def _cleanup_hf_cache_if_not_persistent(self) -> None:
        """Clean up HF cache if it's not on a persistent volume.

        When running in CI without persistent cache, downloaded models accumulate
        and can cause disk/memory exhaustion. This cleans up the model after each
        test if the cache is not persistent.
        """
        import shutil

        hf_home = os.environ.get("HF_HOME", "")
        if not hf_home:
            return

        hf_hub_cache = os.path.join(hf_home, "hub")

        # Check if HF cache is on a persistent volume by looking for a marker file
        # or checking if the directory existed before this test run
        persistent_marker = os.path.join(hf_home, ".persistent_cache")
        if os.path.exists(persistent_marker):
            logger.info("HF cache is persistent, skipping cleanup")
            return

        # Check if the cache directory is empty or was just created
        # If it has very few models, it's likely not persistent
        if not os.path.exists(hf_hub_cache):
            return

        try:
            # Get model cache directories
            model_dirs = [
                d
                for d in os.listdir(hf_hub_cache)
                if d.startswith("models--")
                and os.path.isdir(os.path.join(hf_hub_cache, d))
            ]

            # If there are cached models but no persistent marker, clean up
            # to prevent disk exhaustion in CI
            if model_dirs:
                logger.info(
                    "HF cache appears non-persistent (no .persistent_cache marker), "
                    "cleaning up %d model(s) to prevent disk exhaustion",
                    len(model_dirs),
                )
                for model_dir in model_dirs:
                    model_path = os.path.join(hf_hub_cache, model_dir)
                    try:
                        shutil.rmtree(model_path)
                        logger.info("Cleaned up model cache: %s", model_dir)
                    except Exception as e:
                        logger.warning("Failed to clean up %s: %s", model_dir, e)
        except Exception as e:
            logger.warning("Error during HF cache cleanup: %s", e)

    def _cleanup_rocm_gpu_memory(self) -> None:
        """ROCm-specific cleanup to ensure GPU memory is fully released."""
        import gc

        # Wait for process to fully terminate
        try:
            self.process.wait(timeout=30)
        except Exception:
            pass

        # Force garbage collection multiple times
        for _ in range(3):
            gc.collect()

        # Clear HIP memory on all GPUs
        try:
            import torch

            for i in range(torch.cuda.device_count()):
                with torch.cuda.device(i):
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
        except Exception:
            pass

        # Wait for GPU memory to be released (ROCm can be much slower than CUDA)
        # The GPU driver needs time to reclaim memory from killed processes
        time.sleep(15)


class ServerManager:
    """Manages diffusion server lifecycle."""

    def __init__(
        self,
        model: str,
        port: int,
        wait_deadline: float = 1200.0,
        extra_args: str = "",
        env_vars: dict[str, str] | None = None,
    ):
        self.model = model
        self.port = port
        self.wait_deadline = wait_deadline
        self.extra_args = extra_args
        self.env_vars = env_vars or {}

    def _wait_for_rocm_gpu_memory_clear(self, max_wait: float = 60.0) -> None:
        """ROCm-specific: Wait for GPU memory to be mostly free before starting.

        ROCm GPU memory release from killed processes can be significantly slower
        than CUDA, so we need to wait longer and be more patient.
        """
        try:
            import torch

            if not torch.cuda.is_available():
                return

            start_time = time.time()
            last_total_used = float("inf")

            while time.time() - start_time < max_wait:
                # Check GPU memory usage
                total_used = 0
                for i in range(torch.cuda.device_count()):
                    mem_info = torch.cuda.mem_get_info(i)
                    free, total = mem_info
                    used = total - free
                    total_used += used

                # If less than 5GB is used across all GPUs, we're good
                if total_used < 5 * 1024 * 1024 * 1024:  # 5GB
                    logger.info(
                        "[server-test] ROCm GPU memory is clear (used: %.2f GB)",
                        total_used / (1024**3),
                    )
                    return

                # Log progress
                elapsed = int(time.time() - start_time)
                if total_used < last_total_used:
                    logger.info(
                        "[server-test] ROCm: GPU memory clearing (used: %.2f GB, elapsed: %ds)",
                        total_used / (1024**3),
                        elapsed,
                    )
                else:
                    logger.info(
                        "[server-test] ROCm: Waiting for GPU memory (used: %.2f GB, elapsed: %ds)",
                        total_used / (1024**3),
                        elapsed,
                    )
                last_total_used = total_used
                time.sleep(3)

            # Final warning with detailed GPU info
            logger.warning(
                "[server-test] ROCm GPU memory not fully cleared after %.0fs (used: %.2f GB). "
                "Proceeding anyway - this may cause OOM.",
                max_wait,
                total_used / (1024**3),
            )
        except Exception as e:
            logger.debug("[server-test] Could not check ROCm GPU memory: %s", e)

    def start(self) -> ServerContext:
        """Start the diffusion server and wait for readiness."""
        # ROCm/AMD: Wait for GPU memory to be clear before starting
        # This prevents OOM when running sequential tests on ROCm
        if current_platform.is_hip():
            self._wait_for_rocm_gpu_memory_clear()

        log_dir, perf_log_path = prepare_perf_log()

        safe_model_name = self.model.replace("/", "_")
        stdout_path = (
            Path(tempfile.gettempdir())
            / f"sgl_server_{self.port}_{safe_model_name}.log"
        )
        stdout_path.unlink(missing_ok=True)

        command = [
            "sglang",
            "serve",
            "--model-path",
            self.model,
            "--port",
            str(self.port),
            "--log-level=debug",
        ]
        if self.extra_args.strip():
            command.extend(self.extra_args.strip().split())

        env = os.environ.copy()
        env["SGLANG_DIFFUSION_STAGE_LOGGING"] = "1"
        env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix()

        # Apply custom environment variables
        env.update(self.env_vars)

        cmd_str = shlex.join(command)
        # Use print (not logger) so the command always appears in CI output
        # regardless of log-level configuration.
        print(f"[server-test] Running command: {cmd_str}", flush=True)

        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            env=env,
        )

        log_thread = None
        stdout_fh = stdout_path.open("w", encoding="utf-8", buffering=1)
        if process.stdout:

            def _log_pipe(pipe: Any, file: Any) -> None:
                """Read from pipe and write to file and stdout."""
                try:
                    with pipe:
                        for line in iter(pipe.readline, ""):
                            sys.stdout.write(line)
                            sys.stdout.flush()
                            file.write(line)
                            file.flush()
                except Exception as e:
                    logger.error("Log pipe thread error: %s", e)
                finally:
                    file.close()
                    logger.debug("Log pipe thread finished.")

            log_thread = threading.Thread(
                target=_log_pipe, args=(process.stdout, stdout_fh)
            )
            log_thread.daemon = True
            log_thread.start()

        print(
            f"[server-test] Starting server pid={process.pid}, "
            f"model={self.model}, log={stdout_path}",
            flush=True,
        )

        self._wait_for_ready(process, stdout_path)

        return ServerContext(
            port=self.port,
            process=process,
            model=self.model,
            stdout_file=stdout_path,
            perf_log_path=perf_log_path,
            log_dir=log_dir,
            _stdout_fh=stdout_fh,
            _log_thread=log_thread,
        )

    def _wait_for_ready(self, process: subprocess.Popen, stdout_path: Path) -> None:
        """Wait for server to become ready."""
        start = time.time()
        ready_message = "Application startup complete."
        log_period = 30
        prev_log_period_count = 0

        while time.time() - start < self.wait_deadline:
            if process.poll() is not None:
                tail = self._get_log_tail(stdout_path)
                raise RuntimeError(
                    f"Server exited early (code {process.returncode}).\n{tail}"
                )

            if stdout_path.exists():
                try:
                    content = stdout_path.read_text(encoding="utf-8", errors="ignore")
                    if ready_message in content:
                        logger.info("[server-test] Server ready")
                        return
                except Exception as e:
                    logger.debug("Could not read log yet: %s", e)

            elapsed = int(time.time() - start)
            if (elapsed // log_period) > prev_log_period_count:
                prev_log_period_count = elapsed // log_period
                logger.info("[server-test] Waiting for server... elapsed=%ss", elapsed)
            time.sleep(1)

        tail = self._get_log_tail(stdout_path)
        raise TimeoutError(f"Server not ready within {self.wait_deadline}s.\n{tail}")

    @staticmethod
    def _get_log_tail(path: Path, lines: int = 200) -> str:
        """Get the last N lines from a log file."""
        try:
            content = path.read_text(encoding="utf-8", errors="ignore")
            return "\n".join(content.splitlines()[-lines:])
        except Exception:
            return ""


class PerformanceValidator:
    """Validates performance metrics against expectations."""

    is_video_gen: bool = False

    def __init__(
        self,
        scenario: ScenarioConfig,
        tolerances: ToleranceConfig,
        step_fractions: Sequence[float],
    ):
        self.scenario = scenario
        self.tolerances = tolerances
        self.step_fractions = step_fractions
        self.is_baseline_generation_mode = (
            os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
        )

    def _assert_le(
        self,
        name: str,
        actual: float,
        expected: float,
        tolerance: float,
        min_abs_tolerance_ms: float = 20.0,
    ):
        """Assert that actual is less than or equal to expected within a tolerance.

        Uses the larger of relative tolerance or absolute tolerance to prevent
        flaky failures on very fast operations.

        For AMD GPUs, uses 100% higher tolerance and issues warning instead of assertion.
        """
        # Check if running on AMD GPU
        is_amd = current_platform.is_hip()

        if is_amd:
            # Use 100% higher tolerance for AMD (2x the expected value)
            amd_tolerance = 1.0  # 100%
            upper_bound = calculate_upper_bound(
                expected, amd_tolerance, min_abs_tolerance_ms
            )
            if actual > upper_bound:
                logger.warning(
                    f"[AMD PERF WARNING] Validation would fail for '{name}'.\n"
                    f"  Actual:   {actual:.4f}ms\n"
                    f"  Expected: {expected:.4f}ms\n"
                    f"  AMD Limit: {upper_bound:.4f}ms "
                    f"(rel_tol: {amd_tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)\n"
                    f"  Original tolerance was: {tolerance:.1%}"
                )
        else:
            upper_bound = calculate_upper_bound(
                expected, tolerance, min_abs_tolerance_ms
            )
            assert actual <= upper_bound, (
                f"Validation failed for '{name}'.\n"
                f"  Actual:   {actual:.4f}ms\n"
                f"  Expected: {expected:.4f}ms\n"
                f"  Limit:    {upper_bound:.4f}ms "
                f"(rel_tol: {tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)"
            )

    def validate(
        self, perf_record: RequestPerfRecord, *args, **kwargs
    ) -> PerformanceSummary:
        """Validate all performance metrics and return summary."""
        summary = self.collect_metrics(perf_record)
        if self.is_baseline_generation_mode:
            return summary

        self._validate_e2e(summary)
        self._validate_denoise_agg(summary)
        self._validate_denoise_steps(summary)
        self._validate_stages(summary)

        return summary

    def collect_metrics(
        self,
        perf_record: RequestPerfRecord,
    ) -> PerformanceSummary:
        return PerformanceSummary.from_req_perf_record(perf_record, self.step_fractions)

    def _validate_e2e(self, summary: PerformanceSummary) -> None:
        """Validate end-to-end performance."""
        assert summary.e2e_ms > 0, "E2E duration missing"
        self._assert_le(
            "E2E Latency",
            summary.e2e_ms,
            self.scenario.expected_e2e_ms,
            self.tolerances.e2e,
        )

    def _validate_denoise_agg(self, summary: PerformanceSummary) -> None:
        """Validate aggregate denoising metrics."""
        assert summary.avg_denoise_ms > 0, "Denoising step timings missing"

        self._assert_le(
            "Average Denoise Step",
            summary.avg_denoise_ms,
            self.scenario.expected_avg_denoise_ms,
            self.tolerances.denoise_agg,
        )
        self._assert_le(
            "Median Denoise Step",
            summary.median_denoise_ms,
            self.scenario.expected_median_denoise_ms,
            self.tolerances.denoise_agg,
        )

    def _validate_denoise_steps(self, summary: PerformanceSummary) -> None:
        """Validate individual denoising steps."""
        for idx, actual in summary.sampled_steps.items():
            expected = self.scenario.denoise_step_ms.get(idx)
            if expected is None:
                continue
            if idx == 0:
                # server warmup is generic, so the first real step can still
                # pay request-shape/path lazy init that is not a steady-state signal
                self._assert_le(
                    f"Denoise Step {idx}",
                    actual,
                    expected,
                    FIRST_DENOISE_STEP_TOLERANCE,
                    min_abs_tolerance_ms=FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
                )
                continue

            self._assert_le(
                f"Denoise Step {idx}",
                actual,
                expected,
                self.tolerances.denoise_step,
            )

    def _validate_stages(self, summary: PerformanceSummary) -> None:
        """Validate stage-level metrics."""
        assert summary.stage_metrics, "Stage metrics missing"

        for stage, expected in self.scenario.stages_ms.items():
            if stage == "per_frame_generation" and self.is_video_gen:
                continue
            actual = summary.stage_metrics.get(stage)
            assert actual is not None, f"Stage {stage} timing missing"
            tolerance = (
                self.tolerances.denoise_stage
                if stage == "DenoisingStage"
                else self.tolerances.non_denoise_stage
            )
            if stage.endswith("DecodingStage"):
                tolerance = max(tolerance, 0.9)
                min_abs_tolerance_ms = DECODING_STAGE_MIN_ABS_TOLERANCE_MS
            else:
                min_abs_tolerance_ms = 120.0
            self._assert_le(
                f"Stage '{stage}'",
                actual,
                expected,
                tolerance,
                min_abs_tolerance_ms=min_abs_tolerance_ms,
            )


class VideoPerformanceValidator(PerformanceValidator):
    """Extended validator for video diffusion with frame-level metrics."""

    is_video_gen = True

    def _validate_denoise_steps(self, summary: PerformanceSummary) -> None:
        """Validate individual denoising steps."""
        for idx, actual in summary.sampled_steps.items():
            expected = self.scenario.denoise_step_ms.get(idx)
            if expected is None:
                continue
            if idx == 0:
                self._assert_le(
                    f"Denoise Step {idx}",
                    actual,
                    expected,
                    FIRST_DENOISE_STEP_TOLERANCE,
                    min_abs_tolerance_ms=FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
                )
                continue

            # video per-step samples can catch one-off scheduling/offload jitter;
            # avg and median denoise checks remain the steady-state guard
            self._assert_le(
                f"Denoise Step {idx}",
                actual,
                expected,
                self.tolerances.denoise_step,
                min_abs_tolerance_ms=VIDEO_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
            )

    def validate(
        self,
        perf_record: RequestPerfRecord,
        num_frames: int | None = None,
    ) -> PerformanceSummary:
        """Validate video metrics including frame generation rates."""
        summary = super().validate(perf_record)

        if num_frames and summary.e2e_ms > 0:
            summary.total_frames = num_frames
            summary.avg_frame_time_ms = summary.e2e_ms / num_frames
            summary.frames_per_second = 1000.0 / summary.avg_frame_time_ms

            if not self.is_baseline_generation_mode:
                self._validate_frame_rate(summary)

        return summary

    def _validate_frame_rate(self, summary: PerformanceSummary) -> None:
        """Validate frame generation performance."""
        expected_frame_time = self.scenario.stages_ms.get("per_frame_generation")
        if expected_frame_time and summary.avg_frame_time_ms:
            self._assert_le(
                "Average Frame Time",
                summary.avg_frame_time_ms,
                expected_frame_time,
                self.tolerances.denoise_stage,
            )


class MeshValidator(PerformanceValidator):
    """Validator for 3D mesh generation. Inherits perf validation from PerformanceValidator."""

    pass


# Pinned to a ci-data commit (not main): invalidates the per-URL download cache
# whenever the reference is regenerated, and keeps the mesh GT reproducible.
# Bump this SHA when pushing a new hunyuan3d.glb to ci-data.
HUNYUAN3D_REFERENCE_URL = (
    "https://raw.githubusercontent.com/sgl-project/ci-data/"
    "395f6e49c37d22a57d79fbcd3653d43984099ae2"
    "/diffusion-ci/consistency_gt/1-gpu/hunyuan3d_2_0/hunyuan3d.glb"
)


def _download_reference_mesh(url: str) -> Path:
    """Download a reference mesh from URL, caching in temp dir.

    Validates that the cached/downloaded file actually *loads* as a non-empty
    mesh — not just that a magic/length header looks right. raw.githubusercontent
    can briefly serve a truncated or corrupt response for a just-pushed large
    file, and a prior run may have cached those bytes on a persistent runner; a
    size/magic check can't catch a blob whose byte count matches the declared
    length but whose body is corrupt (exactly what poisoned this CI cache and
    surfaced as a cryptic trimesh "incorrect header on GLB file" deep inside
    validation). Loading via trimesh rejects any such cache (forcing a
    re-download) and turns a bad fresh download into a clear error. The ``v2``
    cache prefix also invalidates blobs written by the earlier, weaker checks.
    """
    import hashlib

    cache_name = f"ref_mesh_v2_{hashlib.md5(url.encode()).hexdigest()}.glb"
    cache_path = Path(tempfile.gettempdir()) / cache_name

    def _loads_as_mesh(path: Path) -> bool:
        try:
            import trimesh

            mesh = trimesh.load(str(path), force="mesh")
            return (
                getattr(mesh, "vertices", None) is not None and len(mesh.vertices) > 0
            )
        except Exception:
            return False

    if cache_path.exists() and _loads_as_mesh(cache_path):
        logger.info(f"Using cached reference mesh: {cache_path}")
        return cache_path

    logger.info(f"Downloading reference mesh from: {url}")
    cache_path.write_bytes(_urlopen_with_retry(url, timeout=60))
    if not _loads_as_mesh(cache_path):
        size = cache_path.stat().st_size if cache_path.exists() else 0
        cache_path.unlink(missing_ok=True)
        raise RuntimeError(
            f"Reference mesh from {url} did not load as a valid mesh "
            f"({size} bytes). The CDN may not have propagated a recently-pushed "
            f"file yet; retry shortly."
        )
    logger.info(f"Reference mesh cached at: {cache_path}")
    return cache_path


def validate_mesh_correctness(
    generated_mesh_path: str,
    reference_url: str = HUNYUAN3D_REFERENCE_URL,
    num_sample_points: int = 4096,
    cd_threshold_ratio: float = 0.01,
    random_seed: int = 42,
):
    """Validate mesh geometric similarity against a reference via Chamfer Distance.

    Downloads the reference mesh from a URL (cached), samples point clouds from
    both meshes, and asserts Chamfer Distance is within threshold.
    """
    import numpy as np

    try:
        import trimesh
    except ImportError:
        pytest.fail("trimesh is required for mesh validation: pip install trimesh")

    from scipy.spatial import cKDTree

    # Load generated mesh
    generated_mesh = trimesh.load(generated_mesh_path)
    if isinstance(generated_mesh, trimesh.Scene):
        generated_mesh = generated_mesh.dump(concatenate=True)

    # Download and load reference mesh
    ref_path = _download_reference_mesh(reference_url)
    reference_mesh = trimesh.load(str(ref_path))
    if isinstance(reference_mesh, trimesh.Scene):
        reference_mesh = reference_mesh.dump(concatenate=True)

    # Bounding box diagonal for threshold normalization
    ref_bbox = reference_mesh.bounding_box.bounds
    bbox_diagonal = float(np.linalg.norm(ref_bbox[1] - ref_bbox[0]))
    cd_threshold = cd_threshold_ratio * bbox_diagonal

    # Sample point clouds
    np.random.seed(random_seed)
    gen_points = np.array(
        generated_mesh.sample(num_sample_points, return_index=True)[0]
    )
    ref_points = np.array(
        reference_mesh.sample(num_sample_points, return_index=True)[0]
    )

    # Bidirectional Chamfer Distance
    tree1 = cKDTree(gen_points)
    tree2 = cKDTree(ref_points)
    forward_cd = float(np.mean(tree2.query(gen_points)[0] ** 2))
    backward_cd = float(np.mean(tree1.query(ref_points)[0] ** 2))
    total_cd = forward_cd + backward_cd

    assert total_cd <= cd_threshold, (
        f"Chamfer Distance check failed: total_cd={total_cd:.6f}, "
        f"threshold={cd_threshold:.6f} ({cd_threshold_ratio * 100:.2f}% of bbox diagonal {bbox_diagonal:.4f})"
    )


# Registry of validators by name
VALIDATOR_REGISTRY = {
    "default": PerformanceValidator,
    "video": VideoPerformanceValidator,
    "mesh": MeshValidator,
    "action": PerformanceValidator,
}


def _extract_async_job_error_message(job: Any) -> str | None:
    error = getattr(job, "error", None)
    if error is None and isinstance(job, dict):
        error = job.get("error")

    if error is None:
        return None

    if isinstance(error, dict):
        for key in ("message", "detail", "error"):
            value = error.get(key)
            if value:
                return str(value)
        return str(error)

    message = getattr(error, "message", None)
    if message:
        return str(message)

    return str(error)


def get_generate_fn(
    model_path: str,
    modality: str,
    sampling_params: DiffusionSamplingParams,
) -> Callable[[str, Client], tuple[str, bytes]]:
    """Return appropriate generation function for the case."""
    # Allow override via environment variable (useful for AMD where large resolutions cause slow VAE)
    output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size)
    n = sampling_params.num_outputs_per_prompt

    def _create_and_download_video(
        client,
        case_id,
        *,
        model: str,
        size: str,
        prompt: str | None = None,
        seconds: int | None = None,
        input_reference: Any | None = None,
        extra_body: dict[Any] | None = None,
        expected_frame_count: int | None = None,
    ) -> str:
        """
        Create a video job via /v1/videos, poll until completion,
        then download the binary content and validate it.

        Returns request-id
        """

        create_kwargs: dict[str, Any] = {
            "model": model,
            "size": size,
        }
        if prompt is not None:
            create_kwargs["prompt"] = prompt
        if seconds is not None:
            create_kwargs["seconds"] = seconds
        if input_reference is not None:
            create_kwargs["input_reference"] = input_reference  # triggers multipart
        if extra_body is not None:
            create_kwargs["extra_body"] = extra_body

        job = client.videos.create(**create_kwargs)  # type: ignore[attr-defined]
        video_id = job.id

        job_completed = False
        is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
        # Check if running on AMD GPU - use longer timeout
        is_amd = current_platform.is_hip()
        if is_baseline_generation_mode:
            timeout = 3600.0
        elif is_amd:
            timeout = 2400.0  # 40 minutes for AMD
        else:
            timeout = 1200.0
        deadline = time.time() + timeout
        while True:
            page = client.videos.list()  # type: ignore[attr-defined]
            item = next((v for v in page.data if v.id == video_id), None)
            status = getattr(item, "status", None) if item is not None else None

            if status == "completed":
                job_completed = True
                break

            if status == "failed":
                error_message = (
                    _extract_async_job_error_message(item) or "unknown error"
                )
                pytest.fail(
                    f"{case_id}: video job {video_id} failed early: {error_message}"
                )

            if status in {"cancelled", "deleted"}:
                pytest.fail(
                    f"{case_id}: video job {video_id} ended with status={status}"
                )

            if time.time() > deadline:
                break

            time.sleep(1)

        if not job_completed:
            if is_baseline_generation_mode:
                logger.warning(
                    f"{case_id}: video job {video_id} timed out during baseline generation. "
                    "Attempting to collect performance data anyway."
                )
                return (video_id, b"")

            if is_amd:
                logger.warning(
                    f"[AMD TIMEOUT WARNING] {case_id}: video job {video_id} did not complete "
                    f"within {timeout}s timeout. This may indicate performance issues on AMD."
                )
                pytest.skip(
                    f"{case_id}: video job timed out on AMD after {timeout}s - skipping"
                )

            pytest.fail(f"{case_id}: video job {video_id} did not complete in time")

        # download video
        resp = client.videos.download_content(video_id=video_id)  # type: ignore[attr-defined]
        content = resp.read()
        validate_openai_video(content)

        expected_filename = f"{video_id}.mp4"
        tmp_path = expected_filename
        with open(tmp_path, "wb") as f:
            f.write(content)

        # Validate output file
        expected_width, expected_height = parse_dimensions(size)
        if (
            extra_body is not None
            and extra_body.get("enable_upscaling")
            and expected_width
            and expected_height
        ):
            scale = extra_body.get("upscaling_scale", 4)
            expected_width *= scale
            expected_height *= scale
        validate_video_file(
            tmp_path, expected_filename, expected_width, expected_height
        )

        if expected_frame_count is not None:
            actual_count = get_video_frame_count(tmp_path)
            assert actual_count == expected_frame_count, (
                f"{case_id}: frame count mismatch after interpolation — "
                f"expected {expected_frame_count}, got {actual_count}"
            )

        upload_file_to_slack(
            case_id=case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            file_path=tmp_path,
            origin_file_path=sampling_params.image_path,
        )
        os.remove(tmp_path)

        return (video_id, content)

    video_seconds = sampling_params.seconds or 4

    def generate_image(case_id, client) -> tuple[str, bytes]:
        """T2I: Text to Image generation."""
        if not sampling_params.prompt:
            pytest.skip(f"{case_id}: no text prompt configured")

        # Request parameters that affect output format
        req_output_format = sampling_params.output_format
        req_background = None  # Not specified in current request

        # Build extra_body for optional features
        extra_body = dict(sampling_params.extras)

        response = client.images.with_raw_response.generate(
            model=model_path,
            prompt=sampling_params.prompt,
            n=n,
            size=output_size,
            response_format="b64_json",
            output_format=req_output_format,
            extra_body=extra_body if extra_body else None,
        )
        result = response.parse()
        validate_image(result.data[0].b64_json)

        rid = result.id

        img_data = base64.b64decode(result.data[0].b64_json)
        # Infer expected format from request parameters
        expected_ext = get_expected_image_format(req_output_format, req_background)
        expected_filename = f"{result.created}.{expected_ext}"
        tmp_path = expected_filename
        with open(tmp_path, "wb") as f:
            f.write(img_data)

        # Validate output file
        expected_width, expected_height = parse_dimensions(output_size)
        if (
            sampling_params.extras.get("enable_upscaling")
            and expected_width
            and expected_height
        ):
            expected_width *= sampling_params.extras.get("upscaling_scale", 4)
            expected_height *= sampling_params.extras.get("upscaling_scale", 4)
        validate_image_file(
            tmp_path,
            expected_filename,
            expected_width,
            expected_height,
            output_format=req_output_format,
            background=req_background,
        )

        upload_file_to_slack(
            case_id=case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            file_path=tmp_path,
        )
        os.remove(tmp_path)

        return (rid, img_data)

    def generate_image_edit(case_id, client) -> tuple[str, bytes]:
        """TI2I: Text + Image -> Image edit."""
        if not sampling_params.prompt or not sampling_params.image_path:
            pytest.skip(f"{case_id}: no edit config")

        image_paths = sampling_params.image_path

        if not isinstance(image_paths, list):
            image_paths = [image_paths]

        new_image_paths = []
        for image_path in image_paths:
            if is_image_url(image_path):
                new_image_paths.append(download_image_from_url(str(image_path)))
            else:
                local_path = Path(image_path)
                new_image_paths.append(local_path)
                if not local_path.exists():
                    pytest.skip(f"{case_id}: file missing: {image_path}")

        image_paths = new_image_paths

        # Request parameters that affect output format
        req_output_format = (
            sampling_params.output_format
        )  # Not specified in current request
        req_background = None  # Not specified in current request

        # Build extra_body for optional features
        extra_body = {"num_frames": sampling_params.num_frames}
        extra_body.update(sampling_params.extras)

        images = [open(image_path, "rb") for image_path in image_paths]
        try:
            response = client.images.with_raw_response.edit(
                model=model_path,
                image=images,
                prompt=sampling_params.prompt,
                n=n,
                size=output_size,
                response_format="b64_json",
                output_format=req_output_format,
                extra_body=extra_body,
            )
        finally:
            for img in images:
                img.close()

        result = response.parse()
        validate_image(result.data[0].b64_json)

        img_data = base64.b64decode(result.data[0].b64_json)
        rid = result.id

        # Infer expected format from request parameters
        expected_ext = get_expected_image_format(req_output_format, req_background)
        expected_filename = f"{rid}.{expected_ext}"
        tmp_path = expected_filename
        with open(tmp_path, "wb") as f:
            f.write(img_data)

        # Validate output file
        expected_width, expected_height = parse_dimensions(output_size)
        validate_image_file(
            tmp_path,
            expected_filename,
            expected_width,
            expected_height,
            output_format=req_output_format,
            background=req_background,
        )

        upload_file_to_slack(
            case_id=case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            file_path=tmp_path,
            origin_file_path=sampling_params.image_path,
        )
        os.remove(tmp_path)

        return (rid, img_data)

    def generate_image_edit_url(case_id, client) -> tuple[str, bytes]:
        """TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download)."""
        if not sampling_params.prompt or not sampling_params.image_path:
            pytest.skip(f"{case_id}: no edit config")
        # Handle both single URL and list of URLs
        image_urls = sampling_params.image_path
        if not isinstance(image_urls, list):
            image_urls = [image_urls]

        # Validate all URLs
        for url in image_urls:
            if not is_image_url(url):
                pytest.skip(
                    f"{case_id}: image_path must be a URL for URL direct test: {url}"
                )

        # Request parameters that affect output format
        req_output_format = (
            sampling_params.output_format
        )  # Not specified in current request
        req_background = None  # Not specified in current request

        response = client.images.with_raw_response.edit(
            model=model_path,
            prompt=sampling_params.prompt,
            image=[],  # Only for OpenAI verification
            n=n,
            size=sampling_params.output_size,
            response_format="b64_json",
            output_format=req_output_format,
            extra_body={"url": image_urls, "num_frames": sampling_params.num_frames},
        )

        result = response.parse()
        rid = result.id

        validate_image(result.data[0].b64_json)

        # Save and upload result for verification
        img_data = base64.b64decode(result.data[0].b64_json)
        # Infer expected format from request parameters
        expected_ext = get_expected_image_format(req_output_format, req_background)
        expected_filename = f"{rid}.{expected_ext}"
        tmp_path = expected_filename
        with open(tmp_path, "wb") as f:
            f.write(img_data)

        # Validate output file
        expected_width, expected_height = parse_dimensions(sampling_params.output_size)
        validate_image_file(
            tmp_path,
            expected_filename,
            expected_width,
            expected_height,
            output_format=req_output_format,
            background=req_background,
        )

        upload_file_to_slack(
            case_id=case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            file_path=tmp_path,
            origin_file_path=str(sampling_params.image_path),
        )
        os.remove(tmp_path)

        return (rid, img_data)

    def generate_video(case_id, client) -> tuple[str, bytes]:
        """T2V: Text ? Video."""
        if not sampling_params.prompt:
            pytest.skip(f"{case_id}: no text prompt configured")

        # Build extra_body for optional features
        extra_body = dict(sampling_params.extras)
        if sampling_params.num_frames:
            extra_body["num_frames"] = sampling_params.num_frames

        # Compute expected output frame count for validation
        expected_frame_count = None
        if (
            sampling_params.extras.get("enable_frame_interpolation")
            and sampling_params.num_frames
        ):
            n = sampling_params.num_frames
            exp = sampling_params.extras.get("frame_interpolation_exp", 1)
            expected_frame_count = (n - 1) * (2**exp) + 1

        return _create_and_download_video(
            client,
            case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            size=output_size,
            seconds=video_seconds,
            extra_body=extra_body if extra_body else None,
            expected_frame_count=expected_frame_count,
        )

    def generate_image_to_video(case_id, client) -> tuple[str, bytes]:
        """I2V: Image -> Video (optional prompt)."""
        if not sampling_params.image_path:
            pytest.skip(f"{case_id}: no input image configured")

        if is_image_url(sampling_params.image_path):
            image_path = download_image_from_url(str(sampling_params.image_path))
        else:
            image_path = Path(sampling_params.image_path)
            if not image_path.exists():
                pytest.skip(f"{case_id}: file missing: {image_path}")

        # Build extra_body for optional features
        extra_body = dict(sampling_params.extras)

        with image_path.open("rb") as fh:
            return _create_and_download_video(
                client,
                case_id,
                model=model_path,
                prompt=sampling_params.prompt,
                size=output_size,
                seconds=video_seconds,
                input_reference=fh,
                extra_body=extra_body if extra_body else None,
            )

    def generate_text_url_image_to_video(case_id, client) -> tuple[str, bytes]:
        if not sampling_params.prompt or not sampling_params.image_path:
            pytest.skip(f"{case_id}: no edit config")

        # Build extra_body for optional features
        extra_body = {"reference_url": sampling_params.image_path}
        extra_body.update(sampling_params.extras)

        return _create_and_download_video(
            client,
            case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            size=sampling_params.output_size,
            seconds=video_seconds,
            extra_body={
                "reference_url": sampling_params.image_path,
                "fps": sampling_params.fps,
                "num_frames": sampling_params.num_frames,
            },
        )

    def generate_text_image_to_video(case_id, client) -> tuple[str, bytes]:
        """TI2V: Text + Image -> Video."""
        if not sampling_params.prompt or not sampling_params.image_path:
            pytest.skip(f"{case_id}: no edit config")

        if is_image_url(sampling_params.image_path):
            image_path = download_image_from_url(str(sampling_params.image_path))
        else:
            image_path = Path(sampling_params.image_path)
            if not image_path.exists():
                pytest.skip(f"{case_id}: file missing: {image_path}")

        # Build extra_body for optional features
        extra_body = dict(sampling_params.extras)

        with image_path.open("rb") as fh:
            return _create_and_download_video(
                client,
                case_id,
                model=model_path,
                prompt=sampling_params.prompt,
                size=output_size,
                seconds=video_seconds,
                input_reference=fh,
                extra_body={
                    "fps": sampling_params.fps,
                    "num_frames": sampling_params.num_frames,
                    **extra_body,
                },
            )

    def generate_realtime_video(case_id, client) -> tuple[str, bytes]:
        """Realtime video generation folded back into mp4 for consistency checks."""
        if not sampling_params.prompt:
            pytest.skip(f"{case_id}: no realtime prompt configured")
        if sampling_params.realtime_num_chunks is None:
            pytest.skip(f"{case_id}: realtime_num_chunks is not configured")
        if sampling_params.realtime_num_chunks <= 0:
            pytest.fail(f"{case_id}: realtime_num_chunks must be positive")

        first_frame = prepare_realtime_first_frame(sampling_params.image_path)
        init_payload = build_realtime_init_payload(
            model_path=model_path,
            sampling_params=sampling_params,
            output_size=output_size,
            first_frame=first_frame,
        )
        realtime_output = asyncio.run(
            collect_realtime_output(
                ws_url=realtime_ws_url(client),
                init_payload=init_payload,
                events=list(sampling_params.realtime_events),
                num_chunks=sampling_params.realtime_num_chunks,
                require_chunk_stats=bool(sampling_params.realtime_perf_thresholds),
            )
        )
        record_realtime_perf_stats(case_id, realtime_output.chunk_stats)
        record_realtime_key_frames(case_id, realtime_output.frames)
        fps = int(sampling_params.fps or 24)
        video_bytes = encode_realtime_frames_to_mp4(realtime_output.frames, fps=fps)
        validate_openai_video(video_bytes)

        rid = f"{case_id}-realtime"
        expected_filename = f"{rid}.mp4"
        tmp_path = expected_filename
        Path(tmp_path).write_bytes(video_bytes)
        expected_width, expected_height = parse_dimensions(output_size)
        validate_video_file(
            tmp_path, expected_filename, expected_width, expected_height
        )
        upload_file_to_slack(
            case_id=case_id,
            model=model_path,
            prompt=sampling_params.prompt,
            file_path=tmp_path,
            origin_file_path=sampling_params.image_path,
        )
        os.remove(tmp_path)
        return (rid, video_bytes)

    def generate_mesh(case_id, client) -> tuple[str, bytes]:
        """I2M: Image to Mesh generation using async /v1/meshes API."""
        import requests as http_requests

        if not sampling_params.image_path:
            pytest.skip(f"{case_id}: no input image configured for mesh generation")

        image_path = sampling_params.image_path
        if isinstance(image_path, str) and is_image_url(image_path):
            image_path = download_image_from_url(image_path)
        elif isinstance(image_path, Path):
            if not image_path.exists():
                pytest.skip(f"{case_id}: image file missing: {image_path}")
        else:
            image_path = Path(str(image_path))
            if not image_path.exists():
                pytest.skip(f"{case_id}: image file missing: {image_path}")

        base_url = str(client.base_url).rstrip("/")
        if base_url.endswith("/v1"):
            base_url = base_url[:-3]

        create_url = f"{base_url}/v1/meshes"

        with open(str(image_path), "rb") as img_file:
            files = {"image": (Path(str(image_path)).name, img_file, "image/png")}
            data = {
                "prompt": "generate 3d mesh",
                "model": model_path,
                "seed": "0",
                "guidance_scale": "5.0",
                "num_inference_steps": "50",
            }

            logger.info(f"[Mesh Gen] Sending request to {create_url}")

            try:
                response = http_requests.post(
                    create_url, files=files, data=data, timeout=60
                )
            except Exception as e:
                pytest.fail(f"{case_id}: mesh creation request failed: {e}")

        if response.status_code != 200:
            pytest.fail(f"{case_id}: mesh creation failed: {response.text}")

        job = response.json()
        mesh_id = job.get("id")
        if not mesh_id:
            pytest.fail(f"{case_id}: no mesh id in response: {job}")

        poll_url = f"{base_url}/v1/meshes/{mesh_id}"
        poll_interval = 5
        max_wait = 1200
        elapsed = 0

        while elapsed < max_wait:
            time.sleep(poll_interval)
            elapsed += poll_interval

            try:
                poll_resp = http_requests.get(poll_url, timeout=30)
            except Exception as e:
                logger.warning(f"[Mesh Gen] Poll failed: {e}")
                continue

            if poll_resp.status_code != 200:
                continue

            status_data = poll_resp.json()
            status = status_data.get("status", "")

            if status == "completed":
                content_url = f"{base_url}/v1/meshes/{mesh_id}/content"
                try:
                    content_resp = http_requests.get(content_url, timeout=60)
                except Exception as e:
                    pytest.fail(f"{case_id}: mesh download failed: {e}")

                if content_resp.status_code != 200:
                    pytest.fail(f"{case_id}: mesh download failed: {content_resp.text}")

                content = content_resp.content
                # Shape-only Hunyuan3D meshes are returned as OBJ, painted meshes
                # as GLB. Pick the extension from the content magic so trimesh.load
                # (which dispatches on the file extension) parses it correctly,
                # instead of raising "incorrect header on GLB file" when an OBJ
                # body is saved under a .glb name.
                ext = ".glb" if content[:4] == b"glTF" else ".obj"
                temp_path = Path(tempfile.gettempdir()) / f"mesh_test_{mesh_id}{ext}"
                temp_path.write_bytes(content)
                MESH_OUTPUT_PATHS[case_id] = str(temp_path)

                logger.info(f"[Mesh Gen] Mesh downloaded to {temp_path}")
                return (mesh_id, b"")
            elif status == "failed":
                error = status_data.get("error", {})
                pytest.fail(f"{case_id}: mesh generation failed: {error}")

        pytest.fail(f"{case_id}: mesh generation timed out after {max_wait}s")

    def generate_action(case_id, client) -> tuple[str, bytes]:
        """VLA action generation using /v1/actions/generations."""
        import numpy as np
        import requests as http_requests

        extra = dict(sampling_params.extras)
        action_horizon = int(extra.get("action_horizon", 50))
        action_dim = int(extra.get("action_dim", 32))
        state_dim = int(extra.get("state_dim", action_dim))
        image_size = int(extra.get("image_size", 64))
        camera_order = tuple(
            extra.get(
                "camera_order",
                ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb"),
            )
        )

        def tensor_payload(array):
            return {
                "dtype": str(array.dtype),
                "shape": list(array.shape),
                "values": array.tolist(),
            }

        def image_payload(camera_index: int):
            y = np.arange(image_size, dtype=np.uint16)[:, None]
            x = np.arange(image_size, dtype=np.uint16)[None, :]
            image = np.stack(
                (
                    (x + camera_index * 17) % 256 + np.zeros_like(y),
                    (y + camera_index * 29) % 256 + np.zeros_like(x),
                    (x + y + camera_index * 41) % 256,
                ),
                axis=-1,
            )
            return tensor_payload(image.astype(np.uint8))

        rng = np.random.default_rng(int(extra.get("seed", 0)))
        request_id = f"{case_id}-{int(time.time() * 1000)}"
        payload = {
            "request_id": request_id,
            "model": model_path,
            "input": {
                "task": sampling_params.prompt or "pick up the blue block",
                "observation": {
                    "images": {
                        camera: image_payload(index)
                        for index, camera in enumerate(camera_order)
                    },
                    "camera_order": list(camera_order),
                    "state": tensor_payload(
                        np.linspace(-0.5, 0.5, state_dim, dtype=np.float32)
                    ),
                    "noise": tensor_payload(
                        rng.standard_normal((action_horizon, action_dim)).astype(
                            np.float32
                        )
                    ),
                },
            },
            "parameters": {
                "action_horizon": action_horizon,
                "action_dim": action_dim,
                "num_inference_steps": int(extra.get("num_inference_steps", 2)),
            },
            "runtime": {
                "return_timing": True,
                "prefix_cache": bool(extra.get("enable_prefix_cache", False)),
                "cuda_graph": bool(extra.get("enable_cuda_graph", True)),
                "output_format": "list",
            },
        }

        base_url = str(client.base_url).rstrip("/")
        endpoint = (
            f"{base_url}/actions/generations"
            if base_url.endswith("/v1")
            else f"{base_url}/v1/actions/generations"
        )
        response = http_requests.post(endpoint, json=payload, timeout=600)
        if response.status_code != 200:
            pytest.fail(f"{case_id}: action generation failed: {response.text}")

        body = response.json()
        action = body["data"][0]["action"]
        if action["shape"] != [action_horizon, action_dim]:
            pytest.fail(
                f"{case_id}: action shape mismatch: {action['shape']} "
                f"!= {[action_horizon, action_dim]}"
            )
        values = action["values"]
        if not all(
            isinstance(value, (int, float)) and np.isfinite(value)
            for row in values
            for value in row
        ):
            pytest.fail(f"{case_id}: action response contains non-finite values")
        return body["id"], response.content

    if modality == "3d":
        fn = generate_mesh
    elif modality == "action":
        fn = generate_action
    elif modality == "video":
        if sampling_params.realtime_num_chunks is not None:
            fn = generate_realtime_video
        elif sampling_params.image_path and sampling_params.prompt:
            if getattr(sampling_params, "direct_url_test", False):
                fn = generate_text_url_image_to_video
            else:
                fn = generate_text_image_to_video
        elif sampling_params.image_path:
            fn = generate_image_to_video
        else:
            fn = generate_video
    elif sampling_params.prompt and sampling_params.image_path:
        if getattr(sampling_params, "direct_url_test", False):
            fn = generate_image_edit_url
        else:
            fn = generate_image_edit
    else:
        fn = generate_image

    return fn