"""
Custom chat template for Xoron-Dev multimodal model.

This template uses our special tokens instead of the pretrained tokenizer's template.
Supports all SOTA features:
- Multimodal inputs (images, videos, audio, documents)
- Chain-of-thought reasoning (think, plan, critique, analysis, observation, reflection)
- Tool/function calling with structured arguments
- Code execution (Jupyter, exec)
- Memory and context management
- Uncertainty/confidence markers
- Generation outputs (gen_image, gen_video)
- Voice/TTS with emotion and prosody markers
- FIM (Fill-in-the-Middle) for code completion
"""

from xorfice.config.special_tokens import SPECIAL_TOKENS

# Clean, standard chat template using only structural Control Tokens.
# Domain tokens (like <|think|>, <|image|>, <|tool_call|>) are treated as vocabulary 
# and generated/parsed naturally within the assistant's or user's content block.
XORON_CHAT_TEMPLATE = """{%- set bos = '<|bos|>' -%}
{%- set eos = '<|eos|>' -%}
{%- set system_start = '<|system|>' -%}
{%- set system_end = '<|/system|>' -%}
{%- set user_start = '<|user|>' -%}
{%- set user_end = '<|/user|>' -%}
{%- set assistant_start = '<|assistant|>' -%}
{%- set assistant_end = '<|/assistant|>' -%}

{{- bos -}}
{%- for message in messages -%}
    {%- if message['role'] == 'system' -%}
        {{- system_start + message['content'] + system_end -}}
    {%- elif message['role'] == 'user' -%}
        {{- user_start + message['content'] + user_end -}}
    {%- elif message['role'] == 'assistant' -%}
        {{- assistant_start + message['content'] -}}
        {%- if not loop.last or add_generation_prompt is not defined or not add_generation_prompt -%}
            {{- assistant_end -}}
        {%- endif -%}
    {%- elif message['role'] == 'tool' or message['role'] == 'exec_result' or message['role'] == 'jupyter' -%}
        {{- user_start + message['content'] + user_end -}}
    {%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt is defined and add_generation_prompt -%}
    {{- assistant_start -}}
{%- endif -%}
"""


XORON_CHAT_TEMPLATE_SIMPLE = XORON_CHAT_TEMPLATE


def get_chat_template(multimodal: bool = True) -> str:
    """
    Get the appropriate chat template.
    
    Args:
        multimodal: If True, returns template with full multimodal support.
                   If False, returns simpler text-only template with basic reasoning.
    
    Returns:
        Jinja2 chat template string
    """
    return XORON_CHAT_TEMPLATE if multimodal else XORON_CHAT_TEMPLATE_SIMPLE


def apply_chat_template_to_tokenizer(tokenizer, multimodal: bool = True):
    """
    Apply the Xoron chat template to a tokenizer.
    
    This replaces any existing chat_template from the pretrained tokenizer
    with our custom template that uses Xoron special tokens.
    
    Args:
        tokenizer: HuggingFace tokenizer instance
        multimodal: Whether to use multimodal template
    
    Returns:
        Modified tokenizer with custom chat_template
    """
    tokenizer.chat_template = get_chat_template(multimodal)
    
    # Also set the special tokens properly
    tokenizer.bos_token = SPECIAL_TOKENS['bos']
    tokenizer.eos_token = SPECIAL_TOKENS['eos']
    tokenizer.pad_token = SPECIAL_TOKENS['pad']
    
    return tokenizer


def format_chat_example(messages: list, tokenizer=None, **kwargs) -> str:
    """
    Format a list of messages using the chat template.
    
    Args:
        messages: List of message dicts with 'role' and 'content' keys
        tokenizer: Optional tokenizer (if None, uses template directly)
        **kwargs: Additional template variables (tools, memory, context, etc.)
    
    Returns:
        Formatted string
    
    Example:
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Hello!", "images": ["<image_placeholder>"]},
            {"role": "assistant", "content": "Hi there!", "thinking": "User greeted me."}
        ]
        
        # With tools
        tools = '[{"name": "search", "description": "Search the web"}]'
        format_chat_example(messages, tools=tools)
    """
    if tokenizer is not None and hasattr(tokenizer, 'apply_chat_template'):
        return tokenizer.apply_chat_template(messages, tokenize=False, **kwargs)
    
    # Manual formatting without tokenizer
    from jinja2 import Template
    template = Template(XORON_CHAT_TEMPLATE)
    return template.render(messages=messages, add_generation_prompt=False, **kwargs)


from typing import Optional, List, Dict

def format_multimodal_message(
    content: str,
    images: Optional[List[str]] = None,
    videos: Optional[List[str]] = None,
    audio: Optional[List[str]] = None,
    documents: Optional[List[str]] = None,
    role: str = "user"
) -> Dict:
    """
    Helper to create a properly formatted multimodal message.
    
    Args:
        content: Text content of the message
        images: List of image placeholders or paths
        videos: List of video placeholders or paths
        audio: List of audio placeholders or paths
        documents: List of document contents
        role: Message role (user, assistant, system)
    
    Returns:
        Formatted message dict
    """
    # Prepend multimodal domain tokens directly into the content
    domain_blocks = []
    
    if images:
        for img in images:
            domain_blocks.append(f"<|image|>{img}<|/image|>")
    if videos:
        for vid in videos:
            domain_blocks.append(f"<|video|>{vid}<|/video|>")
    if audio:
        for aud in audio:
            domain_blocks.append(f"<|audio|>{aud}<|/audio|>")
    if documents:
        for doc in documents:
            domain_blocks.append(f"<|doc|>{doc}<|/doc|>")
            
    full_content = "".join(domain_blocks) + content
    message = {"role": role, "content": full_content}
    return message


def format_assistant_response(
    content: str,
    thinking: Optional[str] = None,
    planning: Optional[str] = None,
    analysis: Optional[str] = None,
    observation: Optional[str] = None,
    reflection: Optional[str] = None,
    critique: Optional[str] = None,
    conclusion: Optional[str] = None,
    tool_calls: Optional[List[Dict]] = None,
    code: Optional[str] = None,
    gen_image: Optional[str] = None,
    gen_video: Optional[str] = None,
    speak: Optional[str] = None,
    uncertain: Optional[str] = None,
    citation: Optional[str] = None,
) -> Dict:
    """
    Helper to create a properly formatted assistant response with reasoning.
    
    Args:
        content: Main response content
        thinking: Chain-of-thought reasoning
        planning: Planning steps
        analysis: Analysis content
        observation: Observations
        reflection: Reflections
        critique: Self-critique
        conclusion: Conclusions
        tool_calls: List of tool calls (dicts with 'name' and 'arguments')
        code: Code to execute
        gen_image: Image generation prompt
        gen_video: Video generation prompt
        speak: TTS content
        uncertain: Uncertain content
        citation: Citation/source
    
    Returns:
        Formatted assistant message dict
    """
    message = {"role": "assistant", "content": content}
    
    # Prepend domain tokens directly into the content since they are vocabulary now.
    domain_blocks = []
    
    if thinking:
        domain_blocks.append(f"<|think|>{thinking}<|/think|>")
    if planning:
        domain_blocks.append(f"<|plan|>{planning}<|/plan|>")
    if analysis:
        domain_blocks.append(f"<|analysis|>{analysis}<|/analysis|>")
    if observation:
        domain_blocks.append(f"<|observation|>{observation}<|/observation|>")
    if reflection:
        domain_blocks.append(f"<|reflection|>{reflection}<|/reflection|>")
    if critique:
        domain_blocks.append(f"<|critique|>{critique}<|/critique|>")
    if conclusion:
        domain_blocks.append(f"<|conclusion|>{conclusion}<|/conclusion|>")
    
    if tool_calls:
        for tool in tool_calls:
            import json
            tool_name = tool.get('name', '')
            args = tool.get('arguments', '')
            if isinstance(args, dict):
                args = json.dumps(args)
            domain_blocks.append(f"<|tool_call|><|function_name|>{tool_name}<|/function_name|><|function_args|>{args}<|/function_args|><|/tool_call|>")
            
    if code:
        domain_blocks.append(f"<|code|>{code}<|/code|>")
    if gen_image:
        domain_blocks.append(f"<|gen_image|>{gen_image}<|/gen_image|>")
    if gen_video:
        domain_blocks.append(f"<|gen_video|>{gen_video}<|/gen_video|>")
    if speak:
        domain_blocks.append(f"<|speak|>{speak}<|/speak|>")
    if uncertain:
        domain_blocks.append(f"<|uncertain|>{uncertain}<|/uncertain|>")
    if citation:
        domain_blocks.append(f"<|cite|>{citation}<|/cite|>")
        
    full_content = "".join(domain_blocks) + content
    message["content"] = full_content
    return message
