from unsloth.chat_templates import CHAT_TEMPLATES

UNSLOTH_TEMPLATES = [*list(CHAT_TEMPLATES.keys()), "deepseek_r1", "deepseek-r1"]

deepseek_r1_template = (
    "{% if not add_generation_prompt is defined %}"
    "{% set add_generation_prompt = false %}"
    "{% endif %}"
    "{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_message='') %}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'system' %}"
    "{% set ns.system_message = message['content'] %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{{bos_token}}"
    "{{ns.system_message}}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'user' %}"
    "{%- set ns.is_tool = false -%}{{'<\uff5cUser\uff5c>' + message['content']}}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is none %}"
    "{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}"
    "{%- if not ns.is_first %}"
    "{{'<\uff5cAssistant\uff5c><\uff5ctool▁calls▁begin\uff5c><\uff5ctool▁call▁begin\uff5c>'"
    " + tool['type'] + '<\uff5ctool▁sep\uff5c>' + tool['function']"
    "['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] +"
    " '\\n' + '```' + '<\uff5ctool▁call▁end\uff5c>'}}"
    "{%- set ns.is_first = true -%}{%- else %}{{'\\n' +"
    " '<\uff5ctool▁call▁begin\uff5c>' + tool['type'] + '<\uff5ctool▁sep\uff5c>' + "
    "tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '"
    "<\uff5ctool▁call▁end\uff5c>'}}{{'<\uff5ctool▁calls▁end\uff5c><\uff5cend▁of▁sentence\uff5c>'}}"
    "{%- endif %}{%- endfor %}{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}"
    "{{'<\uff5ctool▁outputs▁end\uff5c>' + message['content'] + '<\uff5cend▁of▁sentence\uff5c>'}}"
    "{%- set ns.is_tool = false -%}"
    "{%- else %}"
    "{% set content = message['content'] %}"
    "{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}"
    "{{'<\uff5cAssistant\uff5c>' + content + '<\uff5cend▁of▁sentence\uff5c>'}}{%- endif %}{%- endif %}"
    "{%- if message['role'] == 'tool' %}"
    "{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}"
    "{{'<\uff5ctool▁outputs▁begin\uff5c><\uff5ctool▁output▁begin\uff5c>' + message['content'] + "
    "'<\uff5ctool▁output▁end\uff5c>'}}"
    "{%- set ns.is_output_first = false %}{%- else %}"
    "{{'\\n<\uff5ctool▁output▁begin\uff5c>' + message['content'] + '<\uff5ctool▁output▁end\uff5c>'}}"
    "{%- endif %}{%- endif %}"
    "{%- endfor -%}{% if ns.is_tool %}{{'<\uff5ctool▁outputs▁end\uff5c>'}}"
    "{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<\uff5cAssistant\uff5c><think>\\n'}}"
    "{% endif %}"
)

DEEPSEEK_R1_CHAT_TEMPLATE = (
    deepseek_r1_template,
    "<\uff5cend▁of▁sentence\uff5c>",
)
