import gradio as gr
import re
from typing import List, Tuple
import html

def create_collapsible_html(thinking: str) -> str:
    """Create HTML for a collapsible thinking section"""
    escaped_thinking = html.escape(thinking)
    return f'''
    <div class="thinking-container" style="margin: 8px 0;">
        <details class="thinking-details" style="
            border: 1px solid #ddd;
            border-radius: 4px;
            padding: 8px;
            background-color: #f8f9fa;
            cursor: pointer;
        ">
            <summary style="
                font-weight: 500;
                display: flex;
                align-items: center;
                gap: 8px;
            ">
                <span style="transform: rotate(-90deg); display: inline-block;">➤</span>
                Show thoughts
            </summary>
            <div style="
                margin-top: 8px;
                padding: 8px;
                border-top: 1px solid #eee;
                white-space: pre-wrap;
            ">{escaped_thinking}</div>
        </details>
    </div>
    '''

def process_message(message: str) -> str:
    """Process a message to convert thinking sections into collapsible elements"""
    def replace_thinking(match):
        thinking_content = match.group(1).strip()
        return create_collapsible_html(thinking_content)
    
    # Replace thinking sections with collapsible elements
    processed = re.sub(r'<think>(.*?)</think>', replace_thinking, message, flags=re.DOTALL)
    
    # Remove output tags if present
    processed = re.sub(r'<output>(.*?)</output>', r'\1', processed, flags=re.DOTALL)
    
    return processed

class CustomChatInterface(gr.ChatInterface):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
    def render_message(self, message: str, is_user: bool) -> str:
        if not is_user:
            return process_message(message)
        return message

demo = CustomChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You must always include <think> ... </think> <output> </output> tokens.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    css="""
    .thinking-details[open] summary span {
        transform: rotate(0deg) !important;
    }
    .thinking-details summary::-webkit-details-marker {
        display: none;
    }
    """
)

if __name__ == "__main__":
    demo.launch()