#!/usr/bin/env python3
"""Render OpenAI-style requests with the official Ling-3.0 chat template.

The output is the golden file for tests/chat_tools_test.cpp: each request is
rendered by transformers' own Jinja renderer, so the C++ prompt builder is
checked against the template rather than against a reading of it. The
template is inclusionAI/Ling-3.0-tiny@e3a47d5b986e7141b6efd62597d598ebb392060d
chat_template.jinja (pass its local path). Tool-call arguments travel as JSON
strings in the requests, as OpenAI clients send them; the template needs them
as mappings, as vLLM and llama.cpp convert them.
"""
import argparse
import copy
import json
from pathlib import Path

from transformers.utils.chat_template_utils import render_jinja_template

MODEL = "mindnano-ling3-tiny"
READ = {"type": "function", "function": {"name": "read_file", "description": "Read a file's contents.",
        "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}}}
PATCH = {"type": "function", "function": {"name": "patch", "description": "Apply a small text patch to a file.",
         "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "find": {"type": "string"},
                        "replace": {"type": "string"}}, "required": ["path", "find", "replace"]}}}
RUN = {"type": "function", "function": {"name": "run", "description": "运行命令 — 可设超时。",
       "parameters": {"type": "object", "properties": {"command": {"type": "string"}, "timeout": {"type": "integer"},
                      "env": {"type": "object"}, "dry": {"type": "boolean"}}, "required": ["command"]}}}


def call(id_, name, **arguments):
    return {"id": id_, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments, ensure_ascii=False)}}


CASES = [
    ("no_tools_history", False, None, [
        {"role": "system", "content": "请用中文回答"}, {"role": "user", "content": "你好"},
        {"role": "assistant", "content": "你好！"}, {"role": "user", "content": "继续"}]),
    ("tools_no_system", False, [READ], [{"role": "user", "content": "Reply with OK."}]),
    ("tools_system", False, [READ, PATCH], [
        {"role": "system", "content": "You are terse."},
        {"role": "user", "content": "Use the tools to read /tmp/scratch.txt, then patch it to change 'foo' to 'bar'."}]),
    ("tools_system_thinking_phrase", False, [READ], [
        {"role": "system", "content": "detailed thinking off\nBe brief."}, {"role": "user", "content": "Read /tmp/a."}]),
    ("tool_loop", False, [READ, PATCH, RUN], [
        {"role": "system", "content": "You are terse."}, {"role": "user", "content": "Fix /tmp/a.txt."},
        {"role": "assistant", "content": None, "tool_calls": [
            call("call_1", "read_file", path="/tmp/a.txt"),
            call("call_2", "run", command="ls -l", timeout=30, env={"A": "1", "B": [1, 2.5]}, dry=False)]},
        {"role": "tool", "tool_call_id": "call_1", "content": "foo\n"},
        {"role": "tool", "tool_call_id": "call_2", "content": "total 0"},
        {"role": "assistant", "content": "Read it.", "tool_calls": [
            call("call_3", "patch", path="/tmp/a.txt", find="foo", replace="bar\n")]},
        {"role": "tool", "tool_call_id": "call_3", "content": "ok"}]),
    ("thinking_tool_loop", True, [READ], [
        {"role": "user", "content": "What is in /tmp/a?"},
        {"role": "assistant", "content": "", "reasoning_content": "\nI should read it.\n",
         "tool_calls": [call("call_1", "read_file", path="/tmp/a")]},
        {"role": "tool", "tool_call_id": "call_1", "content": "foo"}]),
    ("think_in_content", True, [READ], [
        {"role": "user", "content": "Hi"},
        {"role": "assistant", "content": "<think>\nplan\n</think>\n\nHello."},
        {"role": "user", "content": "读取 /tmp/文件 ✓"}]),
    ("unicode_call", False, [RUN], [
        {"role": "user", "content": "打招呼"},
        {"role": "assistant", "content": None, "tool_calls": [call("call_1", "run", command="echo 你好")]},
        {"role": "tool", "tool_call_id": "call_1", "content": "你好"}]),
]


def template_messages(messages):
    messages = copy.deepcopy(messages)
    for message in messages:
        for tool_call in message.get("tool_calls", []):
            tool_call["function"]["arguments"] = json.loads(tool_call["function"]["arguments"])
    return messages


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--template", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    template = args.template.read_text()
    golden = []
    for name, thinking, tools, messages in CASES:
        rendered, _ = render_jinja_template(conversations=[template_messages(messages)], tools=tools,
                                            chat_template=template, add_generation_prompt=True,
                                            enable_thinking=thinking)
        request = {"model": MODEL, "messages": messages, "enable_thinking": thinking}
        if tools: request["tools"] = tools
        golden.append({"name": name, "request": request, "prompt": rendered[0]})
    args.output.write_text(json.dumps(golden, ensure_ascii=False, indent=1) + "\n")


if __name__ == "__main__":
    main()