#!/usr/bin/env python3
"""
Solstice AI: Sync HyperSynapse-SAS v3.0 "Soft Spectrum" across all model families
with explicit support for regular / vanilla chat templates.
"""

import os
import jinja2

BASE_DIR = "/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates"

# Read user-updated v3.0 Qwen and DeepSeek
with open(os.path.join(BASE_DIR, "qwen.jinja"), "r", encoding="utf-8") as f:
    qwen_v3 = f.read()

with open(os.path.join(BASE_DIR, "deepseek.jinja"), "r", encoding="utf-8") as f:
    deepseek_v3 = f.read()

# Enhance Qwen and DeepSeek with regular/vanilla triggers:
def add_regular_triggers(template_code):
    # Add regular to kwarg effort list
    old_kw = "('none', 'off', '0', 'mortal', 'disabled')"
    new_kw = "('none', 'off', '0', 'mortal', 'disabled', 'regular', 'vanilla', 'standard', 'default', 'clean')"
    template_code = template_code.replace(old_kw, new_kw)
    
    # Add regular to mode_raw list
    old_raw = "('0', 'mortal', 'off', 'none', 'direct', 'chat', 'disabled')"
    new_raw = "('0', 'mortal', 'off', 'none', 'direct', 'chat', 'disabled', 'regular', 'vanilla', 'standard', 'default', 'clean')"
    template_code = template_code.replace(old_raw, new_raw)
    
    # Add inline tags [regular], [vanilla], [standard], [default]
    tag_clean_target = "{%- set t = t | replace('[L0]', '')"
    tag_clean_replacement = "{%- set t = t | replace('[regular]', '') | replace('[REGULAR]', '') | replace('[vanilla]', '') | replace('[standard]', '') | replace('[default]', '') | replace('[L0]', '')"
    template_code = template_code.replace(tag_clean_target, tag_clean_replacement)
    
    # Add support for regular=true kwarg
    regular_kwarg_check = """
{%- set is_explicit_regular = (regular is defined and regular) or (mode is defined and mode in ('regular', 'vanilla', 'standard', 'default', 'clean')) %}
{%- if is_explicit_regular %}
    {%- set _kwarg_effort = 'regular' %}
{%- endif %}
"""
    # Insert right after _kwarg_effort definition
    marker = "{%- set _kwarg_effort = "
    idx = template_code.find(marker)
    if idx != -1:
        end_line = template_code.find("\n", idx)
        template_code = template_code[:end_line+1] + regular_kwarg_check + template_code[end_line+1:]
        
    return template_code

deepseek_v3_enhanced = add_regular_triggers(deepseek_v3)
qwen_v3_enhanced = add_regular_triggers(qwen_v3)

with open(os.path.join(BASE_DIR, "deepseek.jinja"), "w", encoding="utf-8") as f:
    f.write(deepseek_v3_enhanced)

with open(os.path.join(BASE_DIR, "qwen.jinja"), "w", encoding="utf-8") as f:
    f.write(qwen_v3_enhanced)

print("Updated deepseek.jinja and qwen.jinja with regular mode support!")

# Universal.jinja is built on top of enhanced qwen.jinja (ChatML standard)
with open(os.path.join(BASE_DIR, "universal.jinja"), "w", encoding="utf-8") as f:
    f.write(qwen_v3_enhanced)
print("Updated universal.jinja!")

# Create standalone pure regular / vanilla templates for users who want zero overhead:
regular_qwen = """{%- if messages[0]['role'] == 'system' -%}
    {{- '<|im_start|>system\\n' + messages[0]['content'] | trim + '<|im_end|>\\n' -}}
    {%- set offset = 1 -%}
{%- else -%}
    {%- set offset = 0 -%}
{%- endif -%}
{%- for message in messages[offset:] -%}
    {{- '<|im_start|>' + message['role'] + '\\n' + message['content'] | trim + '<|im_end|>\\n' -}}
{%- endfor -%}
{%- if add_generation_prompt -%}
    {{- '<|im_start|>assistant\\n' -}}
{%- endif -%}
"""
with open(os.path.join(BASE_DIR, "regular_qwen.jinja"), "w", encoding="utf-8") as f:
    f.write(regular_qwen)

regular_universal = regular_qwen
with open(os.path.join(BASE_DIR, "regular_universal.jinja"), "w", encoding="utf-8") as f:
    f.write(regular_universal)

regular_deepseek = """{{- bos_token -}}
{%- if messages[0]['role'] == 'system' -%}
    {{- messages[0]['content'] | trim -}}
    {%- set offset = 1 -%}
{%- else -%}
    {%- set offset = 0 -%}
{%- endif -%}
{%- for message in messages[offset:] -%}
    {%- if message['role'] == 'user' -%}
        {{- '<｜User｜>' + message['content'] | trim -}}
    {%- elif message['role'] == 'assistant' -%}
        {{- '<｜Assistant｜>' + message['content'] | trim + eos_token -}}
    {%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
    {{- '<｜Assistant｜>' -}}
{%- endif -%}
"""
with open(os.path.join(BASE_DIR, "regular_deepseek.jinja"), "w", encoding="utf-8") as f:
    f.write(regular_deepseek)

regular_llama3 = """{{- bos_token -}}
{%- if messages[0]['role'] == 'system' -%}
    {{- '<|start_header_id|>system<|end_header_id|>\\n\\n' + messages[0]['content'] | trim + '<|eot_id|>' -}}
    {%- set offset = 1 -%}
{%- else -%}
    {%- set offset = 0 -%}
{%- endif -%}
{%- for message in messages[offset:] -%}
    {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + message['content'] | trim + '<|eot_id|>' -}}
{%- endfor -%}
{%- if add_generation_prompt -%}
    {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}
{%- endif -%}
"""
with open(os.path.join(BASE_DIR, "regular_llama3.jinja"), "w", encoding="utf-8") as f:
    f.write(regular_llama3)

regular_glm = """[gMASK]<sop>
{%- for message in messages -%}
    {%- if message['role'] == 'system' -%}
        {{- '<|system|>\\n' + message['content'] | trim -}}
    {%- elif message['role'] == 'user' -%}
        {{- '<|user|>\\n' + message['content'] | trim -}}
    {%- elif message['role'] == 'assistant' -%}
        {{- '<|assistant|>\\n' + message['content'] | trim -}}
    {%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
    {{- '<|assistant|>\\n' -}}
{%- endif -%}
"""
with open(os.path.join(BASE_DIR, "regular_glm.jinja"), "w", encoding="utf-8") as f:
    f.write(regular_glm)

print("Created all standalone regular_*.jinja templates!")