import re
import jinja2

# 1. Generate universal.jinja from qwen.jinja base with universal tags
with open("/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates/qwen.jinja") as f:
    qwen_content = f.read()

# For universal.jinja:
# Header:
universal_header = """{#- ================================================================= -#}
{#- HyperSynapse-SAS (SelfAgentSwarm) 10-Level Cognitive Architecture -#}
{#- Model-Agnostic / Universal ChatML Chat Template                   -#}
{#- Compatible with any LLM runtime (llama.cpp, vLLM, SGLang, Ollama) -#}
{#- Levels: 0=Mortal/Disabled, 1=Hermes/ulow, 2=Apollo/low,           -#}
{#- 3=Artemis/lmed, 4=Athena/medium, 5=Prometheus/mhigh,              -#}
{#- 6=Solstice/high, 7=Hyperion/xhigh (Default), 8=Einstein (Solo),   -#}
{#- 9=Oracle (Solo). Swarm: swarm (20-Agent), deep-swarm (Council)    -#}
{#- ================================================================= -#}
"""

universal_content = universal_header + qwen_content[qwen_content.find("{%- set inst_hermes -%}"):]
with open("/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates/universal.jinja", "w") as f:
    f.write(universal_content)

print("Created universal.jinja!")

# 2. Build Llama 3 / 4 template
# In Llama 3 / 4, special tokens are:
# <|begin_of_text|>
# <|start_header_id|>system<|end_header_id|>\n\n...<|eot_id|>
# <|start_header_id|>user<|end_header_id|>\n\n...<|eot_id|>
# <|start_header_id|>assistant<|end_header_id|>\n\n...<|eot_id|>

llama_header = """{#- ================================================================= -#}
{#- HyperSynapse-SAS (SelfAgentSwarm) 10-Level Cognitive Architecture -#}
{#- Tailored for Meta Llama 3 / 3.1 / 3.2 / 3.3 / Llama 4             -#}
{#- Levels: 0=Mortal/Disabled, 1=Hermes/ulow, 2=Apollo/low,           -#}
{#- 3=Artemis/lmed, 4=Athena/medium, 5=Prometheus/mhigh,              -#}
{#- 6=Solstice/high, 7=Hyperion/xhigh (Default), 8=Einstein (Solo),   -#}
{#- 9=Oracle (Solo). Swarm: swarm (20-Agent), deep-swarm (Council)    -#}
{#- ================================================================= -#}
"""

# Replace ChatML tokens in qwen rendering section with Llama tokens
qwen_instructions_part = qwen_content[qwen_content.find("{%- set inst_hermes -%}"):qwen_content.find("{%- set ns = namespace(")]
qwen_logic_part = qwen_content[qwen_content.find("{%- set ns = namespace("):]

# Adapt rendering logic for Llama
llama_render = qwen_logic_part.replace("<|im_start|>system\\n", "<|start_header_id|>system<|end_header_id|>\\n\\n")
llama_render = llama_render.replace("<|im_start|>user\\n", "<|start_header_id|>user<|end_header_id|>\\n\\n")
llama_render = llama_render.replace("<|im_start|>assistant\\n", "<|start_header_id|>assistant<|end_header_id|>\\n\\n")
llama_render = llama_render.replace("<|im_start|>developer\\n", "<|start_header_id|>developer<|end_header_id|>\\n\\n")
llama_render = llama_render.replace("<|im_end|>\\n", "<|eot_id|>\\n")
llama_render = llama_render.replace("<|im_end|>", "<|eot_id|>")

llama_content = llama_header + qwen_instructions_part + "{{- bos_token -}}\n" + llama_render
with open("/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates/llama3.jinja", "w") as f:
    f.write(llama_content)

print("Created llama3.jinja!")

# 3. Build Mistral template
mistral_header = """{#- ================================================================= -#}
{#- HyperSynapse-SAS (SelfAgentSwarm) 10-Level Cognitive Architecture -#}
{#- Tailored for Mistral / Mixtral / Codestral / Pixtral              -#}
{#- Levels: 0=Mortal/Disabled, 1=Hermes/ulow, 2=Apollo/low,           -#}
{#- 3=Artemis/lmed, 4=Athena/medium, 5=Prometheus/mhigh,              -#}
{#- 6=Solstice/high, 7=Hyperion/xhigh (Default), 8=Einstein (Solo),   -#}
{#- 9=Oracle (Solo). Swarm: swarm (20-Agent), deep-swarm (Council)    -#}
{#- ================================================================= -#}
"""
# Mistral uses [INST] and [/INST]
mistral_render = qwen_logic_part.replace("<|im_start|>system\\n", "[INST] ")
mistral_render = mistral_render.replace("<|im_start|>user\\n", "[INST] ")
mistral_render = mistral_render.replace("<|im_start|>assistant\\n", "")
mistral_render = mistral_render.replace("<|im_end|>\\n", " [/INST]\n")
mistral_render = mistral_render.replace("<|im_end|>", " [/INST]")

mistral_content = mistral_header + qwen_instructions_part + "{{- bos_token -}}\n" + mistral_render
with open("/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates/mistral.jinja", "w") as f:
    f.write(mistral_content)

print("Created mistral.jinja!")

# 4. Build Gemma template
gemma_header = """{#- ================================================================= -#}
{#- HyperSynapse-SAS (SelfAgentSwarm) 10-Level Cognitive Architecture -#}
{#- Tailored for Google Gemma 2 / Gemma 3 / Gemma 4                   -#}
{#- Levels: 0=Mortal/Disabled, 1=Hermes/ulow, 2=Apollo/low,           -#}
{#- 3=Artemis/lmed, 4=Athena/medium, 5=Prometheus/mhigh,              -#}
{#- 6=Solstice/high, 7=Hyperion/xhigh (Default), 8=Einstein (Solo),   -#}
{#- 9=Oracle (Solo). Swarm: swarm (20-Agent), deep-swarm (Council)    -#}
{#- ================================================================= -#}
"""
gemma_render = qwen_logic_part.replace("<|im_start|>system\\n", "<start_of_turn>system\\n")
gemma_render = gemma_render.replace("<|im_start|>user\\n", "<start_of_turn>user\\n")
gemma_render = gemma_render.replace("<|im_start|>assistant\\n", "<start_of_turn>model\\n")
gemma_render = gemma_render.replace("<|im_start|>developer\\n", "<start_of_turn>system\\n")
gemma_render = gemma_render.replace("<|im_end|>\\n", "<end_of_turn>\\n")
gemma_render = gemma_render.replace("<|im_end|>", "<end_of_turn>")

gemma_content = gemma_header + qwen_instructions_part + "{{- bos_token -}}\n" + gemma_render
with open("/Users/shreyangondaliya/solstice_quants/HyperSynapse-SAS-Chat-Templates/templates/gemma.jinja", "w") as f:
    f.write(gemma_content)

print("Created gemma.jinja!")