# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""On-demand trace builder for parser engine testing and benchmarks.

Generates token sequences programmatically from model-agnostic scenario
definitions.  Each model format handler knows how to render scenarios
into the model's output format, tokenize them with correct special token
IDs, and compute expected parse outputs.

Every generated sample is self-validated by replaying it through the
real parser before being returned.
"""

from __future__ import annotations

import functools
import json
from dataclasses import dataclass
from typing import Any

from tests.parser.engine.replay_harness import (
    MockTokenizer,
    Sample,
    assert_parse_output,
    collect_output,
    replay_streaming,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionToolsParam,
)
from vllm.parser.engine.registered_adapters import (
    DeepSeekV4Parser,
    DeepSeekV32Parser,
    Gemma4Parser,
    Glm47MoeParser,
    KimiK2Parser,
    MinimaxM2Parser,
    NemotronV3Parser,
    Qwen3Parser,
    SeedOssParser,
)

# ── Data structures ──────────────────────────────────────────────────


@dataclass
class ToolCallSpec:
    name: str
    arguments: dict[str, Any]


@dataclass
class Scenario:
    id: str
    description: str
    reasoning: str | None = None
    content: str | None = None
    tool_calls: list[ToolCallSpec] | None = None
    after_tool_response: bool = False


# ── Scenarios ────────────────────────────────────────────────────────

_READ_TOOL = ToolCallSpec("read_file", {"path": "/tmp/test.txt"})
_BASH_TOOL = ToolCallSpec(
    "bash", {"command": "hostname", "description": "Get hostname"}
)
_WEATHER_TOOL = ToolCallSpec(
    "get_weather",
    {"city": "Dallas", "state": "TX", "unit": "fahrenheit"},
)
_COMPLEX_TOOL = ToolCallSpec(
    "search",
    {
        "query": "vllm parser",
        "filters": {"language": "python", "min_stars": 100},
        "tags": ["ml", "inference"],
        "limit": 10,
        "verbose": True,
    },
)

SCENARIOS: list[Scenario] = [
    Scenario(
        id="think-then-tool",
        description="Reasoning then single tool call",
        reasoning="Let me check the file.",
        tool_calls=[_READ_TOOL],
    ),
    Scenario(
        id="think-then-parallel-tools",
        description="Reasoning then two parallel tool calls",
        reasoning="I need to run both commands.",
        tool_calls=[_BASH_TOOL, _WEATHER_TOOL],
    ),
    Scenario(
        id="think-then-content",
        description="Reasoning then content response",
        reasoning="Let me think about this carefully.",
        content="The answer is 42.",
    ),
    Scenario(
        id="content-only",
        description="Plain content response without reasoning",
        content="Hello! How can I help you today?",
    ),
    Scenario(
        id="tool-only",
        description="Tool call without reasoning",
        tool_calls=[_READ_TOOL],
    ),
    Scenario(
        id="complex-json-args",
        description="Tool call with nested objects, arrays, numbers, booleans",
        reasoning="This needs a complex query.",
        tool_calls=[_COMPLEX_TOOL],
    ),
    Scenario(
        id="whitespace-before-tool",
        description="Whitespace-only content before tool call",
        content="\n\n",
        tool_calls=[_WEATHER_TOOL],
    ),
    Scenario(
        id="think-content-tool",
        description="Reasoning, content, then tool call",
        reasoning="Let me analyze and then fetch data.",
        content="Checking the weather now.",
        tool_calls=[_WEATHER_TOOL],
    ),
    Scenario(
        id="think-whitespace-tool",
        description="Reasoning, whitespace-only gap, then tool call",
        reasoning="Let me check the file contents.",
        content="\n\n",
        tool_calls=[_READ_TOOL],
    ),
    Scenario(
        id="empty-reasoning-content",
        description="Empty reasoning section followed by content",
        reasoning="",
        content="The epoch timestamp is 1779111346.",
    ),
    Scenario(
        id="tool-after-tool-response",
        description="Tool call immediately after tool response (agentic flow)",
        tool_calls=[_READ_TOOL],
        after_tool_response=True,
    ),
    Scenario(
        id="empty-tool-block",
        description="Empty tool block followed by content (edge case recovery)",
        content="Content after empty tools.",
        tool_calls=[],
    ),
]


# ── Tokenization ─────────────────────────────────────────────────────


def _word_split(text: str) -> list[str]:
    """Split text into word-like tokens, preserving all characters."""
    if not text:
        return []
    parts: list[str] = []
    current = ""
    for ch in text:
        if ch in " \t\n\r" and current and current[-1] not in " \t\n\r":
            parts.append(current)
            current = ch
        else:
            current += ch
    if current:
        parts.append(current)
    return parts


def _tokenize(
    segments: list[tuple[str, bool]],
    vocab: dict[str, int],
    start_id: int = 100,
) -> list[tuple[int, str]]:
    """Build token list from segments.

    Each segment is ``(text, is_special)``.  Special segments use vocab
    IDs; content segments are word-split with sequential IDs.
    """
    tokens: list[tuple[int, str]] = []
    next_id = start_id

    for text, is_special in segments:
        if not text:
            continue
        if is_special:
            tid = vocab.get(text)
            if tid is None:
                raise ValueError(f"Special token {text!r} not in vocab")
            tokens.append((tid, text))
        else:
            for word in _word_split(text):
                tokens.append((next_id, word))
                next_id += 1

    return tokens


# ── Tool definitions ─────────────────────────────────────────────────


def _infer_schema(value: object) -> dict:
    """Infer a JSON Schema from a Python value, recursing into dicts/lists."""
    if isinstance(value, bool):
        return {"type": "boolean"}
    if isinstance(value, int):
        return {"type": "integer"}
    if isinstance(value, float):
        return {"type": "number"}
    if isinstance(value, str):
        return {"type": "string"}
    if isinstance(value, dict):
        return {
            "type": "object",
            "properties": {k: _infer_schema(v) for k, v in value.items()},
        }
    if isinstance(value, list) and value:
        return {"type": "array", "items": _infer_schema(value[0])}
    if isinstance(value, list):
        return {"type": "array"}
    return {}


def _tool_defs(tool_calls: list[ToolCallSpec]) -> list[dict]:
    """Generate OpenAI-style tool definitions from tool call specs."""
    seen: set[str] = set()
    tools: list[dict] = []
    for tc in tool_calls:
        if tc.name in seen:
            continue
        seen.add(tc.name)
        properties = {k: _infer_schema(v) for k, v in tc.arguments.items()}
        tools.append(
            {
                "type": "function",
                "function": {
                    "name": tc.name,
                    "parameters": {
                        "type": "object",
                        "properties": properties,
                    },
                },
            }
        )
    return tools


# ── Format handlers ──────────────────────────────────────────────────


def _expected_tc(scenario: Scenario) -> list[dict] | None:
    if not scenario.tool_calls:
        return None
    return [{"name": tc.name, "arguments": tc.arguments} for tc in scenario.tool_calls]


def _expected_tools(scenario: Scenario) -> list[dict] | None:
    return _tool_defs(scenario.tool_calls) if scenario.tool_calls else None


def _validate_sample(sample: Sample, parser_cls: type, **kwargs) -> None:
    """Replay sample through the real parser and assert correctness."""
    tokenizer = MockTokenizer(vocab=dict(sample.vocab), tokens=sample.tokens)
    parser = parser_cls(tokenizer, sample.tools, **kwargs)
    deltas = replay_streaming(
        parser,
        sample.tokens,
        chunk_size=1,
        tools=sample.tools,
        prompt_token_ids=sample.prompt_token_ids,
    )
    output = collect_output(deltas)
    assert_parse_output(output, sample)


def _validate_tools(
    tools: list[dict] | None,
) -> list[ChatCompletionToolsParam] | None:
    if not tools:
        return None
    return [ChatCompletionToolsParam.model_validate(t) for t in tools]


def _make_sample(
    sample_id: str,
    description: str,
    vocab: dict[str, int],
    segments: list[tuple[str, bool]],
    expected_reasoning: str | None,
    expected_content: str | None,
    expected_tool_calls: list[dict] | None,
    tools: list[dict] | None,
    chat_template_kwargs: dict | None = None,
    prompt_token_ids: list[int] | None = None,
) -> Sample:
    tokens = _tokenize(segments, vocab)
    return Sample(
        id=sample_id,
        description=description,
        source="trace-builder",
        vocab=dict(vocab),
        tokens=tokens,
        expected_reasoning=expected_reasoning,
        expected_content=expected_content,
        expected_tool_calls=expected_tool_calls,
        tools=_validate_tools(tools),
        chat_template_kwargs=chat_template_kwargs,
        prompt_token_ids=prompt_token_ids,
    )


# ── Qwen3 (XML tool format, starts in REASONING) ────────────────────

_QWEN3_VOCAB: dict[str, int] = {
    "<think>": 50,
    "</think>": 51,
    "<tool_call>": 60,
    "</tool_call>": 61,
}


def _qwen3_arg_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return str(value)
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False)


def _qwen3_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
    parts = [f"\n<function={tc.name}>"]
    for key, value in tc.arguments.items():
        parts.append(f"\n<parameter={key}>{_qwen3_arg_value(value)}</parameter>")
    parts.append("\n</function>\n")
    return [
        ("<tool_call>", True),
        ("".join(parts), False),
        ("</tool_call>", True),
    ]


def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append((scenario.reasoning, False))
    if scenario.content is not None or scenario.tool_calls is not None:
        segs.append(("</think>", True))
    if scenario.tool_calls is not None and not scenario.tool_calls:
        segs.append(("<tool_call>", True))
        segs.append(("</tool_call>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls:
        for tc in scenario.tool_calls:
            segs.extend(_qwen3_tool_segments(tc))
    return segs


def _qwen3_expected_content(scenario: Scenario) -> str | None:
    if (
        scenario.content is not None
        and scenario.tool_calls
        and not scenario.content.strip()
    ):
        return ""
    return scenario.content


def _build_qwen3(
    scenario: Scenario,
    name: str = "qwen3",
    parser_cls: type = Qwen3Parser,
    strip_trailing_ws: bool = False,
    validate: bool = True,
) -> Sample:
    expected_reasoning: str | None
    if scenario.reasoning is not None:
        r = scenario.reasoning
        if strip_trailing_ws:
            r = r.rstrip()
        expected_reasoning = r
    else:
        expected_reasoning = ""

    sample = _make_sample(
        sample_id=f"{name}-{scenario.id}",
        description=scenario.description,
        vocab=_QWEN3_VOCAB,
        segments=_qwen3_segments(scenario),
        expected_reasoning=expected_reasoning,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
    )
    if validate:
        _validate_sample(sample, parser_cls)
    return sample


# ── MiniMax M2 (XML invoke format, starts in REASONING) ──────────────

_MINIMAX_M2_VOCAB: dict[str, int] = {
    "<think>": 50,
    "</think>": 51,
    "<minimax:tool_call>": 60,
    "</minimax:tool_call>": 61,
}


def _minimax_m2_arg_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return str(value)
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False)


def _minimax_m2_tool_segments(tool_calls: list[ToolCallSpec]) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = [("<minimax:tool_call>", True)]
    for tc in tool_calls:
        segs.append((f'<invoke name="{tc.name}">', False))
        for key, value in tc.arguments.items():
            segs.append(
                (
                    f'<parameter name="{key}">'
                    f"{_minimax_m2_arg_value(value)}"
                    "</parameter>",
                    False,
                )
            )
        segs.append(("</invoke>", False))
    segs.append(("</minimax:tool_call>", True))
    return segs


def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append((scenario.reasoning, False))
    if scenario.content is not None or scenario.tool_calls is not None:
        segs.append(("</think>", True))
    if scenario.tool_calls is not None and not scenario.tool_calls:
        segs.append(("<minimax:tool_call>", True))
        segs.append(("</minimax:tool_call>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls:
        segs.extend(_minimax_m2_tool_segments(scenario.tool_calls))
    return segs


def _build_minimax_m2(scenario: Scenario, validate: bool = True) -> Sample:
    expected_reasoning: str | None
    if scenario.reasoning is not None:
        expected_reasoning = scenario.reasoning.rstrip()
    else:
        expected_reasoning = ""

    sample = _make_sample(
        sample_id=f"minimax_m2-{scenario.id}",
        description=scenario.description,
        vocab=_MINIMAX_M2_VOCAB,
        segments=_minimax_m2_segments(scenario),
        expected_reasoning=expected_reasoning,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
    )
    if validate:
        _validate_sample(sample, MinimaxM2Parser)
    return sample


# ── Gemma4 (channel reasoning, custom arg format) ────────────────────

_GEMMA4_VOCAB: dict[str, int] = {
    "<|channel>": 50,
    "<channel|>": 51,
    "<|tool_call>": 48,
    "<tool_call|>": 49,
    '<|"|>': 52,
    "<|turn>": 53,
    "<|tool_response>": 54,
}
_GEMMA4_THOUGHT_PREFIX = "thought\n"
_GEMMA4_QUOTE = '<|"|>'


def _gemma4_value_segments(value: Any) -> list[tuple[str, bool]]:
    """Render a value in Gemma4 arg format as segments."""
    if isinstance(value, str):
        return [(_GEMMA4_QUOTE, True), (value, False), (_GEMMA4_QUOTE, True)]
    if isinstance(value, bool):
        return [("true" if value else "false", False)]
    if isinstance(value, (int, float)):
        return [(str(value), False)]
    if isinstance(value, dict):
        segs: list[tuple[str, bool]] = [("{", False)]
        for i, (k, v) in enumerate(value.items()):
            if i > 0:
                segs.append((",", False))
            segs.append((f"{k}:", False))
            segs.extend(_gemma4_value_segments(v))
        segs.append(("}", False))
        return segs
    if isinstance(value, list):
        segs = [("[", False)]
        for i, item in enumerate(value):
            if i > 0:
                segs.append((",", False))
            segs.extend(_gemma4_value_segments(item))
        segs.append(("]", False))
        return segs
    return [(json.dumps(value, ensure_ascii=False), False)]


def _gemma4_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = [
        ("<|tool_call>", True),
        (f"call:{tc.name}", False),
        ("{", False),
    ]
    for i, (key, value) in enumerate(tc.arguments.items()):
        if i > 0:
            segs.append((",", False))
        segs.append((f"{key}:", False))
        segs.extend(_gemma4_value_segments(value))
    segs.append(("}", False))
    segs.append(("<tool_call|>", True))
    return segs


def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append(("<|channel>", True))
        segs.append((_GEMMA4_THOUGHT_PREFIX, False))
        segs.append((scenario.reasoning, False))
        segs.append(("<channel|>", True))
    if scenario.tool_calls is not None and not scenario.tool_calls:
        segs.append(("<|tool_call>", True))
        segs.append(("<tool_call|>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls:
        for tc in scenario.tool_calls:
            segs.extend(_gemma4_tool_segments(tc))
    return segs


def _build_gemma4(scenario: Scenario, validate: bool = True) -> Sample:
    prompt_token_ids = None
    if scenario.after_tool_response:
        prompt_token_ids = [_GEMMA4_VOCAB["<|tool_response>"]]
    sample = _make_sample(
        sample_id=f"gemma4-{scenario.id}",
        description=scenario.description,
        vocab=_GEMMA4_VOCAB,
        segments=_gemma4_segments(scenario),
        expected_reasoning=scenario.reasoning,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
        prompt_token_ids=prompt_token_ids,
    )
    if validate:
        _validate_sample(sample, Gemma4Parser)
    return sample


def _build_nemotron_v3(scenario: Scenario, validate: bool = True) -> Sample:
    return _build_qwen3(
        scenario,
        name="nemotron_v3",
        parser_cls=NemotronV3Parser,
        strip_trailing_ws=True,
        validate=validate,
    )


# ── Seed-OSS (Qwen3 XML grammar with Seed wrapper tokens) ────────────

_SEED_OSS_VOCAB: dict[str, int] = {
    "<seed:think>": 50,
    "</seed:think>": 51,
    "<seed:tool_call>": 60,
    "</seed:tool_call>": 61,
}


def _seed_oss_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
    parts = [f"\n<function={tc.name}>"]
    for key, value in tc.arguments.items():
        parts.append(f"\n<parameter={key}>{_qwen3_arg_value(value)}</parameter>")
    parts.append("\n</function>\n")
    return [
        ("<seed:tool_call>", True),
        ("".join(parts), False),
        ("</seed:tool_call>", True),
    ]


def _seed_oss_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append((scenario.reasoning, False))
    if scenario.content is not None or scenario.tool_calls is not None:
        segs.append(("</seed:think>", True))
    if scenario.tool_calls is not None and not scenario.tool_calls:
        segs.append(("<seed:tool_call>", True))
        segs.append(("</seed:tool_call>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls:
        for tc in scenario.tool_calls:
            segs.extend(_seed_oss_tool_segments(tc))
    return segs


def _build_seed_oss(scenario: Scenario, validate: bool = True) -> Sample:
    sample = _make_sample(
        sample_id=f"seed_oss-{scenario.id}",
        description=scenario.description,
        vocab=_SEED_OSS_VOCAB,
        segments=_seed_oss_segments(scenario),
        expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "",
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
    )
    if validate:
        _validate_sample(sample, SeedOssParser)
    return sample


# ── DeepSeek V4 (DSML tool format) ──────────────────────────────────

_DSML = "｜DSML｜"
_DSV4_VOCAB: dict[str, int] = {
    "<think>": 128821,
    "</think>": 128822,
    f"<{_DSML}tool_calls>": 128823,
    f"</{_DSML}tool_calls>": 128824,
}


def _dsv4_param_text(key: str, value: Any) -> str:
    is_string = isinstance(value, str)
    if is_string:
        val_str = value
    elif isinstance(value, bool):
        val_str = "true" if value else "false"
    elif isinstance(value, (int, float)):
        val_str = str(value)
    else:
        val_str = json.dumps(value, ensure_ascii=False)
    string_attr = "true" if is_string else "false"
    return (
        f'<{_DSML}parameter name="{key}" string="{string_attr}">'
        f"{val_str}</{_DSML}parameter>\n"
    )


def _dsv4_tool_text(tc: ToolCallSpec) -> str:
    parts = [f'<{_DSML}invoke name="{tc.name}">\n']
    for key, value in tc.arguments.items():
        parts.append(_dsv4_param_text(key, value))
    parts.append(f"</{_DSML}invoke>\n")
    return "".join(parts)


def _dsml_tool_segs(
    scenario: Scenario,
    tag: str,
) -> list[tuple[str, bool]]:
    if not scenario.tool_calls:
        return []
    parts = ["\n"]
    for tc in scenario.tool_calls:
        parts.append(_dsv4_tool_text(tc))
    return [
        (f"<{_DSML}{tag}>", True),
        ("".join(parts), False),
        (f"</{_DSML}{tag}>", True),
    ]


def _dsv4_segments(scenario: Scenario, thinking: bool) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []

    if thinking:
        if scenario.reasoning is not None:
            segs.append((scenario.reasoning, False))
        if scenario.content is not None or scenario.tool_calls:
            segs.append(("</think>", True))
    else:
        if scenario.reasoning is not None:
            segs.append(("<think>", True))
            segs.append((scenario.reasoning, False))
            segs.append(("</think>", True))

    if scenario.content is not None:
        segs.append((scenario.content, False))

    segs.extend(_dsml_tool_segs(scenario, "tool_calls"))
    return segs


def _build_deepseek_v4(scenario: Scenario, validate: bool = True) -> Sample:
    thinking = scenario.reasoning is not None
    chat_kwargs = {"thinking": True} if thinking else None

    if thinking:
        expected_reasoning: str | None = scenario.reasoning or ""
    else:
        expected_reasoning = None

    sample = _make_sample(
        sample_id=f"deepseek_v4-{scenario.id}",
        description=scenario.description,
        vocab=_DSV4_VOCAB,
        segments=_dsv4_segments(scenario, thinking),
        expected_reasoning=expected_reasoning,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
        chat_template_kwargs=chat_kwargs,
    )
    if validate:
        kwargs = {}
        if chat_kwargs:
            kwargs["chat_template_kwargs"] = chat_kwargs
        _validate_sample(sample, DeepSeekV4Parser, **kwargs)
    return sample


# ── DeepSeek V3.2 (DSML tool format, no reasoning) ──────────────────

_DSV32_VOCAB: dict[str, int] = {
    f"<{_DSML}function_calls>": 128830,
    f"</{_DSML}function_calls>": 128831,
}


def _dsv32_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []

    if scenario.content is not None:
        segs.append((scenario.content, False))

    segs.extend(_dsml_tool_segs(scenario, "function_calls"))
    return segs


def _build_deepseek_v32(scenario: Scenario, validate: bool = True) -> Sample | None:
    if scenario.reasoning is not None:
        return None

    sample = _make_sample(
        sample_id=f"deepseek_v32-{scenario.id}",
        description=scenario.description,
        vocab=_DSV32_VOCAB,
        segments=_dsv32_segments(scenario),
        expected_reasoning=None,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
    )
    if validate:
        _validate_sample(sample, DeepSeekV32Parser)
    return sample


# ── GLM-4.7 MoE (XML tool format, starts in REASONING) ──────────────

_GLM47_MOE_VOCAB: dict[str, int] = {
    "<think>": 50,
    "</think>": 51,
    "<tool_call>": 60,
    "</tool_call>": 61,
    "<arg_key>": 62,
    "</arg_key>": 63,
    "<arg_value>": 64,
    "</arg_value>": 65,
}


def _glm47_moe_arg_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return str(value)
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False)


def _glm47_moe_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = [
        ("<tool_call>", True),
        (tc.name, False),
    ]
    for key, value in tc.arguments.items():
        segs.extend(
            [
                ("<arg_key>", True),
                (key, False),
                ("</arg_key>", True),
                ("<arg_value>", True),
                (_glm47_moe_arg_value(value), False),
                ("</arg_value>", True),
            ]
        )
    segs.append(("</tool_call>", True))
    return segs


def _glm47_moe_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append((scenario.reasoning, False))
    if scenario.content is not None or scenario.tool_calls:
        segs.append(("</think>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls:
        for tc in scenario.tool_calls:
            segs.extend(_glm47_moe_tool_segments(tc))
    return segs


def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample:
    sample = _make_sample(
        sample_id=f"glm47_moe-{scenario.id}",
        description=scenario.description,
        vocab=_GLM47_MOE_VOCAB,
        segments=_glm47_moe_segments(scenario),
        expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "",
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
    )
    if validate:
        _validate_sample(sample, Glm47MoeParser)
    return sample


# ── Kimi K2 (native tool-call section, starts in REASONING) ──────────

_KIMI_K2_VOCAB: dict[str, int] = {
    "<think>": 50,
    "</think>": 51,
    "<|tool_calls_section_begin|>": 60,
    "<|tool_calls_section_end|>": 61,
    "<|tool_call_begin|>": 62,
    "<|tool_call_end|>": 63,
    "<|tool_call_argument_begin|>": 64,
}


def _kimi_k2_tool_segments(
    tool_calls: list[ToolCallSpec],
) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = [("<|tool_calls_section_begin|>", True)]
    for index, tc in enumerate(tool_calls):
        args = json.dumps(tc.arguments, ensure_ascii=False, separators=(",", ":"))
        segs.extend(
            [
                ("<|tool_call_begin|>", True),
                (f"functions.{tc.name}:{index}\n", False),
                ("<|tool_call_argument_begin|>", True),
                (args, False),
                ("<|tool_call_end|>", True),
            ]
        )
    segs.append(("<|tool_calls_section_end|>", True))
    return segs


def _kimi_k2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
    segs: list[tuple[str, bool]] = []
    if scenario.reasoning is not None:
        segs.append(("<think>", True))
        segs.append((scenario.reasoning, False))
    if scenario.content is not None or scenario.tool_calls is not None:
        segs.append(("</think>", True))
    if scenario.content is not None:
        segs.append((scenario.content, False))
    if scenario.tool_calls is not None:
        segs.extend(_kimi_k2_tool_segments(scenario.tool_calls))
    return segs


def _build_kimi_k2(
    scenario: Scenario,
    validate: bool = True,
    thinking: bool = True,
) -> Sample:
    expected_reasoning = (
        scenario.reasoning.rstrip()
        if (thinking and scenario.reasoning is not None)
        else None
    )
    if thinking and scenario.reasoning is None:
        expected_reasoning = ""

    sample = _make_sample(
        sample_id=f"kimi_k2-{scenario.id}",
        description=scenario.description,
        vocab=_KIMI_K2_VOCAB,
        segments=_kimi_k2_segments(scenario),
        expected_reasoning=expected_reasoning,
        expected_content=_qwen3_expected_content(scenario),
        expected_tool_calls=_expected_tc(scenario),
        tools=_expected_tools(scenario),
        chat_template_kwargs=None if thinking else {"thinking": False},
    )
    if validate:
        _validate_sample(
            sample,
            KimiK2Parser,
            chat_template_kwargs=sample.chat_template_kwargs,
        )
    return sample


_KIMI_K2_SCENARIOS = [
    *SCENARIOS,
    Scenario(
        id="trailing-reasoning-whitespace",
        description="Reasoning trailing whitespace is stripped",
        reasoning="Reasoning with trailing whitespace. \n\t",
        content="Done.",
    ),
]


# ── Registry and public API ──────────────────────────────────────────

_BUILDERS: dict[str, Any] = {
    "deepseek_v32": _build_deepseek_v32,
    "deepseek_v4": _build_deepseek_v4,
    "gemma4": _build_gemma4,
    "minimax_m2": _build_minimax_m2,
    "nemotron_v3": _build_nemotron_v3,
    "seed_oss": _build_seed_oss,
    "glm47_moe": _build_glm47_moe,
    "kimi_k2": _build_kimi_k2,
    "qwen3": _build_qwen3,
}


@functools.cache
def build_samples(model: str) -> tuple[Sample, ...]:
    """Build all scenario samples for a model, self-validated."""
    builder = _BUILDERS[model]
    scenarios = _KIMI_K2_SCENARIOS if model == "kimi_k2" else SCENARIOS
    return tuple(s for s in (builder(sc) for sc in scenarios) if s is not None)


def build_sample(model: str, scenario: Scenario) -> Sample | None:
    """Build a single sample for one model + scenario."""
    return _BUILDERS[model](scenario)


def build_scaling_sample(
    model: str, token_count: int, validate: bool = False
) -> Sample:
    """Build a sample with approximately *token_count* tokens."""
    sentence = "The quick brown fox jumps over the lazy dog. "
    text = sentence * (token_count // 10 + 1)
    scenario = Scenario(
        id=f"scaling-{token_count}",
        description=f"Scaling test with ~{token_count} tokens",
        reasoning=text,
        tool_calls=[_READ_TOOL],
    )
    return _BUILDERS[model](scenario, validate=validate)