# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations

import ast
import os
import pathlib

import pytest

from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
    _TASK_PRESETS,
    HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
    available_bot_tasks,
    available_tasks,
    build_prompt,
    build_prompt_tokens,
    resolve_stop_token_ids,
)

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


class FakeTokenizer:
    SPECIAL = {
        "<|startoftext|>": 1,
        "<img>": 2,
        "<think>": 3,
        "<recaption>": 4,
        "<|endoftext|>": 5,
        "</recaption>": 6,
        "</answer>": 7,
        "<boi>": 8,
        "</think>": 9,
        **{f"<img_ratio_{i}>": 1000 + i for i in range(33)},
    }

    def __init__(self) -> None:
        self.encode_calls: list[str] = []
        self.eos_token_id = self.SPECIAL["<|endoftext|>"]

    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]:
        self.encode_calls.append(text)
        return list(range(100, 100 + len(text)))


def test_available_tasks_covers_all_modalities():
    assert set(available_tasks()) == {"t2t", "i2t", "it2i", "t2i"}


def test_available_bot_tasks_covers_all_modes():
    assert set(available_bot_tasks()) == {None, "think", "recaption", "think_recaption", "vanilla"}


def test_legacy_task_presets_still_available():
    assert {
        "it2i_think",
        "it2i_recaption",
        "it2i_think_recaption",
        "t2i_think",
        "t2i_recaption",
        "t2i_vanilla",
    } <= set(_TASK_PRESETS)


def test_legacy_base_task_omitted_bot_task_keeps_plain_mode():
    prompt = build_prompt("HELLO", task="i2t")
    assert prompt.endswith("Assistant: ")
    assert not prompt.endswith("<think>")

    result = build_prompt_tokens("hi", FakeTokenizer(), task="i2t")
    assert result.system_prompt_type == "en_unified"
    assert result.token_ids[-1] not in {
        FakeTokenizer.SPECIAL["<think>"],
        FakeTokenizer.SPECIAL["<recaption>"],
    }


def test_legacy_composite_task_with_none_bot_task_keeps_encoded_mode():
    prompt = build_prompt("HELLO", task="it2i_think", bot_task=None)
    assert prompt.endswith("Assistant: <think>")

    result = build_prompt_tokens("hi", FakeTokenizer(), task="it2i_recaption", bot_task=None)
    assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<recaption>"]


def test_default_prompt_still_uses_it2i_think_mode():
    prompt = build_prompt("HELLO")
    assert prompt.endswith("Assistant: <think>")

    result = build_prompt_tokens("hi", FakeTokenizer())
    assert result.system_prompt_type == "en_unified"
    assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<think>"]


def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range():
    """Image-output tasks stop on any ``<img_ratio_*>`` token.

    Mirrors upstream ``modeling_hunyuan_image_3.py::generate_image``
    (line 3289-3303): when ``need_ratio`` is true,
    ``final_stop_tokens = list(range(start_ratio, end_ratio + 1)) +
    ratio_token_other_slices``. AR stops AT the ratio token sampled
    after ``<img_size_*>``; the bridge then strips the trailing ratio
    token before passing the cot to DiT.
    """
    tok = FakeTokenizer()

    start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
    expected = list(range(start, end + 1)) + list(range(other_start, other_end + 1))

    # Image-output: t2i / it2i stop on the full ratio token range.
    for bot in ("think", "recaption", "think_recaption", "vanilla"):
        assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok) == expected
        assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok) == expected

    # Text-output: i2t / t2t comprehension stops on <answer> (response sits inside).
    answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id]
    assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id]

    # image_size="auto" is equivalent to None (need_ratio=True).
    for bot in ("think", "recaption", "think_recaption", "vanilla"):
        assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected
        assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected


def test_resolve_stop_token_ids_explicit_image_size_stops_on_terminator():
    """When image_size is an explicit "WxH" string, need_ratio=False: AR
    already knows the size and stops at the cot terminator instead of the
    ratio range.

    Mirrors upstream logic where ``need_ratio`` is false when the user
    supplies a target size, so ``final_stop_tokens`` is set to the
    recaption/think terminator (not the ratio tokens).
    """
    tok = FakeTokenizer()

    end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
    end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]

    # think_recaption / recaption -> [end_of_recaption]
    for bot in ("think_recaption", "recaption"):
        for task in ("t2i", "it2i"):
            result = resolve_stop_token_ids(task=task, bot_task=bot, tokenizer=tok, image_size="1024x1024")
            assert result == [end_of_recaption_id], (
                f"task={task}, bot_task={bot}, image_size='1024x1024': expected [{end_of_recaption_id}], got {result}"
            )

    # think -> [end_of_think, end_of_recaption]
    for task in ("t2i", "it2i"):
        result = resolve_stop_token_ids(task=task, bot_task="think", tokenizer=tok, image_size="1024x1024")
        assert result == [end_of_think_id, end_of_recaption_id], (
            f"task={task}, bot_task='think', image_size='1024x1024': "
            f"expected [{end_of_think_id}, {end_of_recaption_id}], got {result}"
        )

    # vanilla -> no early return, falls through to ratio range
    start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
    expected_ratio = list(range(start, end + 1)) + list(range(other_start, other_end + 1))
    for task in ("t2i", "it2i"):
        result = resolve_stop_token_ids(task=task, bot_task="vanilla", tokenizer=tok, image_size="1024x1024")
        assert result == expected_ratio, (
            f"task={task}, bot_task='vanilla', image_size='1024x1024': expected ratio range, got {result}"
        )

    # Text tasks are unaffected by image_size.
    answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok, image_size="1024x1024") == [answer_id]


@pytest.mark.parametrize(
    "task,bot_task",
    [
        ("t2t", None),
        ("i2t", None),
        ("it2i", "think"),
        ("it2i", "recaption"),
        ("it2i", "think_recaption"),
        ("t2i", "think"),
        ("t2i", "recaption"),
        ("t2i", "think_recaption"),
    ],
)
def test_build_prompt_string_structure_chat_template(task: str, bot_task: str | None):
    s = build_prompt("HELLO", task=task, bot_task=bot_task)
    assert s.startswith("<|startoftext|>")
    assert "User: " in s
    assert "Assistant: " in s
    assert s.index("User: ") < s.index("HELLO") < s.index("Assistant: ")

    if task in ("i2t", "it2i"):
        assert s.index("User: ") < s.index("<img>") < s.index("HELLO")
    else:
        assert "<img>" not in s

    if bot_task in ("think", "think_recaption"):
        assert s.endswith("Assistant: <think>")
    elif bot_task == "recaption":
        assert s.endswith("Assistant: <recaption>")
    elif bot_task is None:
        assert s.endswith("Assistant: ")


def test_build_prompt_vanilla_uses_pretrain_template():
    s = build_prompt("HELLO", task="t2i", bot_task="vanilla")
    assert s.startswith("<|startoftext|>")
    assert "User: " not in s
    assert "Assistant: " not in s
    assert s.endswith("HELLO")


def test_build_prompt_vanilla_rejects_non_t2i_task():
    with pytest.raises(ValueError, match="bot_task='vanilla'"):
        build_prompt("x", task="it2i", bot_task="vanilla")
    with pytest.raises(ValueError, match="bot_task='vanilla'"):
        build_prompt_tokens("x", FakeTokenizer(), task="i2t", bot_task="vanilla")


def test_build_prompt_unknown_task_raises():
    with pytest.raises(ValueError, match="Unknown task"):
        build_prompt("x", task="bogus")
    with pytest.raises(ValueError, match="Unknown task"):
        build_prompt_tokens("x", FakeTokenizer(), task="bogus")


def test_build_prompt_unknown_bot_task_raises():
    with pytest.raises(ValueError, match="Unknown bot_task"):
        build_prompt("x", task="t2i", bot_task="bogus")
    with pytest.raises(ValueError, match="Unknown bot_task"):
        build_prompt_tokens("x", FakeTokenizer(), task="t2i", bot_task="bogus")


def test_build_prompt_tokens_segments_each_boundary():
    tok = FakeTokenizer()
    build_prompt_tokens("写诗。", tok, task="i2t", bot_task=None)
    assert "User: " in tok.encode_calls
    assert "写诗。" in tok.encode_calls
    assert "\n\nAssistant: " in tok.encode_calls
    for call in tok.encode_calls:
        if call != "写诗。":
            assert "写诗。" not in call


def test_build_prompt_tokens_image_placeholder_present_for_image_tasks():
    tok = FakeTokenizer()
    result = build_prompt_tokens("hi", tok, task="i2t", bot_task=None)
    ids = result.token_ids
    assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"]
    assert FakeTokenizer.SPECIAL["<img>"] in ids


def test_build_prompt_tokens_no_image_for_text_only_tasks():
    tok = FakeTokenizer()
    result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
    ids = result.token_ids
    assert FakeTokenizer.SPECIAL["<img>"] not in ids


@pytest.mark.parametrize(
    "task,bot_task,trigger_id",
    [
        ("it2i", "think", FakeTokenizer.SPECIAL["<think>"]),
        ("t2i", "think", FakeTokenizer.SPECIAL["<think>"]),
        ("t2i", "think_recaption", FakeTokenizer.SPECIAL["<think>"]),
        ("it2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
        ("t2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
        ("it2i_think", None, FakeTokenizer.SPECIAL["<think>"]),
        ("it2i_recaption", None, FakeTokenizer.SPECIAL["<recaption>"]),
    ],
)
def test_build_prompt_tokens_trigger_is_last_token(task: str, bot_task: str | None, trigger_id: int):
    tok = FakeTokenizer()
    result = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task)
    assert result.token_ids[-1] == trigger_id


def test_build_prompt_tokens_no_trigger_for_plain_tasks():
    tok = FakeTokenizer()
    result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
    assert result.token_ids[-1] not in {
        FakeTokenizer.SPECIAL["<think>"],
        FakeTokenizer.SPECIAL["<recaption>"],
    }


def _repo_root() -> pathlib.Path:
    return pathlib.Path(__file__).resolve().parents[4]


def test_end2end_routes_through_shared_prompt_utils():
    end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
    tree = ast.parse(end2end_path.read_text(encoding="utf-8"))

    local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)}
    assert not (local_func_names & {"build_prompt", "build_prompt_tokens"})

    imported_from_prompt_utils: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"):
            imported_from_prompt_utils.update(alias.name for alias in node.names)
    expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"}
    assert expected_imports <= imported_from_prompt_utils


_HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"


def _hf_cached(model_id: str) -> bool:
    hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")
    snap_dir = os.path.join(hf_home, "hub", f"models--{model_id.replace('/', '--')}", "snapshots")
    return os.path.isdir(snap_dir) and any(os.scandir(snap_dir))


@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
def test_segment_tokenize_diverges_from_full_string_encode():
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
    user_prompt = "写一首关于夜的诗。"
    result = build_prompt_tokens(user_prompt, tok, task="i2t", bot_task=None)
    seg_ids = result.token_ids
    full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False)
    assert seg_ids != full_ids
    assert len(seg_ids) >= len(full_ids)