#!/usr/bin/env python3
"""Generate the PRODUCTION fixed chat_template.jinja for SmolLM3-3B.
Uses a hybrid approach: raw strings for Jinja2, concat for special tokens."""
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B")

THINK_S = tok.decode([128002])
THINK_E = tok.decode([128003])
RESP_S = tok.decode([128013])
RESP_E = tok.decode([128014])
TC_S = tok.decode([128015])
TC_E = tok.decode([128016])

# Build the template as a list of chunks
# Key: any line with special tokens gets its own append with string concat
T = []

# ─── System header (no special tokens) ───
T.append('{# ───── defaults ───── #}\n')
T.append('{%- if enable_thinking is not defined -%}\n')
T.append('{%- set enable_thinking = true -%}\n')
T.append('{%- endif -%}\n\n')
T.append('{# ───── reasoning mode ───── #}\n')
T.append('{%- if enable_thinking -%}\n')
T.append('  {%- set reasoning_mode = "/think" -%}\n')
T.append('{%- else -%}\n')
T.append('  {%- set reasoning_mode = "/no_think" -%}\n')
T.append('{%- endif -%}\n\n')
T.append('{# ───── header (system message) ───── #}\n')
T.append('{{- "<|im_start|>system\\n" -}}\n\n')
T.append('{%- if messages[0].role == "system" -%}\n')
T.append('  {%- set system_message = messages[0].content -%}\n')
T.append('  {%- if "/no_think" in system_message -%}\n')
T.append('    {%- set reasoning_mode = "/no_think" -%}\n')
T.append('  {%- elif "/think" in system_message -%}\n')
T.append('    {%- set reasoning_mode = "/think" -%}\n')
T.append('  {%- endif -%}\n')
T.append('  {%- set custom_instructions = system_message.replace("/no_think", "").replace("/think", "").rstrip() -%}\n')
T.append('{%- endif -%}\n\n')
T.append('{%- if "/system_override" in system_message -%}\n')
T.append('  {{- custom_instructions.replace("/system_override", "").rstrip() -}}\n')
T.append('{%- else -%}\n')
T.append('  {{- "## Metadata\\n\\n" -}}\n')
T.append('  {{- "Knowledge Cutoff Date: June 2025\\n" -}}\n')
T.append('  {%- set today = strftime_now("%d %B %Y") -%}\n')
T.append('  {{- "Today Date: " ~ today ~ "\\n" -}}\n')
T.append('  {{- "Reasoning Mode: " + reasoning_mode + "\\n\\n" -}}\n\n')
T.append('  {{- "## Custom Instructions\\n\\n" -}}\n')
T.append('  {%- if custom_instructions -%}\n')
T.append('    {{- custom_instructions + "\\n\\n" -}}\n')
T.append('  {%- elif reasoning_mode == "/think" -%}\n')
T.append('    {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\\n\\n" -}}\n')
T.append('  {%- else -%}\n')
T.append('    {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\\n\\n" -}}\n')
T.append('  {%- endif -%}\n\n')
T.append('  {%- if xml_tools or python_tools or tools -%}\n')
T.append('    {{- "### Tools\\n\\n" -}}\n')
T.append('    {%- if xml_tools or tools -%}\n')
T.append('      {%- if tools -%}\n')
T.append('        {%- set xml_tools = tools -%}\n')
T.append('      {%- endif -%}\n')
T.append('      {%- set ns = namespace(xml_tool_string="You may call one or more functions to assist with the user query.\\nYou are provided with function signatures within <tools></tools> XML tags:\\n\\n<tools>\\n") -%}\n')
T.append('      {%- for tool in xml_tools[:] -%}\n')
T.append('        {%- set ns.xml_tool_string = ns.xml_tool_string ~ (tool | tojson) ~ "\\n" -%}\n')
T.append('      {%- endfor -%}\n')

# Tool format instruction - has special tokens
T.append('      {%- set xml_tool_string = ns.xml_tool_string + "</tools>\\n\\nFor each function call, return a json object with function name and arguments within ' + TC_S + ' XML tags:\\n' + TC_S + '\\n{\\"name\\": <function-name>, \\"arguments\\": <args-json-object>}\\n' + TC_E + '" -%}\n')

T.append('      {{- xml_tool_string -}}\n')
T.append('    {%- endif -%}\n')
T.append('    {%- if python_tools -%}\n')
T.append('      {%- set ns = namespace(python_tool_string="You may call one or more functions as python tools.\\n<tools>\\n") -%}\n')
T.append('      {%- for tool in python_tools[:] -%}\n')
T.append('        {%- set ns.python_tool_string = ns.python_tool_string ~ (tool | string) ~ "\\n" -%}\n')
T.append('      {%- endfor -%}\n')
T.append('      {%- set python_tool_string = ns.python_tool_string + "</tools>\\n\\nThe state persists between code executions." -%}\n')
T.append('      {{- python_tool_string -}}\n')
T.append('    {%- endif -%}\n')
T.append('    {{- "\\n\\n" -}}\n')
T.append('  {%- endif -%}\n')
T.append('{%- endif -%}\n')
T.append('{{- "<|im_end|>\\n" -}}\n\n')

# ─── Main loop ───
T.append('{# ───── main loop ───── #}\n')
T.append('{%- for message in messages -%}\n')
T.append('    {%- if message.role == "user" -%}\n')
T.append('        {{ "<|im_start|>user\\n" + message.content + "<|im_end|>\\n" }}\n')
T.append('    {%- elif message.role == "assistant" -%}\n')
T.append('        {% generation %}\n')
T.append('        {%- if message.tool_calls -%}\n')
T.append('            {%- set ns = namespace(tc_text="") -%}\n')
T.append('            {%- for tc in message.tool_calls -%}\n')

# FIX: Render tool calls with TC_S/TC_E tokens using ~ (Jinja2 concat)
T.append('                {%- set ns.tc_text = ns.tc_text ~ "' + TC_S + '\\n{\\"name\\": \\"" ~ tc.function.name ~ "\\", \\"arguments\\": " ~ tc.function.arguments ~ "}\\n' + TC_E + '" -%}\n')

T.append('            {%- endfor -%}\n')
T.append('            {{ "<|im_start|>assistant\\n" ~ (message.content if message.content is string else "") ~ ns.tc_text ~ "<|im_end|>\\n" }}\n')
T.append('        {%- else -%}\n')

# FIX: /think = use think tags, /no_think = plain text (was inverted in original)
T.append('            {%- if reasoning_mode == "/think" -%}\n')
T.append('                {{ "<|im_start|>assistant\\n' + THINK_S + '\\n" ~ (message.content if message.content is string else "") ~ "\\n' + THINK_E + '<|im_end|>\\n" }}\n')
T.append('            {%- else -%}\n')
T.append('                {{ "<|im_start|>assistant\\n" ~ (message.content if message.content is string else "") ~ "<|im_end|>\\n" }}\n')
T.append('            {%- endif -%}\n')
T.append('        {%- endif -%}\n')
T.append('        {% endgeneration %}\n')

# FIX: Tool role with RESP_S/RESP_E tokens
T.append('    {%- elif message.role == "tool" -%}\n')
T.append('        {{ "<|im_start|>user\\n' + RESP_S + '\\n" ~ (message.content if message.content is string else "") ~ "\\n' + RESP_E + '<|im_end|>\\n" }}\n')
T.append('    {%- endif -%}\n')
T.append('{%- endfor -%}\n\n')

# ─── Generation prompt ───
T.append('{# ───── generation prompt ───── #}\n')
T.append('{%- if add_generation_prompt -%}\n')
T.append('    {%- if reasoning_mode == "/think" -%}\n')
T.append('        {{ "<|im_start|>assistant\\n' + THINK_S + '\\n" }}\n')
T.append('    {%- else -%}\n')
T.append('        {{ "<|im_start|>assistant\\n" }}\n')
T.append('    {%- endif -%}\n')
T.append('{%- endif -%}\n')

template = ''.join(T)

with open('/root/chat_template.jinja', 'w', encoding='utf-8') as f:
    f.write(template)

print("Production template written to /root/chat_template.jinja")
print(f"Length: {len(template)} bytes")