# SPDX-License-Identifier: Apache-2.0
"""Shared prompt-template construction for HunyuanImage-3.0-Instruct.

Single source of truth for the AR-prefill prompt format used by the
example scripts and any downstream caller that needs to build
HunyuanImage3 chat-template token sequences without invoking the full
diffusion pipeline tokenizer wrapper.

The DiT pipeline (`pipeline_hunyuan_image3.py`) builds prompts through
`TokenizerWrapper.apply_chat_template`, which eagerly consumes
`JointImageInfo` objects produced by image preprocessing. The example
flow uses an `<img>` placeholder + `multi_modal_data` instead, so it
needs a lighter-weight builder that only requires a HF tokenizer. This
module provides that builder; the (task, bot_task) -> template mapping
below is the canonical mapping for both flows.

Two orthogonal axes:

  * `task` selects the I/O modality combination, which only controls
    whether `<img>` placeholders are emitted between `User: ` and the
    user prompt: ``i2t`` / ``it2i`` produce them, ``t2t`` / ``t2i`` do
    not.

  * `bot_task` selects the prompting mode and drives both the system
    prompt and the trigger tag appended after ``Assistant: ``. ``None``
    (default) gives a plain Assistant turn under the unified prompt;
    ``think`` / ``recaption`` switch the trigger tag to ``<think>`` /
    ``<recaption>``; ``think_recaption`` swaps the system prompt for
    the dedicated combined-mode template; ``vanilla`` drops the chat
    structure entirely (pretrain template, ``t2i`` only).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from .system_prompt import get_system_prompt

# HunyuanImage-3.0-Instruct special token ids from tokenizer.json.
# Keep offline AR prompt/stop-token behavior independent of runtime
# tokenizer lookup for these fixed control tokens.
HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS: dict[str, int] = {
    "<|endoftext|>": 127957,
    "<|startoftext|>": 127958,
    "<boi>": 128000,
    "<eoi>": 128001,
    "<img>": 128006,
    "<cfg>": 128010,
    "<recaption>": 128018,
    "</recaption>": 128019,
    "<think>": 128023,
    "</think>": 128024,
    "<answer>": 128025,
    "</answer>": 128026,
    "<img_size_1024>": 128037,
    "<img_ratio_0>": 128044,
    "<img_ratio_32>": 128076,
    "<img_ratio_33>": 130103,
    "<img_ratio_36>": 130106,
}


def validate_special_token_ids(tokenizer: Any) -> None:
    """Fail fast if a real tokenizer's ids drift from the hardcoded table.

    ``HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS`` exists so stop-token / trigger-tag
    resolution doesn't need a tokenizer round-trip, but that only holds as
    long as the ids match ``tokenizer.json`` for whatever model repo/revision
    is actually loaded. Call this once a real tokenizer is available (the
    example scripts do, right after loading it) to catch drift immediately
    instead of silently producing a request that runs to completion with the
    wrong stop tokens. Not wired into ``build_prompt_tokens`` itself: callers
    that intentionally pass a stand-in tokenizer with different ids (e.g.
    this module's own unit tests) must not be forced through this check.

    Both ``tencent/HunyuanImage-3.0`` (base) and ``-Instruct`` report the
    same ``model_class_name`` (``HunyuanImage3ForCausalMM``), so this hook
    fires for either checkpoint -- but the base tokenizer's vocab doesn't
    include the Instruct-only ``<img_ratio_33>``-``<img_ratio_36>`` tokens
    (added for it2i's second aspect-ratio range). A token the tokenizer
    doesn't recognize at all (``convert_tokens_to_ids`` returning
    ``unk_token_id``, or ``None`` for tokenizers without an unk concept) is
    "not part of this checkpoint's vocab," not "drifted"; only an id that
    resolves to something *else* indicates the hardcoded table is stale.
    """
    unk_id = getattr(tokenizer, "unk_token_id", None)
    mismatches: dict[str, tuple[int, int]] = {}
    for tok, expected_id in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.items():
        actual_id = tokenizer.convert_tokens_to_ids(tok)
        if actual_id is None or (unk_id is not None and actual_id == unk_id):
            continue  # not in this checkpoint's vocab -- nothing to compare
        if actual_id != expected_id:
            mismatches[tok] = (expected_id, actual_id)
    if mismatches:
        details = ", ".join(f"{tok!r}: expected {exp}, got {act}" for tok, (exp, act) in mismatches.items())
        raise ValueError(
            "HunyuanImage-3.0 tokenizer special-token ids no longer match "
            f"HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS ({details}). The loaded "
            "tokenizer/model revision has likely changed; update the hardcoded "
            "table in prompt_utils.py to match."
        )


# bot_task -> (sys_type, trigger_tag).
# ``vanilla`` is special-cased downstream: it bypasses the chat template
# (no ``User:`` / ``Assistant:`` framing) and is only valid with
# ``task='t2i'``.
_BOT_TASK_PRESETS: dict[str | None, tuple[str, str | None]] = {
    None: ("en_unified", None),
    "think": ("en_unified", "<think>"),
    "recaption": ("en_unified", "<recaption>"),
    "think_recaption": ("en_think_recaption", "<think>"),
    "vanilla": ("en_vanilla", None),
}

_TASKS: frozenset[str] = frozenset({"t2t", "i2t", "it2i", "t2i"})


class _DefaultBotTask:
    pass


_DEFAULT_BOT_TASK = _DefaultBotTask()

# Legacy composite task alias -> (task, bot_task). Keep this during rebase so
# older callers and intermediate commits still resolve cleanly.
_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
    "t2t": ("en_unified", None, None),
    "i2t": ("en_unified", None, None),
    "it2i_think": ("en_unified", "think", "<think>"),
    "it2i_recaption": ("en_unified", "recaption", "<recaption>"),
    "it2i_think_recaption": ("en_unified", "think_recaption", "<think>"),
    "t2i": ("en_unified", None, None),
    "t2i_vanilla": ("en_vanilla", "vanilla", None),
    "t2i_think": ("en_unified", "think", "<think>"),
    "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
}

_LEGACY_COMPOSITE_TASKS: frozenset[str] = frozenset(_TASK_PRESETS) - {"t2t", "i2t", "t2i"}


def _normalize_task_and_bot_task(
    task: str,
    bot_task: str | None | _DefaultBotTask,
) -> tuple[str, str | None]:
    bot_task_was_omitted = bot_task is _DEFAULT_BOT_TASK
    if task in _TASK_PRESETS:
        _, legacy_bot_task, _ = _TASK_PRESETS[task]
        base_task = task.split("_", 1)[0]
        if base_task == "t2i" and task == "t2i":
            base_task = "t2i"
        if task in ("t2t", "i2t", "t2i"):
            base_task = task
        if bot_task_was_omitted:
            bot_task = legacy_bot_task
        elif task in _LEGACY_COMPOSITE_TASKS and bot_task is None:
            # Composite task names already encode the legacy bot_task. Keep
            # calls like build_prompt_tokens(task="it2i_think", bot_task=None)
            # on their historical meaning; explicit None is the plain-mode
            # escape hatch only for the new two-axis base tasks.
            bot_task = legacy_bot_task
        task = base_task
    elif bot_task_was_omitted:
        bot_task = "think"
    return task, bot_task


def available_tasks() -> list[str]:
    """Sorted list of `task` values accepted by the prompt builders."""
    return sorted(_TASKS)


def available_bot_tasks() -> list[str | None]:
    """Sorted list of `bot_task` values (with ``None`` first)."""
    rest = sorted(k for k in _BOT_TASK_PRESETS if k is not None)
    return [None, *rest]


def resolve_sys_type(bot_task: str | None) -> str:
    """Default system-prompt type for a given ``bot_task``."""
    if bot_task not in _BOT_TASK_PRESETS:
        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
    return _BOT_TASK_PRESETS[bot_task][0]


def resolve_stop_token_ids(
    task: str = "it2i",
    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
    tokenizer: Any | None = None,
    image_size: str | None = None,
) -> list[int]:
    """AR stop-token ids for a given (task, bot_task) generation request.

    Image-output tasks (``it2i`` / ``t2i``) stop on any ``<img_ratio_*>``
    token. Upstream ``modeling_hunyuan_image_3.py::generate_image``
    (line 3289-3303) sets ``final_stop_tokens`` to the full ratio token
    range when ``need_ratio`` is true, then strips the trailing ratio
    token before passing the cot to the image stage. AR's natural
    trajectory under ``_stage_transitions`` is
    ``</recaption><answer><boi><img_size_base><img_ratio_X>``; stopping
    AT the ratio token means KV ends exactly at the prefix DiT reuses,
    and ``ar2diffusion`` can read the ratio off the last sampled token
    without AR wasting decode steps on ``<|endoftext|>``.

    Text-output tasks (``i2t`` / ``t2t``) stop on ``<answer>`` -- the AR
    is the final stage, and the comprehension response sits inside the
    ``<answer>`` body so the answer-open is the natural cot/recaption
    terminator.
    """
    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
    if task not in _TASKS:
        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
    if bot_task not in _BOT_TASK_PRESETS:
        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
    if task in ("it2i", "t2i"):
        need_ratio = image_size is None or image_size == "auto"
        if not need_ratio:
            end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
            end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
            if bot_task in ("think_recaption", "recaption"):
                return [end_of_recaption_id]
            if bot_task == "think":
                return [end_of_think_id, end_of_recaption_id]
        # Main ratio range: <img_ratio_0> .. <img_ratio_32>.
        start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
        end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
        stops = list(range(start, end + 1))
        # Other slices (upstream tokenizer ``ratio_token_other_slices``):
        # <img_ratio_33> .. <img_ratio_36>.
        other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
        other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
        stops.extend(range(other_start, other_end + 1))
        return stops
    return [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]]


# Upstream "Multi-Image Fusion" caps reference images at 3 per request.
MAX_IMAGES_PER_REQUEST = 3


def _validate_num_images(num_images: int) -> None:
    if not (1 <= num_images <= MAX_IMAGES_PER_REQUEST):
        raise ValueError(f"num_images must be in [1, {MAX_IMAGES_PER_REQUEST}], got {num_images}")


def _resolve_preset(task: str, bot_task: str | None) -> tuple[str, str | None]:
    """Validate (task, bot_task) and return ``(sys_type, trigger_tag)``."""
    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
    if task not in _TASKS:
        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
    if bot_task not in _BOT_TASK_PRESETS:
        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
    if bot_task == "vanilla" and task != "t2i":
        raise ValueError(f"bot_task='vanilla' is only valid with task='t2i' (pretrain template); got task={task!r}")
    return _BOT_TASK_PRESETS[bot_task]


def build_prompt(
    user_prompt: str,
    task: str = "it2i",
    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
    sys_type: str | None = None,
    custom_system_prompt: str | None = None,
    num_images: int = 1,
) -> str:
    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path)."""
    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
    preset_sys_type, trigger_tag = _resolve_preset(task, bot_task)
    effective_sys_type = sys_type or preset_sys_type

    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
    sys_text = system_prompt.strip() if system_prompt is not None else ""

    has_image_input = task in ("i2t", "it2i")
    if has_image_input:
        _validate_num_images(num_images)

    if bot_task == "vanilla":
        parts = ["<|startoftext|>"]
        if sys_text:
            parts.append(sys_text)
        parts.append(user_prompt)
        return "".join(parts)

    parts = ["<|startoftext|>"]
    if sys_text:
        parts.append(f"{sys_text}\n\n")
    parts.append("User: ")
    if has_image_input:
        parts.extend(["<img>"] * num_images)
    parts.append(user_prompt)
    parts.append("\n\nAssistant: ")
    if trigger_tag:
        parts.append(trigger_tag)
    return "".join(parts)


@dataclass
class PromptTokensResult:
    token_ids: list[int]
    system_prompt_type: str


def build_prompt_tokens(
    user_prompt: str,
    tokenizer,
    task: str = "it2i",
    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
    sys_type: str | None = None,
    custom_system_prompt: str | None = None,
    num_images: int = 1,
) -> PromptTokensResult:
    """Segment-by-segment tokenization that matches HF apply_chat_template."""
    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
    preset_sys_type, trigger_tag = _resolve_preset(task, bot_task)
    effective_sys_type = sys_type or preset_sys_type

    bos_id = tokenizer.convert_tokens_to_ids("<|startoftext|>")
    img_id = tokenizer.convert_tokens_to_ids("<img>")
    trig_id = tokenizer.convert_tokens_to_ids(trigger_tag) if trigger_tag else None

    has_image_input = task in ("i2t", "it2i")
    if has_image_input:
        _validate_num_images(num_images)

    if bot_task == "vanilla":
        s = build_prompt(
            user_prompt,
            task=task,
            bot_task=bot_task,
            sys_type=sys_type,
            custom_system_prompt=custom_system_prompt,
        )
        token_ids = tokenizer.encode(s, add_special_tokens=False)
        return PromptTokensResult(
            token_ids=token_ids,
            system_prompt_type=effective_sys_type,
        )

    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
    sys_text = system_prompt.strip() if system_prompt is not None else ""

    ids: list[int] = [bos_id]
    if sys_text:
        ids += tokenizer.encode(sys_text, add_special_tokens=False)
        ids += tokenizer.encode("\n\n", add_special_tokens=False)
    ids += tokenizer.encode("User: ", add_special_tokens=False)
    if has_image_input:
        ids += [img_id] * num_images
    ids += tokenizer.encode(user_prompt, add_special_tokens=False)
    ids += tokenizer.encode("\n\nAssistant: ", add_special_tokens=False)
    if trig_id is not None:
        ids += [trig_id]

    return PromptTokensResult(
        token_ids=ids,
        system_prompt_type=effective_sys_type,
    )


@dataclass
class ARPromptInputs:
    """Fully-resolved AR-stage inputs for one HunyuanImage3 request.

    Consolidates the three things every caller (offline example, OpenAI
    server) must compute before submitting an AR request: the AR prefill
    (as segment-tokenized ``prompt_token_ids`` when a tokenizer is
    available, else a formatted ``prompt`` string), the resolved DiT
    system-prompt type, and the AR ``stop_token_ids``.
    """

    prompt: str | None
    prompt_token_ids: list[int] | None
    system_prompt_type: str
    stop_token_ids: list[int]


def build_ar_prompt_inputs(
    user_prompt: str,
    task: str = "t2i",
    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
    sys_type: str | None = None,
    custom_system_prompt: str | None = None,
    num_images: int = 1,
    tokenizer: Any | None = None,
    image_size: str | None = None,
) -> ARPromptInputs:
    """Build the complete AR-stage inputs for a HunyuanImage3 request.

    Wraps :func:`build_prompt_tokens` (byte-for-byte HF parity when
    ``tokenizer`` is given), the :func:`build_prompt` string fallback (when
    no tokenizer is available), and :func:`resolve_stop_token_ids` behind
    one call, so a caller doesn't have to re-derive the AR prefill/stop-token
    logic -- and, critically, doesn't have to independently resolve
    ``bot_task`` for the prompt-building call vs. the stop-token call, which
    can silently diverge if a caller passes the same nominal ``bot_task``
    value to both but resolves "omitted" vs. "explicit None" differently for
    each (this seam always forwards the identical resolved value to both).

    Used by all four offline shared task examples: ``text_to_image.py`` /
    ``image_edit.py`` (via ``model_extras.hunyuan_image3.build_ar_stage_inputs``)
    and ``x_to_text.py`` (via ``build_x_to_text_prompt``, which also calls
    ``build_ar_stage_inputs``). The OpenAI server's ``serving_chat.py`` does
    not currently call this seam -- it still independently calls
    ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids``, and
    can hit exactly the divergence described above when ``bot_task`` is
    omitted from ``extra_body`` (pre-existing, unrelated to this module).
    ``image_size`` follows the same convention as :func:`resolve_stop_token_ids`
    (``None``/``"auto"`` lets the AR predict the aspect ratio; an explicit
    ``"{w}x{h}"`` makes it stop at the terminator).
    """
    if tokenizer is not None:
        result = build_prompt_tokens(
            user_prompt,
            tokenizer,
            task=task,
            bot_task=bot_task,
            sys_type=sys_type,
            custom_system_prompt=custom_system_prompt,
            num_images=num_images,
        )
        prompt_str: str | None = None
        prompt_token_ids: list[int] | None = result.token_ids
        system_prompt_type = result.system_prompt_type
    else:
        prompt_str = build_prompt(
            user_prompt,
            task=task,
            bot_task=bot_task,
            sys_type=sys_type,
            custom_system_prompt=custom_system_prompt,
            num_images=num_images,
        )
        prompt_token_ids = None
        _, resolved_bot_task = _normalize_task_and_bot_task(task, bot_task)
        system_prompt_type = sys_type or resolve_sys_type(resolved_bot_task)

    stop_token_ids = resolve_stop_token_ids(
        task=task,
        bot_task=bot_task,
        tokenizer=tokenizer,
        image_size=image_size,
    )
    return ARPromptInputs(
        prompt=prompt_str,
        prompt_token_ids=prompt_token_ids,
        system_prompt_type=system_prompt_type,
        stop_token_ids=stop_token_ids,
    )


__all__ = [
    "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS",
    "validate_special_token_ids",
    "MAX_IMAGES_PER_REQUEST",
    "_TASK_PRESETS",
    "ARPromptInputs",
    "available_bot_tasks",
    "available_tasks",
    "build_ar_prompt_inputs",
    "build_prompt",
    "build_prompt_tokens",
    "resolve_stop_token_ids",
    "resolve_sys_type",
]