# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Integration coverage for the shared task examples' ``_apply_ar_stage_inputs``.

The unit tests in ``test_model_extras.py`` cover ``get_ar_input_builder`` and
``build_ar_stage_inputs`` in isolation, but nothing previously exercised the
*glue* in ``text_to_image.py`` / ``image_edit.py`` that wires them together:
resolving the model's declared ``ar_input_builder`` via the registry, loading
the AR tokenizer, and writing the result onto ``prompt_dict`` /
``sampling_params_list``. A regression in that glue -- e.g. the registry key
mismatch fixed in this PR (``HunyuanImage3ForCausalMM`` vs
``HunyuanImage3Pipeline``), or the ``trust_remote_code`` opt-in being
silently bypassed -- would not have been caught by the isolated unit tests
above, since they call ``build_ar_stage_inputs`` directly rather than going
through the shared script's own code path.

These tests dynamically load the example scripts by file path (``examples/``
is not an installed package) and call ``_apply_ar_stage_inputs`` the same way
``main()`` does, with a fake tokenizer standing in for a real HF download.
"""

from __future__ import annotations

import importlib.util
import sys
from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import pytest

from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

_REPO_ROOT = Path(__file__).resolve().parents[2]


@pytest.mark.parametrize("with_ar_stage", [False, True])
def test_text_to_image_only_builds_ar_inputs_for_ar_deployments(monkeypatch, tmp_path, with_ar_stage):
    from PIL import Image
    from vllm import SamplingParams

    from vllm_omni.config.omni_config import normalize_and_validate_diffusion_engine_ingress_kwargs

    module = _load_module("_shared_text_to_image_deploy", "examples/offline_inference/text_to_image/text_to_image.py")
    tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module)
    monkeypatch.setattr(module, "get_ar_tokenizer_validator", lambda _: None)
    monkeypatch.setattr(
        module, "current_omni_platform", SimpleNamespace(device_type="cpu", pre_register_and_update=lambda parser: None)
    )
    monkeypatch.setattr(module, "is_nextstep_model", lambda _: False)
    monkeypatch.setattr(module, "get_model_class_name", lambda _: "HunyuanImage3ForCausalMM")
    defaults = [OmniDiffusionSamplingParams()]
    if with_ar_stage:
        defaults.insert(0, SamplingParams())
    captured: dict[str, Any] = {}

    def generate(prompt, *, sampling_params_list):
        captured.update(prompt=prompt, params=sampling_params_list)
        return [SimpleNamespace(images=[Image.new("RGB", (8, 8))])]

    def create_omni(**kwargs):
        # Exercise the real ingress validator so stale example-only options
        # cannot silently pass this test while breaking model initialization.
        normalize_and_validate_diffusion_engine_ingress_kwargs(kwargs, stage_id=0)
        return SimpleNamespace(default_sampling_params_list=defaults, generate=generate)

    monkeypatch.setattr(module, "Omni", create_omni)
    output = tmp_path / "output.png"
    monkeypatch.setattr(
        sys,
        "argv",
        [
            "text_to_image.py",
            "--model",
            "test-hunyuan",
            "--prompt",
            "A dog",
            "--output",
            str(output),
            "--extra-body",
            '{"bot_task": null}' if with_ar_stage else '{"bot_task": "none"}',
            "--trust-remote-code",
        ],
    )
    module.main()

    assert output.is_file()
    assert bool(tokenizer_calls) == with_ar_stage
    assert ("prompt_token_ids" in captured["prompt"]) == with_ar_stage
    if with_ar_stage:
        assert captured["params"][0].stop_token_ids
    else:
        assert captured["prompt"] == {"prompt": "A dog", "modalities": ["image"]}
        assert captured["params"][0].extra_args["bot_task"] == "none"


def test_hunyuan_accuracy_offline_cli_contract(monkeypatch, tmp_path):
    import subprocess

    from PIL import Image

    accuracy = _load_module("_hunyuan_accuracy", "tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py")
    module = _load_module("_hunyuan_accuracy_cli", "examples/offline_inference/text_to_image/text_to_image.py")
    output = tmp_path / "offline.png"

    def run(argv, *, check):
        assert check
        assert Path(argv[1]).is_file()
        monkeypatch.setattr(sys, "argv", argv[1:])
        args = module.parse_args()
        assert args.enable_expert_parallel
        assert args.extra_body == {"bot_task": "none"}
        assert args.use_system_prompt == "en_unified"
        assert args.trust_remote_code
        assert args.seed == accuracy.SEED
        assert args.num_inference_steps == accuracy.NUM_INFERENCE_STEPS
        assert args.output == str(output)
        Image.new("RGB", (8, 8)).save(args.output)

    monkeypatch.setattr(subprocess, "run", run)
    image = accuracy._run_vllm_omni_hunyuan_image3_offline(
        model="test-hunyuan", deploy_config="deploy.yaml", output_path=output
    )
    assert image.size == (8, 8)


def _load_module(name: str, relpath: str):
    path = _REPO_ROOT / relpath
    spec = importlib.util.spec_from_file_location(name, path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


class _FakeHunyuanTokenizer:
    """Mirrors test_model_extras.py's fixture: distinct ids, no network I/O."""

    SPECIAL = {"<|startoftext|>": 1, "<img>": 2, "<think>": 3, "<recaption>": 4}

    def convert_tokens_to_ids(self, tok: str) -> int:
        return self.SPECIAL.get(tok, 0)

    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
        return [500 + len(text)]


@dataclass
class _FakeARSamplingParams:
    """Stands in for an LLM (non-diffusion) stage's sampling params."""

    stop_token_ids: list[int] | None = None


def _patch_tokenizer_loader(monkeypatch: pytest.MonkeyPatch, module: Any) -> list[dict[str, Any]]:
    """Patch the ``AutoTokenizer`` the module imports locally, recording calls."""
    calls: list[dict[str, Any]] = []

    class _FakeAutoTokenizer:
        @staticmethod
        def from_pretrained(model: str, trust_remote_code: bool = False) -> _FakeHunyuanTokenizer:
            calls.append({"model": model, "trust_remote_code": trust_remote_code})
            return _FakeHunyuanTokenizer()

    monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", _FakeAutoTokenizer.from_pretrained)
    return calls


def _patch_failing_tokenizer_loader(monkeypatch: pytest.MonkeyPatch) -> None:
    """Patch AutoTokenizer.from_pretrained to always raise, simulating a
    missing --trust-remote-code / network / cache failure."""

    class _FailingAutoTokenizer:
        @staticmethod
        def from_pretrained(model: str, trust_remote_code: bool = False) -> Any:
            raise OSError("simulated tokenizer load failure")

    monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", _FailingAutoTokenizer.from_pretrained)


@pytest.mark.parametrize(
    "module_name,relpath",
    [
        ("_shared_text_to_image", "examples/offline_inference/text_to_image/text_to_image.py"),
        ("_shared_image_edit", "examples/offline_inference/image_to_image/image_edit.py"),
    ],
)
def test_apply_ar_stage_inputs_routes_through_real_registry(
    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
) -> None:
    """The shared script's glue must resolve HunyuanImage3's *real* registry
    entry (keyed on the runtime ``model_class_name``) and write its output
    onto the request -- not a stub, not a hardcoded id."""
    from vllm_omni.model_extras import get_ar_input_builder

    module = _load_module(module_name, relpath)
    tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module)

    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
    assert ar_input_builder is not None, (
        "get_ar_input_builder('HunyuanImage3ForCausalMM') returned None -- "
        "if this regresses, the shared script silently skips AR-stage "
        "input building entirely (this is exactly how the registry-key bug "
        "this PR fixes went undetected)."
    )

    prompt_dict: dict[str, Any] = {}
    ar_params = _FakeARSamplingParams()
    diffusion_params = OmniDiffusionSamplingParams()
    sampling_params_list: list[Any] = [ar_params, diffusion_params]

    module._apply_ar_stage_inputs(
        ar_input_builder,
        model="tencent/HunyuanImage-3.0-Instruct",
        prompt_text="a red panda",
        extra_body={"bot_task": "think", "use_system_prompt": "en_recaption"},
        num_images=0,
        height=1024,
        width=1024,
        prompt_dict=prompt_dict,
        sampling_params_list=sampling_params_list,
        trust_remote_code=False,
    )

    # The tokenizer was loaded exactly once, for the model the user asked for.
    assert len(tokenizer_calls) == 1
    assert tokenizer_calls[0]["model"] == "tencent/HunyuanImage-3.0-Instruct"

    # Byte-for-byte tokenizer path: prompt_token_ids populated, not the string form.
    assert prompt_dict.get("prompt_token_ids"), "AR prefill token-ids were not written onto the request"
    assert "prompt" not in prompt_dict
    assert prompt_dict.get("use_system_prompt") == "en_recaption"
    assert prompt_dict.get("modalities") == ["image"]

    # stop_token_ids reached the AR-stage params, and only those. The diffusion
    # stage must be left untouched: OmniDiffusionSamplingParams has no
    # stop_token_ids field, so the helper skipping it means the attribute is
    # never created (reading it would raise AttributeError, not return None).
    assert ar_params.stop_token_ids, "AR stage did not receive stop_token_ids"
    assert not hasattr(diffusion_params, "stop_token_ids"), "diffusion stage params must not receive stop_token_ids"


@pytest.mark.parametrize(
    "module_name,relpath,cli_flag",
    [
        ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", False),
        ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", True),
    ],
)
def test_apply_ar_stage_inputs_respects_trust_remote_code_option(
    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str, cli_flag: bool
) -> None:
    """The tokenizer load must use the caller's resolved --trust-remote-code
    value, not silently default to True regardless of user intent."""
    from vllm_omni.model_extras import get_ar_input_builder

    module = _load_module(module_name, relpath)
    tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module)
    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")

    module._apply_ar_stage_inputs(
        ar_input_builder,
        model="tencent/HunyuanImage-3.0-Instruct",
        prompt_text="a red panda",
        extra_body={},
        num_images=0,
        height=1024,
        width=1024,
        prompt_dict={},
        sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
        trust_remote_code=cli_flag,
    )

    assert tokenizer_calls[0]["trust_remote_code"] is cli_flag


@pytest.mark.parametrize(
    "module_name,relpath",
    [
        ("_shared_text_to_image_validate", "examples/offline_inference/text_to_image/text_to_image.py"),
        ("_shared_image_edit_validate", "examples/offline_inference/image_to_image/image_edit.py"),
    ],
)
def test_apply_ar_stage_inputs_runs_declared_tokenizer_validator(
    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
) -> None:
    """get_ar_tokenizer_validator's hook must run on a real tokenizer load,
    and a validation failure must propagate -- not get swallowed by the
    tokenizer-load try/except (which would silently fall back to the string
    prompt instead of surfacing the drift)."""
    from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator

    module = _load_module(module_name, relpath)
    _patch_tokenizer_loader(monkeypatch, module)
    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
    validator = get_ar_tokenizer_validator("HunyuanImage3ForCausalMM")
    assert validator is not None

    # The fake tokenizer's ids intentionally don't match
    # HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, so the real validator must reject it.
    with pytest.raises(ValueError, match="no longer match"):
        module._apply_ar_stage_inputs(
            ar_input_builder,
            model="tencent/HunyuanImage-3.0-Instruct",
            prompt_text="a red panda",
            extra_body={},
            num_images=0,
            height=1024,
            width=1024,
            prompt_dict={},
            sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
            trust_remote_code=False,
            validate_tokenizer=validator,
        )


@pytest.mark.parametrize(
    "module_name,relpath,required_argv",
    [
        ("_shared_text_to_image_cli", "examples/offline_inference/text_to_image/text_to_image.py", []),
        # image_edit.py's parser declares --image/--prompt as required (no defaults);
        # omitting them makes parse_args() raise SystemExit(2) before the assertion
        # below is ever reached.
        (
            "_shared_image_edit_cli",
            "examples/offline_inference/image_to_image/image_edit.py",
            ["--image", "dummy.png", "--prompt", "x"],
        ),
    ],
)
@pytest.mark.parametrize("argv,expected", [([], False), (["--trust-remote-code"], True)])
def test_trust_remote_code_flag_parses_from_real_cli(
    monkeypatch: pytest.MonkeyPatch,
    module_name: str,
    relpath: str,
    required_argv: list[str],
    argv: list[str],
    expected: bool,
) -> None:
    """The other tests in this file call _apply_ar_stage_inputs directly with
    an already-resolved trust_remote_code bool, which never exercises
    argparse itself. Closing that specific gap doesn't need a GPU or a model
    load -- parse_args() only touches sys.argv -- so cover it here instead of
    the (expensive, one-model-load-per-invocation) full CLI subprocess a
    stricter GPU e2e test would need."""
    module = _load_module(module_name, relpath)
    monkeypatch.setattr(sys, "argv", ["prog", *required_argv, *argv])
    args = module.parse_args()
    assert args.trust_remote_code is expected


@pytest.mark.parametrize(
    "module_name,relpath",
    [
        ("_shared_text_to_image_failfast", "examples/offline_inference/text_to_image/text_to_image.py"),
        ("_shared_image_edit_failfast", "examples/offline_inference/image_to_image/image_edit.py"),
    ],
)
def test_apply_ar_stage_inputs_tokenizer_load_failure_raises_by_default(
    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
) -> None:
    """A tokenizer load failure must surface as an error by default, not
    silently degrade to the string-prompt form -- that degraded form can
    complete "successfully" with subtly wrong stop tokens, masking the real
    problem (e.g. a missing --trust-remote-code)."""
    from vllm_omni.model_extras import get_ar_input_builder

    module = _load_module(module_name, relpath)
    _patch_failing_tokenizer_loader(monkeypatch)
    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")

    with pytest.raises(OSError, match="simulated tokenizer load failure"):
        module._apply_ar_stage_inputs(
            ar_input_builder,
            model="tencent/HunyuanImage-3.0-Instruct",
            prompt_text="a red panda",
            extra_body={},
            num_images=0,
            height=1024,
            width=1024,
            prompt_dict={},
            sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
            trust_remote_code=False,
        )


@pytest.mark.parametrize(
    "module_name,relpath",
    [
        ("_shared_text_to_image_fallback", "examples/offline_inference/text_to_image/text_to_image.py"),
        ("_shared_image_edit_fallback", "examples/offline_inference/image_to_image/image_edit.py"),
    ],
)
def test_apply_ar_stage_inputs_allow_tokenizer_fallback_option(
    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
) -> None:
    """With allow_tokenizer_fallback=True (the CLI's --allow-tokenizer-fallback),
    a tokenizer load failure degrades to the string-prompt form instead of
    raising -- for a deliberate offline-compat run, not the default."""
    from vllm_omni.model_extras import get_ar_input_builder

    module = _load_module(module_name, relpath)
    _patch_failing_tokenizer_loader(monkeypatch)
    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")

    prompt_dict: dict[str, Any] = {}
    module._apply_ar_stage_inputs(
        ar_input_builder,
        model="tencent/HunyuanImage-3.0-Instruct",
        prompt_text="a red panda",
        extra_body={"bot_task": "recaption"},
        num_images=0,
        height=1024,
        width=1024,
        prompt_dict=prompt_dict,
        sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
        trust_remote_code=False,
        allow_tokenizer_fallback=True,
    )

    # No tokenizer -> string-prompt form, not the byte-for-byte token-ids form.
    assert "prompt_token_ids" not in prompt_dict
    assert prompt_dict.get("prompt") and "a red panda" in prompt_dict["prompt"]