#!/usr/bin/python3
# -*- coding: utf-8 -*-
import json

import gradio as gr
from transformers import AutoTokenizer


def run_chat_template(conversation: str, model_name: str, add_generation_prompt: bool = False):
    conversation = json.loads(conversation)

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    result = tokenizer.apply_chat_template(
        conversation,
        # tools=None,
        tokenize=False,
        add_generation_prompt=add_generation_prompt,
        # enable_thinking=True,
    )
    return result


def get_chat_template_tab():
    with gr.TabItem("chat_template"):
        model_name_choices = [
            "Qwen/Qwen3-8B",
            "meta-llama/Llama-3.1-8B",
            "meta-llama/Llama-3.1-8B-Instruct",
            "openai/gpt-oss-20b",
            "jingyaogong/MiniMind2",

        ]
        ct_model_name = gr.Dropdown(choices=model_name_choices, value=model_name_choices[0], label="model_name", allow_custom_value=True)
        ct_conversation = gr.Textbox(label="conversation")
        ct_add_generation_prompt = gr.Checkbox(label="add_generation_prompt")
        ct_tokenize = gr.Button("tokenize", variant="primary")
        ct_output = gr.Textbox(label="output", max_lines=100)

        ct_tokenize.click(
            run_chat_template,
            inputs=[ct_conversation, ct_model_name, ct_add_generation_prompt],
            outputs=[ct_output],
        )

        gr.Examples(
            examples=[
                [
                    json.dumps([{"role": "user", "content": "帮我识别出文本中的关键词：\n凉山彝族社会中的\"尔普\"(份子钱)是一种礼物交换形式.对\"尔普\"的研究和分析,可有助于人们理解凉山彝族社会.\"尔普\"本来是维系彝族传统社会宗族内部亲属组织的纽带,由于文化变迁的原因,后来发展出了跨宗族的\"尔普\"新形式,又由于族群互动的原因,还产生了跨越族群的\"尔普\"形式.\"尔普\"形式的变迁是族群互动下的一种文化变迁形式,其动力来源于彝、汉两族的互动关系.彝族社会中\"尔普\"的变迁形式是人类学关于族群互动下的文化变迁理论的鲜活事例."}, {"role": "assistant", "content": "彝族;尔普;礼物交换;族群互动"}], ensure_ascii=False, indent=2),
                    "Qwen/Qwen3-8B",
                    True,
                ],
                [
                    json.dumps([{"role": "system", "content": "你是一个有帮助的助手，可以调用工具来获取信息。可用工具：[工具描述JSON]"}, {"role": "user", "content": "查询北京和上海的天气对比"}, {"role": "assistant", "content": "用户想比较两地的天气，我需要分别查询。", "tool_calls": [{"type": "function", "function": {"name": "get_weather", "arguments": "{\"location\": \"北京\"}"}}, {"type": "function", "function": {"name": "get_weather", "arguments": "{\"location\": \"上海\"}"}}]}, {"role": "tool", "content": "{\"weather\": \"sunny\", \"temperature\": 22}"}, {"role": "tool", "content": "{\"weather\": \"cloudy\", \"temperature\": 24}"}, {"role": "assistant", "content": "北京：晴天，22°C；上海：多云，24°C。两地温差不大，上海稍暖一些。"}], ensure_ascii=False, indent=2),
                    "Qwen/Qwen3-8B",
                    True,
                ],
                [
                    json.dumps([{"role": "user", "content": "请先思考再回答：太阳为什么从东边升起？"}, {"role": "assistant", "content": "<think>用户问的是一个天文物理现象。我需要从地球自转的方向来解释。地球绕着地轴自西向东自转。从地球上的观察者视角来看，当地球转向太阳时，太阳就好像从东方的地平线出现。‘东方’的定义本身就与地球自转方向有关。所以，根本原因是地球的自转方向决定了太阳视运动的方向。</think>太阳从东边升起是由于地球的自转方向是自西向东。"}], ensure_ascii=False, indent=2),
                    "Qwen/Qwen3-8B",
                    True,
                ],
                [
                    json.dumps([{"role": "system", "content": "你是一个智能助理，可以调用计算工具。可用工具：\n\n- 工具名：`calculate`\n  描述：计算一个数学表达式的结果。\n  参数：`expression` (字符串): 要计算的表达式，如'3 + 5 * 2'。"}, {"role": "user", "content": "请帮我计算一下三加五乘以二等于多少。"}, {"role": "assistant", "content": "", "tool_calls": [{"function": {"name": "calculate", "arguments": "{\"expression\": \"3 + 5 * 2\"}"}}]}], ensure_ascii=False, indent=2),
                    "Qwen/Qwen3-8B",
                    True,
                ]
            ],
            inputs=[ct_conversation, ct_model_name, ct_add_generation_prompt],
            outputs=[ct_output],
            fn=run_chat_template,
        )

    return locals()


if __name__ == "__main__":
    pass