from pathlib import Path

from jinja2 import Environment, StrictUndefined


def raise_exception(message: str):
    raise RuntimeError(message)


def main() -> None:
    template = Path("Qwen3-Coder-Next.jinja").read_text()
    env = Environment(undefined=StrictUndefined)
    env.globals.update(raise_exception=raise_exception)
    compiled = env.from_string(template)

    cases = [
        (
            "plain_developer",
            [
                {"role": "developer", "content": "dev rules"},
                {"role": "user", "content": "hello"},
            ],
            [],
        ),
        (
            "typed_user_and_developer",
            [
                {
                    "role": "developer",
                    "content": [{"type": "input_text", "text": "dev typed"}],
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "hello"},
                        {"type": "image_url", "image_url": {"url": "x"}},
                    ],
                },
            ],
            [],
        ),
        (
            "tool_mapping_args_and_typed_tool_response",
            [
                {"role": "user", "content": "run"},
                {
                    "role": "assistant",
                    "content": "",
                    "tool_calls": [
                        {
                            "function": {
                                "name": "bash",
                                "arguments": {"cmd": "ls", "limit": 3},
                            }
                        }
                    ],
                },
                {"role": "tool", "content": [{"type": "text", "text": "ok"}]},
            ],
            [
                {
                    "type": "function",
                    "function": {
                        "name": "bash",
                        "description": "run shell",
                        "parameters": {
                            "type": "object",
                            "properties": {"cmd": {"type": "string"}},
                        },
                    },
                }
            ],
        ),
        (
            "tool_string_args_and_typed_assistant_content",
            [
                {"role": "user", "content": "run"},
                {
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": "I will run it."}],
                    "tool_calls": [
                        {"function": {"name": "bash", "arguments": '{"cmd":"ls"}'}}
                    ],
                },
            ],
            [],
        ),
        (
            "none_content",
            [
                {"role": "developer", "content": None},
                {"role": "user", "content": None},
            ],
            [],
        ),
    ]

    for name, messages, tools in cases:
        rendered = compiled.render(
            messages=messages,
            tools=tools,
            add_generation_prompt=True,
        )
        assert "<|im_start|>assistant\n" in rendered, name
        print(f"CASE {name}: OK len={len(rendered)}")

    print("ALL_OK")


if __name__ == "__main__":
    main()