{% extends "admin/base.html" %}

{% block title %}Chat Playground{% endblock %}

{# This block is intentionally empty to remove the default page title and save vertical space. #}
{% block header_title %}{% endblock %}

{% block content %}
<style>
    /* Make the playground's parent containers fill the height and disable their scrolling/padding */
    main {
        overflow-y: hidden !important;
        padding: 0 !important;
        display: flex;
        flex-direction: column;
        height: 100%;
    }
    main > .max-w-7xl {
        max-width: none !important;
        margin: 0 !important;
        height: 100%;
        flex-grow: 1;
        display: flex;
        flex-direction: column;
    }
</style>

<div class="flex h-full w-full bg-[#0a0a0f] rounded-lg overflow-hidden border border-white/10 shadow-2xl">
    <!-- Chat History Sidebar -->
    <div id="chat-history-sidebar" class="w-64 border-r border-white/10 flex-col bg-black/40 hidden md:flex flex-shrink-0">
        <div class="p-4 border-b border-white/10 space-y-2">
            <button id="new-chat-btn" class="w-full bg-indigo-600 hover:bg-indigo-500 text-white px-4 py-2 rounded-lg font-black tracking-widest text-xs uppercase transition-colors shadow-lg shadow-indigo-500/20">+ New Chat</button>
            <button onclick="if(confirm('Clear all local chat history?')){localStorage.removeItem(HISTORY_STORAGE_KEY); window.location.reload();}" class="w-full bg-white/5 hover:bg-red-900/20 text-gray-500 hover:text-red-400 py-1.5 rounded-lg text-[10px] font-bold uppercase transition-all border border-transparent hover:border-red-500/30">Purge Local Cache</button>
        </div>
        <div id="chat-history-list" class="flex-1 overflow-y-auto custom-scrollbar p-2 space-y-1">
            <!-- History items injected here -->
        </div>
    </div>
    
    <!-- Main Chat Area -->
    <div class="flex-1 flex flex-col h-full relative min-w-0">
        <!-- Header -->
        <div class="card-header p-2 md:p-4 flex items-center justify-between flex-shrink-0 border-b border-white/10 bg-black/20">
        <div class="flex items-center gap-2 md:gap-4">
            <h2 class="text-xl font-bold">Chat</h2>
            <div id="system-prompt-indicator" class="hidden items-center gap-1 text-xs px-2 py-1 bg-sky-500/20 text-sky-300 rounded-full cursor-pointer" title="A system prompt is active. Click to edit.">
                <svg class="w-3 h-3" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0 3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.608 2.296.07 2.572-1.065z"></path><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"></path></svg>
                <span>System Prompt</span>
            </div>
        </div>
        <div class="flex items-center gap-2">
            <div id="thinking-toggle-container" class="hidden md:flex items-center gap-2 text-sm pr-2 border-r border-gray-600" title="Enable model thinking. This will show the reasoning trace before the final answer.">
                <span class="text-gray-400">Thinking</span>
                <button type="button" id="thinking-toggle" class="relative inline-flex flex-shrink-0 h-6 w-11 border-2 border-transparent rounded-full cursor-pointer transition-colors ease-in-out duration-200 focus:outline-none bg-gray-600">
                    <span id="thinking-toggle-knob" class="pointer-events-none inline-block h-5 w-5 rounded-full bg-white shadow-lg transform ring-0 transition ease-in-out duration-200 translate-x-0"></span>
                </button>
            </div>
            <div id="thinking-level-container" class="hidden md:flex items-center gap-2 text-sm pr-2 border-r border-gray-600" title="Set the thinking trace length for GPT-OSS models.">
                <span class="text-gray-400">Think Level</span>
                <select id="thinking-level-select" class="px-2 py-1 rounded-md text-sm">
                    <option value="low">Low</option>
                    <option value="medium" selected>Medium</option>
                    <option value="high">High</option>
                </select>
            </div>
            <div id="image-annotation-toggle-container" class="hidden md:flex items-center gap-2 text-sm pr-2 border-r border-gray-600" title="Toggle image annotation mode. This adds a special system prompt to enable object detection and bounding box drawing on images.">
                <span class="text-gray-400">Annotate</span>
                <button type="button" id="image-annotation-toggle" class="relative inline-flex flex-shrink-0 h-6 w-11 border-2 border-transparent rounded-full cursor-pointer transition-colors ease-in-out duration-200 focus:outline-none bg-gray-600">
                    <span id="image-annotation-toggle-knob" class="pointer-events-none inline-block h-5 w-5 rounded-full bg-white shadow-lg transform ring-0 transition ease-in-out duration-200 translate-x-0"></span>
                </button>
            </div>
            <div class="relative" id="add-message-menu-container">
                <button id="add-message-btn" title="Add Message" class="p-2 rounded hover:bg-white/10">
                    <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M12 9v3m0 0v3m0-3h3m-3 0H9m12 0a9 9 0 11-18 0 9 9 0 0118 0z"></path></svg>
                </button>
                <div id="add-message-dropdown" class="absolute bottom-full right-0 mb-2 w-48 rounded-md shadow-lg bg-gray-800 ring-1 ring-black ring-opacity-5 hidden z-10">
                    <a href="#" id="add-user-msg-btn" class="block px-4 py-2 text-sm text-gray-300 hover:bg-gray-700">Add User Message</a>
                    <a href="#" id="add-ai-msg-btn" class="block px-4 py-2 text-sm text-gray-300 hover:bg-gray-700">Add AI Message</a>
                </div>
            </div>
            <button id="system-prompt-btn" title="Set System Prompt" class="p-2 rounded hover:bg-white/10">
                <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0 3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.608 2.296.07 2.572-1.065z"></path><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"></path></svg>
            </button>
            <select id="model-selector" class="px-3 py-2 rounded-md shadow-sm focus:outline-none focus:ring-[var(--color-primary-500)] focus:border-[var(--color-primary-500)] sm:text-sm">
                {% if model_groups %}
                    {% for server_name, models in model_groups.items() %}
                        <optgroup label="{{ server_name }}">
                        {% for model in models %}
                            <option value="{{ model }}" {% if model == selected_model %}selected{% endif %}>{{ model }}</option>
                        {% endfor %}
                        </optgroup>
                    {% endfor %}
                {% else %}
                    <option value="" disabled>No chat models found</option>
                {% endif %}
            </select>
            <button id="test-prompts-btn" title="Test Prompts" class="p-2 rounded hover:bg-white/10">
                 <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z"></path></svg>
            </button>
            <input type="file" id="import-chat-input" class="hidden" accept=".json">
            <button id="import-chat-btn" title="Import Chat" class="p-2 rounded hover:bg-white/10">
                <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-8l-4-4m0 0L8 8m4-4v12"></path></svg>
            </button>
            <button id="export-chat-btn" title="Export Chat" class="p-2 rounded hover:bg-white/10">
                <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"></path></svg>
            </button>
            <button id="clear-chat-btn" class="px-3 py-2 bg-red-600 text-white rounded hover:bg-red-700 focus:outline-none">Clear</button>
        </div>
    </div>

    <!-- Chat Messages -->
    <div id="chat-container" class="flex-grow p-2 md:p-4 overflow-y-auto space-y-4">
        <!-- Messages will be injected here -->
        <div class="text-center text-gray-400">Select a model and start chatting.</div>
    </div>
    
    <!-- Global Stats Bar -->
    <div id="stats-bar" class="p-2 flex-shrink-0 border-t border-b border-white/10 text-xs text-gray-400 flex items-center justify-center gap-6">
        <div>Mean TTFT: <span id="mean-ttft" class="font-semibold text-gray-300">N/A</span></div>
        <div>Mean TPS: <span id="mean-tps" class="font-semibold text-gray-300">N/A</span></div>
    </div>

    <!-- Message Input -->
    <div class="p-2 md:p-4 flex-shrink-0">
        <div id="doc-previews" class="doc-previews"></div>
        <div id="image-previews" class="image-previews"></div>
        <form id="chat-form" class="flex items-end gap-2 md:gap-4">
            <button type="button" id="attach-img-btn" class="p-2 rounded-md hover:bg-white/10 flex-shrink-0" title="Attach Image">
                <svg class="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z"></path></svg>
            </button>
            <button type="button" onclick="openPlaygroundPaint()" class="p-2 rounded-md hover:bg-white/10 text-indigo-400 flex-shrink-0" title="Sketch/Annotate">
                <svg class="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15.232 5.232l3.536 3.536m-2.036-5.036a2.5 2.5 0 113.536 3.536L6.5 21.036H3v-3.5L15.232 5.232z"></path></svg>
            </button>
            <input type="file" id="image-input" class="hidden" multiple accept="image/*">
            <button type="button" id="attach-doc-btn" class="p-2 rounded-md hover:bg-white/10 flex-shrink-0" title="Attach Document">
                <svg class="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 12h6m-6 4h6m2 5H7a2 2 0 01-2-2V5a2 2 0 012-2h5.586a1 1 0 01.707.293l5.414 5.414a1 1 0 01.293.707V19a2 2 0 01-2 2z"></path></svg>
            </button>
            <button type="button" id="global-import-btn" class="p-2 rounded-md hover:bg-indigo-600/20 text-indigo-400 flex-shrink-0" title="Search & Import External Knowledge">
                <svg class="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z"></path></svg>
            </button>
            <input type="file" id="doc-input" class="hidden" multiple accept=".txt,.md,.json,.csv,.py,.js,.html,.css,.xml,.yaml,.yml,.log,.ini,.toml,.sh,.bat">
            <textarea id="message-input" placeholder="Type your message here, or paste an image..." rows="1" class="flex-grow block w-full px-3 py-2 rounded-md shadow-sm focus:outline-none focus:ring-[var(--color-primary-500)] focus:border-[var(--color-primary-500)] sm:text-sm resize-none"></textarea>
            <button type="submit" id="send-btn" class="justify-center py-2 px-6 border border-transparent rounded-md shadow-sm text-sm font-medium text-white bg-[var(--color-primary-600)] hover:bg-[var(--color-primary-700)] focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-[var(--color-primary-500)] self-end">Send</button>
        </form>
    </div>
    </div>
</div>

<script src="{{ url_for('static', path='vendor/marked.min.js') }}"></script>
<script src="{{ url_for('static', path='vendor/mermaid.min.js') }}"></script>
<script>
document.addEventListener('DOMContentLoaded', () => {
    mermaid.initialize({ startOnLoad: false, theme: 'dark' });
    
    // --- Marked.js & Renderer Setup ---
    const renderer = new marked.Renderer();
    const originalCodeRenderer = renderer.code;
    renderer.code = function(code, lang, escaped) {
      const originalHtml = originalCodeRenderer.call(this, code, lang, escaped);
      // We add a wrapper div that our script will later use to inject a copy button.
      return `<div class="code-block-wrapper">${originalHtml}</div>`;
    };
    marked.setOptions({
        renderer: renderer,
        gfm: true,
        breaks: true,
        smartypants: false
    });

    // --- Main Rendering Function ---
    function highlightAndMathify(element) {
        element.querySelectorAll('pre code').forEach((block) => {
            hljs.highlightElement(block);
        });
        if (window.renderMathInElement) {
            renderMathInElement(element, {
                delimiters: [
                    {left: '$$', right: '$$', display: true}, {left: '$', right: '$', display: false},
                    {left: '\\[', right: '\\]', display: true}, {left: '\\(', right: '\\)', display: false}
                ],
                throwOnError: false
            });
        }
    }
    
    // --- DOM Elements ---
    const chatContainer = document.getElementById('chat-container');
    const chatForm = document.getElementById('chat-form');
    const messageInput = document.getElementById('message-input');
    const modelSelector = document.getElementById('model-selector');
    const sendBtn = document.getElementById('send-btn');
    const clearChatBtn = document.getElementById('clear-chat-btn');
    const importChatBtn = document.getElementById('import-chat-btn');
    const importChatInput = document.getElementById('import-chat-input');
    const exportChatBtn = document.getElementById('export-chat-btn');
    const testPromptsBtn = document.getElementById('test-prompts-btn');
    const systemPromptBtn = document.getElementById('system-prompt-btn');
    const systemPromptIndicator = document.getElementById('system-prompt-indicator');
    const attachImgBtn = document.getElementById('attach-img-btn');
    const imageInput = document.getElementById('image-input');
    const imagePreviews = document.getElementById('image-previews');
    const attachDocBtn = document.getElementById('attach-doc-btn');
    const docInput = document.getElementById('doc-input');
    const docPreviews = document.getElementById('doc-previews');
    const annotationToggle = document.getElementById('image-annotation-toggle');
    const annotationToggleKnob = document.getElementById('image-annotation-toggle-knob');
    const thinkingToggleContainer = document.getElementById('thinking-toggle-container');
    const thinkingToggle = document.getElementById('thinking-toggle');
    const thinkingToggleKnob = document.getElementById('thinking-toggle-knob');
    const thinkingLevelContainer = document.getElementById('thinking-level-container');
    const thinkingLevelSelect = document.getElementById('thinking-level-select');
    const addMessageMenuContainer = document.getElementById('add-message-menu-container');
    const addMessageBtn = document.getElementById('add-message-btn');
    const addMessageDropdown = document.getElementById('add-message-dropdown');
    const addUserMsgBtn = document.getElementById('add-user-msg-btn');
    const addAiMsgBtn = document.getElementById('add-ai-msg-btn');


    // --- Chat History Management ---
    const HISTORY_STORAGE_KEY = 'lollms_playground_chats';
    let chatHistory = JSON.parse(localStorage.getItem(HISTORY_STORAGE_KEY) || '[]');
    let currentChatId = null;

    function saveHistory() {
        try {
            localStorage.setItem(HISTORY_STORAGE_KEY, JSON.stringify(chatHistory));
            renderHistoryList();
        } catch (e) {
            if (e.name === 'QuotaExceededError' || e.name === 'NS_ERROR_DOM_QUOTA_REACHED') {
                console.warn("LocalStorage quota exceeded. Pruning oldest chats...");
                
                // 1. If we have more than 1 chat, remove the oldest one and retry
                if (chatHistory.length > 1) {
                    chatHistory.pop(); // Remove the last (oldest) item
                    saveHistory();     // Recursive retry
                } else {
                    // 2. If even a single chat is too big (likely due to too many images), 
                    // we must inform the user and stop the crash.
                    console.error("Critical: Single chat session exceeds browser storage limits.");
                    alert("Storage Limit Reached: This conversation is too large to save locally (likely due to high-resolution images). Previous chats have been cleared to make room, but you may need to clear this chat or use fewer images to continue saving history.");
                }
            } else {
                console.error("LocalStorage Save Error:", e);
            }
        }
    }

    window.createNewChat = () => {
        currentChatId = Date.now().toString();
        messages = [];
        attachedImages = [];
        attachedDocs = [];
        lastSentImageURL = null;
        chatHistory.unshift({
            id: currentChatId,
            title: 'New Chat',
            updatedAt: Date.now(),
            messages: [],
            systemPrompt: userSystemPrompt,
            model: modelSelector.value
        });
        saveHistory();
        renderChat();
        renderImagePreviews();
        renderDocPreviews();
        updateGlobalStats();
    };

    window.updateCurrentChat = () => {
        if (!currentChatId) createNewChat();
        const chat = chatHistory.find(c => c.id === currentChatId);
        if (chat) {
            chat.messages = messages;
            chat.systemPrompt = userSystemPrompt;
            chat.model = modelSelector.value;
            chat.updatedAt = Date.now();
            
            if (chat.title === 'New Chat' && messages.length > 0) {
                const firstUserMsg = messages.find(m => m.role === 'user');
                if (firstUserMsg) {
                    let text = typeof firstUserMsg.content === 'string' ? firstUserMsg.content : (firstUserMsg.content.find(c => c.type === 'text')?.text || 'Image Chat');
                    chat.title = text.substring(0, 30).replace(/\n/g, ' ') + (text.length > 30 ? '...' : '');
                }
            }
        }
        saveHistory();
    };

    window.loadChat = (id) => {
        const chat = chatHistory.find(c => c.id === id);
        if (!chat) return;
        currentChatId = id;
        messages = chat.messages || [];
        userSystemPrompt = chat.systemPrompt || '';
        if (chat.model && modelSelector.querySelector(`option[value="${chat.model}"]`)) {
            modelSelector.value = chat.model;
        }
        updateSystemPromptIndicator();
        renderChat();
        updateGlobalStats();
        renderHistoryList(); // Highlight active
    };

    window.deleteChat = (id, e) => {
        e.stopPropagation();
        if(!confirm('Delete this chat?')) return;
        chatHistory = chatHistory.filter(c => c.id !== id);
        if (currentChatId === id) {
            if (chatHistory.length > 0) {
                loadChat(chatHistory[0].id);
            } else {
                createNewChat();
            }
        } else {
            saveHistory();
        }
    };

    function renderHistoryList() {
        const list = document.getElementById('chat-history-list');
        if (!list) return;
        list.innerHTML = chatHistory.map(chat => `
            <div class="group flex items-center justify-between p-2 rounded-lg cursor-pointer transition-colors ${chat.id === currentChatId ? 'bg-indigo-500/20 text-white border border-indigo-500/30' : 'text-gray-400 hover:bg-white/5 border border-transparent'}" onclick="loadChat('${chat.id}')">
                <div class="truncate text-xs font-medium flex-1 mr-2">${chat.title}</div>
                <button class="opacity-0 group-hover:opacity-100 text-gray-500 hover:text-red-400 p-1" onclick="deleteChat('${chat.id}', event)">
                    <svg class="w-3 h-3" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 7l-.867 12.142A2 2 0 0116.138 21H7.862a2 2 0 01-1.995-1.858L5 7m5 4v6m4-6v6m1-10V4a1 1 0 00-1-1h-4a1 1 0 00-1 1v3M4 7h16"></path></svg>
                </button>
            </div>
        `).join('');
    }

    // --- Constants & State ---
    const ANNOTATION_PROMPT = `You are an advanced AI assistant for image analysis and annotation. Your primary function is to process images and return structured data for various tasks like object detection, text translation, and redaction.

CRITICAL INSTRUCTIONS:
1.  Your response MUST contain a single JSON array of annotation objects wrapped in <annotate> tags.
2.  You can add conversational text AFTER the closing </annotate> tag to answer questions, but the JSON part is mandatory.
3.  All coordinates MUST be normalized (values between 0.0 and 1.0). (0,0) is top-left, (1,1) is bottom-right.
4.  If no objects are found or the task cannot be performed, return an empty array: <annotate>[]</annotate>

SUPPORTED ANNOTATION SCHEMAS:

1.  **Bounding Box (\`bbox\`)**: To draw a box around an object.
    -   \`type\`: "bbox"
    -   \`bbox\`: [x_min, y_min, x_max, y_max] (array of 4 numbers)
    -   \`label\`: (string) The name of the object.
    -   \`display\`: (optional object)
        -   \`border_color\`: (string, hex code) e.g., "#ff00ff"
        -   \`fill_opacity\`: (number, 0.0 to 1.0) Opacity of the box fill.

    *Example*: Find a cat.
    <annotate>
    [
      {"type": "bbox", "bbox": [0.1, 0.2, 0.4, 0.5], "label": "cat", "display": {"border_color": "#00ff00"}}
    ]
    </annotate>

2.  **Point (\`point\`)**: To mark a specific point on the image.
    -   \`type\`: "point"
    -   \`coords\`: [x, y] (array of 2 numbers)
    -   \`label\`: (string) The name of the point.
    -   \`display\`: (optional object) with \`border_color\`.

    *Example*: Mark the cat's eye.
    <annotate>
    [
      {"type": "point", "coords": [0.25, 0.3], "label": "cat's left eye"}
    ]
    </annotate>

3.  **Polygon (\`polygon\`)**: To outline a complex shape. (Note: The UI will render this as a dashed bounding box around the points).
    -   \`type\`: "polygon"
    -   \`points\`: [[x1, y1], [x2, y2], ...] (array of point arrays)
    -   \`label\`: (string) The name of the shape.

    *Example*: Outline the shape of the roof.
    <annotate>
    [
      {"type": "polygon", "points": [[0.2, 0.1], [0.8, 0.1], [0.5, 0.4]], "label": "roof"}
    ]
    </annotate>

4.  **Redaction (\`redaction\`)**: To blur a sensitive area.
    -   \`type\`: "redaction"
    -   \`bbox\`: [x_min, y_min, x_max, y_max]
    -   \`action\`: "blur"
    -   \`strength\`: (optional number) Blur intensity, e.g., 10.

    *Example*: Hide the face.
    <annotate>
    [
      {"type": "redaction", "bbox": [0.6, 0.2, 0.8, 0.4], "action": "blur", "strength": 8}
    ]
    </annotate>

5.  **Replace (\`replace\`)**: To cover an area with a solid color and optional text.
    -   \`type\`: "replace"
    -   \`bbox\`: [x_min, y_min, x_max, y_max]
    -   \`replacement_text\`: (string) Text to display in the box.
    -   \`color\`: (optional string, hex code) Background color of the box.

    *Example*: Cover the license plate.
    <annotate>
    [
      {"type": "replace", "bbox": [0.4, 0.7, 0.6, 0.8], "replacement_text": "REDACTED", "color": "#000000"}
    ]
    </annotate>

6.  **Translated Text (\`translated_text\`)**: To overlay translated text on the image.
    -   \`type\`: "translated_text"
    -   \`bbox\`: [x_min, y_min, x_max, y_max] (The box where the original text was).
    -   \`translated_text\`: (string) The translated text.

    *Example*: Translate the sign to Spanish.
    <annotate>
    [
      {"type": "translated_text", "bbox": [0.3, 0.5, 0.7, 0.6], "translated_text": "Abierto"}
    ]
    </annotate>
    
7.  **Text Overlay (\`text\`)**: To place text directly onto the image.
    -   \`type\`: "text"
    -   \`text\`: (string) The text to display.
    -   \`position\`: [x, y] (array of 2 numbers) The top-left corner for the text.
    -   \`style\`: (optional object) A dictionary of CSS styles.
        -   Supported keys: \`font_size\`, \`font_weight\`, \`color\`, \`background_color\`, \`padding\`, \`border_radius\`.

    *Example*: Add a title and a caption.
    <annotate>
    [
      {"type": "text", "text": "Sunset Dreams", "position": [0.05, 0.1], "style": {"font_size": "32px", "color": "white", "font_weight": "bold"}},
      {"type": "text", "text": "A beautiful evening", "position": [0.05, 0.9], "style": {"font_size": "16px", "color": "white", "background_color": "rgba(0,0,0,0.5)", "padding": "5px"}}
    ]
    </annotate>`;
    
    const PREBUILT_SYSTEM_PROMPTS = [
        { name: "Standard Assistant", prompt: "You are a helpful, respectful and honest assistant." },
        { name: "Chain of Thought", prompt: "When you are asked a question, first provide a step-by-step plan of how you will answer the question. Then, produce the answer, explaining your reasoning at each step." },
        { name: "Image Bounding Box Detection", prompt: ANNOTATION_PROMPT },
        { name: "Python Expert", prompt: "You are an expert Python programmer. Provide concise, efficient, and well-documented code. Your answers should only contain code, unless asked otherwise." },
        { name: "Creative Writer", prompt: "You are a creative writer, tasked with writing a compelling and imaginative story based on the user's prompt. Focus on vivid descriptions and engaging narrative." }
    ];
    
    let messages = [];
    let isLoading = false;
    let abortController = null;
    let attachedImages = [];
    let attachedDocs = [];
    let lastSentImageURL = null;
    let userSystemPrompt = '';
    let isAnnotationMode = false;
    let isThinkingMode = false;
    const CUSTOM_PROMPTS_KEY = 'ollama_proxy_custom_prompts';
    const LAST_MODEL_KEY = 'ollama_proxy_last_model';
    const LAST_SYSTEM_PROMPT_KEY = 'ollama_proxy_last_system_prompt';
    const CUSTOM_SYSTEM_PROMPTS_KEY = 'ollama_proxy_custom_system_prompts';
    const ANNOTATION_MODE_KEY = 'ollama_proxy_annotation_mode';
    const THINKING_MODE_KEY = 'ollama_proxy_thinking_mode';
    const langExtMap = {
        python: 'py', py: 'py', javascript: 'js', js: 'js', typescript: 'ts',
        ts: 'ts', html: 'html', css: 'css', json: 'json', markdown: 'md', md: 'md',
        bash: 'sh', sh: 'sh', shell: 'sh', sql: 'sql', yaml: 'yaml', yml: 'yml',
        dockerfile: 'Dockerfile', java: 'java', c: 'c', cpp: 'cpp', csharp: 'cs',
        go: 'go', rust: 'rs', ruby: 'rb', php: 'php', perl: 'pl', xml: 'xml',
        ini: 'ini', toml: 'toml', plaintext: 'txt', text: 'txt', txt: 'txt'
    };


    // --- File & Image Handling ---
    function handleImageFiles(files) {
        for (const file of files) {
            if (!file.type.startsWith('image/')){ continue }
            const reader = new FileReader();
            reader.onload = e => {
                if (attachedImages.length < 5) {
                    attachedImages.push(e.target.result);
                    // Treat as Input Artifact
                    window.ArtifactRegistry.register(`<artifact name="${file.name}" type="image" path="${e.target.result}"/>`);
                    renderImagePreviews();
                } else {
                    alert("You can attach a maximum of 5 images.");
                }
            };
            reader.readAsDataURL(file);
        }
    }
    
    function handleDocFiles(files) {
        for (const file of files) {
            const reader = new FileReader();
            reader.onload = e => {
                if (attachedDocs.length < 5) {
                    attachedDocs.push({ name: file.name, content: e.target.result });
                    // Treat as Input Artifact
                    window.ArtifactRegistry.register(`<artifact name="${file.name}" type="file">${e.target.result}</artifact>`);
                    renderDocPreviews();
                } else {
                    alert("You can attach a maximum of 5 documents.");
                }
            };
            reader.onerror = () => alert(`Error reading file ${file.name}`);
            reader.readAsText(file);
        }
    }

    function renderImagePreviews() {
        imagePreviews.innerHTML = '';
        attachedImages.forEach((imgData, index) => {
            const previewContainer = document.createElement('div');
            previewContainer.className = 'img-preview-container';
            const img = document.createElement('img');
            img.src = imgData;
            img.className = 'img-preview';
            const removeBtn = document.createElement('button');
            removeBtn.className = 'remove-img-btn';
            removeBtn.innerHTML = '&times;';
            removeBtn.onclick = () => {
                attachedImages.splice(index, 1);
                renderImagePreviews();
            };
            previewContainer.appendChild(img);
            previewContainer.appendChild(removeBtn);
            imagePreviews.appendChild(previewContainer);
        });
    }
    
    function renderDocPreviews() {
        docPreviews.innerHTML = '';
        attachedDocs.forEach((doc, index) => {
            const el = document.createElement('div');
            el.className = 'doc-preview-container';
            el.innerHTML = `
                <svg class="w-4 h-4 text-gray-400 flex-shrink-0" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 12h6m-6 4h6m2 5H7a2 2 0 01-2-2V5a2 2 0 012-2h5.586a1 1 0 01.707.293l5.414 5.414a1 1 0 01.293.707V19a2 2 0 01-2 2z"></path></svg>
                <span class="doc-preview-name" title="${doc.name}">${doc.name}</span>
                <button class="remove-doc-btn" data-index="${index}" title="Remove file">&times;</button>
            `;
            docPreviews.appendChild(el);
        });
    }

    function handlePaste(event) {
        const files = event.clipboardData?.files;
        if (files && files.length > 0) {
            handleImageFiles(files);
            event.preventDefault();
        }
    }
    
    // --- Test Prompts Modal ---
    async function showTestPrompts() {
        const response = await fetch("{{ url_for('admin_get_test_prompts') }}");
        const prebuiltPrompts = await response.json();
        const customPrompts = JSON.parse(localStorage.getItem(CUSTOM_PROMPTS_KEY) || '[]');
        
        let modalBody = `
        <div class="space-y-4">
            <details class="border border-gray-700 rounded-md">
                <summary class="cursor-pointer p-3 font-semibold hover:bg-white/5">Add New Test Prompt</summary>
                <div id="add-prompt-form" class="p-3 border-t border-gray-700 space-y-3">
                    <input id="new-prompt-category" type="text" placeholder="Category (e.g., Coding)" class="w-full px-3 py-2 rounded-md">
                    <input id="new-prompt-title" type="text" placeholder="Title (e.g., Python FizzBuzz)" class="w-full px-3 py-2 rounded-md">
                    <textarea id="new-prompt-prompt" placeholder="Prompt text..." rows="3" class="w-full px-3 py-2 rounded-md"></textarea>
                    <textarea id="new-prompt-expected" placeholder="Expected outcome..." rows="2" class="w-full px-3 py-2 rounded-md"></textarea>
                    <button id="save-new-prompt-btn" class="py-2 px-4 bg-blue-600 text-white rounded hover:bg-blue-700">Save Custom Prompt</button>
                </div>
            </details>
            <div class="prompts-list-container" style="max-height: 60vh; overflow-y: auto; padding-right: 0.5rem;"></div>
        </div>`;
        
        window.showModal('Test Prompts Library', modalBody);
        renderPromptsList(prebuiltPrompts, customPrompts);
    }

    function renderPromptsList(prebuilt, custom) {
        const container = document.querySelector('.prompts-list-container');
        if (!container) return;

        const allPrompts = [...prebuilt.map(p => ({...p, isCustom: false})), ...custom.map(p => ({...p, isCustom: true}))];
        const categories = allPrompts.reduce((acc, p) => {
            if (!acc[p.category]) acc[p.category] = [];
            acc[p.category].push(p);
            return acc;
        }, {});

        let listHtml = '<div class="space-y-4">';
        for (const category in categories) {
            listHtml += `<h4 class="text-lg font-semibold mt-4">${category}</h4><div class="space-y-2">`;
            categories[category].forEach(p => {
                const deleteBtn = p.isCustom ? `<button class="delete-prompt-btn text-red-500 hover:text-red-700 text-lg" data-prompt-title="${escape(p.title)}" title="Delete">&times;</button>` : '';
                listHtml += `
                    <div class="p-3 border border-gray-700 rounded-md">
                        <div class="flex justify-between items-start">
                           <p class="font-semibold">${p.title} ${p.isCustom ? '<span class="text-xs font-normal text-blue-400">(Custom)</span>' : ''}</p>
                           ${deleteBtn}
                        </div>
                        <p class="text-sm text-gray-400 italic mt-1"><strong>Expected:</strong> ${p.expected_outcome}</p>
                        <button class="use-prompt-btn mt-2 text-sm text-[var(--color-primary-500)] hover:underline" data-prompt="${escape(p.prompt)}">Use this prompt</button>
                    </div>`;
            });
            listHtml += '</div>';
        }
        listHtml += '</div>';
        container.innerHTML = listHtml;
    }

    // --- System Prompt & Annotation Logic ---
    function getFinalSystemPrompt() {
        const prompts = [];
        if (userSystemPrompt) {
            prompts.push(userSystemPrompt);
        }
        if (isAnnotationMode) {
            prompts.push(ANNOTATION_PROMPT);
        }
        return prompts.join('\n\n');
    }
    
    function showSystemPromptModal() {
        const customSystemPrompts = JSON.parse(localStorage.getItem(CUSTOM_SYSTEM_PROMPTS_KEY) || '[]');
        
        let modalBody = `<div class="space-y-4">
            <div>
                <label for="system-prompt-textarea" class="block text-sm font-medium mb-1">Current System Prompt</label>
                <textarea id="system-prompt-textarea" rows="6" class="w-full p-2 rounded-md">${userSystemPrompt}</textarea>
            </div>
            <div class="flex flex-wrap gap-2 items-center">
                <div class="relative">
                     <button id="load-system-prompt-btn" class="py-2 px-4 text-sm bg-gray-600 rounded-md">Load Preset</button>
                     <div id="system-prompt-dropdown" class="absolute bottom-full mb-2 w-72 rounded-md shadow-lg bg-gray-800 ring-1 ring-black ring-opacity-5 hidden z-20"></div>
                </div>
                <input id="new-system-prompt-name" type="text" placeholder="Save current as preset..." class="flex-grow p-2 rounded-md text-sm">
                <button id="save-custom-system-prompt-btn" class="py-2 px-4 text-sm bg-blue-600 rounded-md">Save as Preset</button>
            </div>
             <div class="flex justify-end gap-4 pt-4 border-t border-gray-700">
                <button id="clear-system-prompt-btn" class="py-2 px-4 bg-red-600 text-white rounded">Clear & Close</button>
                <button id="apply-system-prompt-btn" class="py-2 px-4 bg-green-600 text-white rounded">Apply & Close</button>
            </div>
        </div>`;

        window.showModal('System Prompt', modalBody);

        const dropdown = document.getElementById('system-prompt-dropdown');
        let dropdownHTML = '<div class="py-1">';
        dropdownHTML += '<div class="px-4 py-2 text-xs uppercase text-gray-400">Pre-built</div>';
        PREBUILT_SYSTEM_PROMPTS.forEach(p => {
            dropdownHTML += `<a href="#" class="block px-4 py-2 text-sm text-gray-300 hover:bg-gray-700 system-prompt-item" data-prompt="${escape(p.prompt)}">${p.name}</a>`;
        });
        if (customSystemPrompts.length > 0) {
            dropdownHTML += '<div class="px-4 py-2 mt-2 text-xs uppercase text-gray-400 border-t border-gray-600">Custom</div>';
            customSystemPrompts.forEach(p => {
                 dropdownHTML += `<div class="flex items-center justify-between px-4 py-2 text-sm text-gray-300 hover:bg-gray-700">
                                    <a href="#" class="flex-grow system-prompt-item" data-prompt="${escape(p.prompt)}">${p.name}</a>
                                    <button class="delete-system-prompt-btn text-red-500 hover:text-red-400 ml-2" data-name="${escape(p.name)}">&times;</button>
                                 </div>`;
            });
        }
        dropdownHTML += '</div>';
        dropdown.innerHTML = dropdownHTML;
    }

    function updateSystemPromptIndicator() {
        const finalPrompt = getFinalSystemPrompt();
        if (finalPrompt) {
            systemPromptIndicator.classList.remove('hidden');
            systemPromptIndicator.classList.add('flex');
        } else {
            systemPromptIndicator.classList.add('hidden');
            systemPromptIndicator.classList.remove('flex');
        }
    }
    
    function updateAnnotationToggleUI() {
        if (isAnnotationMode) {
            annotationToggle.classList.remove('bg-gray-600');
            annotationToggle.classList.add('bg-[var(--color-primary-600)]');
            annotationToggleKnob.classList.remove('translate-x-0');
            annotationToggleKnob.classList.add('translate-x-5');
        } else {
            annotationToggle.classList.add('bg-gray-600');
            annotationToggle.classList.remove('bg-[var(--color-primary-600)]');
            annotationToggleKnob.classList.add('translate-x-0');
            annotationToggleKnob.classList.remove('translate-x-5');
        }
    }
    
    function updateThinkingToggleUI() {
        if (isThinkingMode) {
            thinkingToggle.classList.remove('bg-gray-600');
            thinkingToggle.classList.add('bg-[var(--color-primary-600)]');
            thinkingToggleKnob.classList.remove('translate-x-0');
            thinkingToggleKnob.classList.add('translate-x-5');
        } else {
            thinkingToggle.classList.add('bg-gray-600');
            thinkingToggle.classList.remove('bg-[var(--color-primary-600)]');
            thinkingToggleKnob.classList.add('translate-x-0');
            thinkingToggleKnob.classList.remove('translate-x-5');
        }
    }

    // --- Core Chat Functions ---
    function renderMermaid(container) {
        container.querySelectorAll('pre code.language-mermaid').forEach(async (block) => {
            const wrapper = block.parentElement.parentElement;
            const id = 'mm-' + Math.random().toString(36).substr(2, 9);
            const div = document.createElement('div');
            div.id = id;
            div.className = 'mermaid';
            div.textContent = block.innerText;
            wrapper.parentNode.replaceChild(div, wrapper);
            try { await mermaid.run({ nodes: [div] }); } catch (e) { console.error("Mermaid error", e); }
        });
    }

    async function fetchStreamedResponse() {
        if (isLoading) return;
        abortController = new AbortController();
        toggleLoading(true);
        renderChat();

        const botMessageWrapper = findLastBotMessageWrapper();
        const botContentElement = botMessageWrapper.querySelector('.message-content');
        const botStatsElement = botMessageWrapper.querySelector('.message-stats');
        
        const isScrolledToBottom = () => {
            return chatContainer.scrollHeight - chatContainer.clientHeight <= chatContainer.scrollTop + 50;
        };

        let fullResponse = '';
        let firstChunk = true;
        const startTime = performance.now();
        let streamError = null;
        let currentTTFT = null;
        let currentTPS = null;
        let buffer = '';

        // Create a CLEAN deep copy of messages to send to the API
        let messagesForApi = JSON.parse(JSON.stringify(messages));
        
        // Handle image formatting for Ollama compatibility
        messagesForApi = messagesForApi.map(msg => 
            Array.isArray(msg.content) && !msg.content.some(item => item.type === 'image_url') 
            ? { ...msg, content: msg.content.find(item => item.type === 'text')?.text || '' } 
            : msg
        );
        
        // PREPEND UNIVERSAL SYSTEM PROMPT
        const finalSystemPrompt = getFinalSystemPrompt();
        if (finalSystemPrompt && finalSystemPrompt.trim().length > 0) {
            // Check if there is already a system prompt in history (e.g. from an imported chat)
            const hasExistingSystem = messagesForApi.some(m => m.role === 'system');
            if (!hasExistingSystem) {
                messagesForApi.unshift({ role: "system", content: finalSystemPrompt });
            }
        }

        const body = { 
            model: modelSelector.value, 
            messages: messagesForApi
        };

        let thinkOption = null;
        if (!thinkingToggleContainer.classList.contains('hidden')) {
            if (isThinkingMode) {
                if (!thinkingLevelContainer.classList.contains('hidden')) {
                    thinkOption = thinkingLevelSelect.value;
                } else {
                    thinkOption = true;
                }
            } else {
                // Explicitly send false to disable native auto-thinking features
                thinkOption = false;
            }
        }
        if (thinkOption !== null) {
            body.think_option = thinkOption;
        }

        try {
            const response = await fetch("{{ url_for('admin_playground_stream') }}", {
                method: 'POST',
                headers: { 'Content-Type': 'application/json', 'X-CSRF-Token': '{{ csrf_token }}' },
                body: JSON.stringify(body),
                signal: abortController.signal
            });
            if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`);

            const reader = response.body.getReader();
            const decoder = new TextDecoder();

            while (true) {
                const shouldScroll = isScrolledToBottom(); // Check before DOM change
                const { value, done } = await reader.read();
                if (done) break;
                
                buffer += decoder.decode(value, { stream: true });
                const lines = buffer.split('\n');
                buffer = lines.pop(); // Keep partial line for next chunk

                for (const line of lines) {
                    if (line.trim() === '') continue;
                    try {
                        const data = JSON.parse(line);

                        if (firstChunk && data.message && typeof data.message.content === 'string') {
                            firstChunk = false;
                            botContentElement.innerHTML = ''; // Clear placeholder on first message with a content key
                        }

                        if (data.message?.content) { // Process only non-empty content for stats and rendering
                            if (currentTTFT === null) { // If TTFT not set, this is the first real token
                                const ttft = Math.round(performance.now() - startTime);
                                currentTTFT = ttft;
                                botStatsElement.innerHTML = `<span>TTFT: ${ttft}ms</span>`;
                                botStatsElement.classList.remove('hidden');
                            }
                            fullResponse += data.message.content;
                            renderBotMessage(botContentElement, fullResponse, lastSentImageURL);
                            highlightAndMathify(botContentElement);
                        }
                        
                        if (data.done && data.hasOwnProperty('eval_count') && data.hasOwnProperty('eval_duration')) {
                            const tokensPerSecond = data.eval_duration > 0 ? (data.eval_count / (data.eval_duration / 1e9)) : 0;
                            currentTPS = tokensPerSecond;
                            if (botStatsElement.innerHTML.includes('TTFT')) {
                                botStatsElement.innerHTML += ` | <span>Speed: ${tokensPerSecond.toFixed(2)} t/s</span>`;
                            } else {
                                botStatsElement.innerHTML = `<span>Speed: ${tokensPerSecond.toFixed(2)} t/s</span>`;
                                botStatsElement.classList.remove('hidden');
                            }
                        }

                        if (data.error) {
                            botContentElement.innerHTML += `<div class="p-3 bg-red-900/20 border border-red-500/50 rounded-lg text-red-400 mt-2"><strong>Orchestration Error:</strong> ${data.details || data.error}</div>`;
                            isLoading = false;
                            toggleLoading(false);
                        }
                    } catch (err) { 
                        if (!line.includes('processing')) {
                            console.warn("JSON parsing error in stream:", err, "Line:", line); 
                        }
                    }
                }
                if (shouldScroll) {
                    chatContainer.scrollTop = chatContainer.scrollHeight;
                }
            }
        } catch (error) {
            if (error.name === 'AbortError') {
                console.log('Stream cancelled by user.');
                streamError = null;
                fullResponse += "\n\n*[Generation stopped by user]*";
            } else {
                streamError = error;
                console.error('Error during chat stream:', error);
                botContentElement.innerHTML = `<p class="text-red-400"><strong>Error:</strong> ${error.message}</p>`;
                botStatsElement.classList.add('hidden');
            }
        } finally {
            toggleLoading(false);
            if (!streamError) {
                const newMessage = {
                    role: 'assistant',
                    content: fullResponse,
                    stats: (currentTTFT !== null && currentTPS !== null) ? { ttft: currentTTFT, tps: currentTPS } : null
                };
                if (fullResponse.includes('<annotate>')) {
                    newMessage.annotated_image_url = lastSentImageURL;
                }
                messages.push(newMessage);
                updateGlobalStats();
                updateCurrentChat();
            }
            renderChat(); 
            chatContainer.scrollTop = chatContainer.scrollHeight;
        }
    }
    
    async function handleChatSubmit(e) {
        if (e) e.preventDefault();
        const userInput = messageInput.value.trim();
        if ((!userInput && attachedImages.length === 0 && attachedDocs.length === 0) || isLoading) return;

        let content = [];
        let textForContent = userInput;
        
        if (attachedDocs.length > 0) {
            const fileContentPrefix = attachedDocs.map(doc => `File content from "${doc.name}":\n\`\`\`\n${doc.content}\n\`\`\``).join('\n\n');
            textForContent = `${fileContentPrefix}\n\n${textForContent}`;
        }
        
        if (textForContent) content.push({ type: 'text', text: textForContent });
        
        if (attachedImages.length > 0) {
            lastSentImageURL = attachedImages[attachedImages.length - 1]; // Store last image for annotation
            attachedImages.forEach(img => content.push({ type: 'image_url', image_url: { url: img }}));
        }
        
        messages.push({ role: 'user', content: content });
        updateCurrentChat();
        
        messageInput.value = '';
        messageInput.style.height = 'auto';
        attachedImages = [];
        attachedDocs = [];
        
        // --- UX FIX: Reset physical file inputs to allow re-selection of same files ---
        imageInput.value = "";
        docInput.value = "";
        
        renderImagePreviews();
        renderDocPreviews();

        await fetchStreamedResponse();
    }
    
    function renderBotMessage(container, rawContent, imageUrlForAnnotation) {
        container.dataset.rawContent = rawContent;
        container.innerHTML = window.LCPProcessor.render(rawContent, imageUrlForAnnotation);
    }
    
    function processAnnotations() {
        const annotationJson = part.slice(10, -11);
        const safeJson = annotationJson.replace(/'/g, "&apos;");

        if (!imageUrlForAnnotation) {
            finalHtml += `<div class="text-red-400">Error: Could not find an image to annotate for this message.</div>`;
        } else {
            const uniqueId = `annotation-${Date.now()}-${Math.random().toString(36).slice(2)}`;
            finalHtml += `<div id="${uniqueId}" class="annotated-image-container relative mt-2 group"
                data-annotations='${safeJson}' data-image-url="${imageUrlForAnnotation}">
                <img src="${imageUrlForAnnotation}" class="block max-w-full h-auto rounded-md">
                <div class="annotation-overlay"></div>
            </div>`;
        }
    }

    function renderChat() {
        chatContainer.innerHTML = '';
        if (messages.length === 0 && !isLoading) {
            chatContainer.innerHTML = '<div class="text-center text-gray-400">Select a model and start chatting.</div>';
            return;
        }
        messages.forEach((msg, index) => appendMessage(msg, index));
        if (isLoading) {
            appendMessage({role: 'assistant', content: ''}, messages.length, true);
        }
        // Force scroll to bottom on new render
        chatContainer.scrollTop = chatContainer.scrollHeight;
        highlightAndMathify(chatContainer);
        processAnnotations();
    }
    
    function appendMessage(msg, messageIndex, isPlaceholder = false) {
        const { role, content, stats } = msg;
        const messageWrapper = document.createElement('div');
        messageWrapper.className = `flex gap-4 message-wrapper group ${role === 'user' ? 'justify-end' : 'justify-start'}`;
        messageWrapper.dataset.messageIndex = messageIndex;

        const contentWrapper = document.createElement('div');
        // Theming the bubbles: User gets primary color, Assistant gets a tinted transparent background
        const bubbleBaseClass = "max-w-3xl w-full p-4 rounded-2xl relative shadow-sm transition-all";
        const bubbleThemeClass = role === 'user' 
            ? 'bg-[var(--color-primary-600)] text-white ml-12' 
            : 'bg-white/5 border border-white/10 mr-12';
            
        contentWrapper.className = `${bubbleBaseClass} ${bubbleThemeClass}`;
        
        const actionsDiv = document.createElement('div');
        actionsDiv.className = 'message-actions absolute -top-3 right-2 flex items-center gap-1 opacity-0 group-hover:opacity-100 transition-opacity';
        
        const contentDiv = document.createElement('div');
        contentDiv.className = 'message-content prose prose-invert max-w-none prose-p:my-2 prose-pre:my-2 prose-code:font-sans';

        // Detect HTML
        const isHtml = typeof content === 'string' && /<[a-z][\s\S]*>/i.test(content);
        if (isHtml) {
            const viewBtn = document.createElement('button');
            viewBtn.className = 'view-html-btn p-1 rounded-full bg-indigo-600 hover:bg-indigo-500 text-white';
            viewBtn.title = 'View in New Tab';
            viewBtn.innerHTML = '<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M10 6H6a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2v-4M14 4h6m0 0v6m0-6L10 14"></path></svg>';
            viewBtn.onclick = () => {
                const blob = new Blob([content], {type: 'text/html'});
                window.open(URL.createObjectURL(blob), '_blank');
            };
            actionsDiv.appendChild(viewBtn);
        }

        const statsDiv = document.createElement('div');
        statsDiv.className = 'message-stats text-xs text-gray-400 mt-2 space-x-2 hidden';

        if (isPlaceholder) {
            contentDiv.innerHTML = `<div class="flex items-center gap-2 text-gray-400"><svg class="animate-spin h-5 w-5" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24"><circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle><path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg><span>Waiting for response...</span></div>`;
        } else {
            let actionsHTML = `
                <button class="copy-btn p-1 rounded-full bg-gray-800 hover:bg-gray-700 text-gray-400 hover:text-white" title="Copy"><svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"></path></svg></button>
                <button class="edit-btn p-1 rounded-full bg-gray-800 hover:bg-gray-700 text-gray-400 hover:text-white" title="Edit"><svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15.232 5.232l3.536 3.536m-2.036-5.036a2.5 2.5 0 113.536 3.536L6.5 21.036H3v-3.5L15.232 5.232z"></path></svg></button>
                <button class="delete-btn p-1 rounded-full bg-gray-800 hover:bg-gray-700 text-gray-400 hover:text-white" title="Delete"><svg class="w-4 h-4" fill="currentColor" viewBox="0 0 24 24"><path d="M6 19c0 1.1.9 2 2 2h8c1.1 0 2-.9 2-2V7H6v12zM19 4h-3.5l-1-1h-5l-1 1H5v2h14V4z"></path></svg></button>`;
            
            if (role === 'assistant') {
                actionsHTML += `<button class="regenerate-btn p-1 rounded-full bg-gray-800 hover:bg-gray-700 text-gray-400 hover:text-white" title="Regenerate"><svg class="w-4 h-4" fill="currentColor" viewBox="0 0 24 24"><path d="M17.65 6.35C16.2 4.9 14.21 4 12 4c-4.42 0-7.99 3.58-7.99 8s3.57 8 7.99 8c3.73 0 6.84-2.55 7.73-6h-2.08c-.82 2.33-3.04 4-5.65 4-3.31 0-6-2.69-6-6s2.69-6 6-6c1.66 0 3.14.69 4.22 1.78L13 11h7V4l-2.35 2.35z"></path></svg></button>`;
                
                let rating = msg.rating || 0;
                let thumbUpClass = rating === 1 ? 'text-green-400 bg-green-500/20' : 'text-gray-400 bg-gray-800 hover:bg-gray-700';
                let thumbDownClass = rating === -1 ? 'text-red-400 bg-red-500/20' : 'text-gray-400 bg-gray-800 hover:bg-gray-700';
                
                actionsHTML += `
                    <div class="w-px h-4 bg-white/20 mx-1"></div>
                    <button class="thumb-up-btn p-1 rounded-full transition-colors ${thumbUpClass}" title="Good Response"><svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M14 10h4.764a2 2 0 011.789 2.894l-3.5 7A2 2 0 0115.263 21h-4.017c-.163 0-.326-.02-.485-.06L7 20m7-10V5a2 2 0 00-2-2h-.095c-.5 0-.905.405-.905.905 0 .714-.211 1.412-.608 2.006L7 11v9m7-10h-2M7 20H5a2 2 0 01-2-2v-6a2 2 0 012-2h2.514"></path></svg></button>
                    <button class="thumb-down-btn p-1 rounded-full transition-colors ${thumbDownClass}" title="Bad Response"><svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M10 14H5.236a2 2 0 01-1.789-2.894l3.5-7A2 2 0 018.736 3h4.018a2 2 0 01.485.06l3.76.94m-7 10v5a2 2 0 002 2h.096c.5 0 .905-.405.905-.904 0-.715.211-1.413.608-2.008L17 13V4m-7 10h2m5-10h2a2 2 0 012 2v6a2 2 0 01-2 2h-2.514"></path></svg></button>
                `;
            }
            actionsDiv.innerHTML = actionsHTML;

            if (role === 'assistant') {
                renderBotMessage(contentDiv, content, msg.annotated_image_url);
                if (stats && stats.ttft !== null && stats.tps !== null) {
                    statsDiv.innerHTML = `<span>TTFT: ${stats.ttft}ms</span> | <span>Speed: ${stats.tps.toFixed(2)} t/s</span>`;
                    statsDiv.classList.remove('hidden');
                }
            } else {
                let textContent = '';
                contentDiv.innerHTML = ''; 
                if (Array.isArray(content)) {
                    content.forEach(part => {
                        if (part.type === 'text' && part.text) {
                            textContent = part.text;
                            contentDiv.insertAdjacentHTML('beforeend', marked.parse(part.text));
                        } else if (part.type === 'image_url') {
                            const imageContainer = document.createElement('div');
                            imageContainer.className = 'image-message-container';
                            imageContainer.innerHTML = `<img src="${part.image_url.url}" class="block max-w-full h-auto rounded-md mt-2 cursor-pointer hover:ring-2 ring-indigo-500 transition-all" onclick="openPlaygroundPaint('${part.image_url.url}', ${messageIndex})">`;
                            contentDiv.appendChild(imageContainer);
                        }
                    });
                } else if (typeof content === 'string') {
                    textContent = content;
                    contentDiv.innerHTML = marked.parse(content);
                }
                contentDiv.dataset.rawContent = textContent;
            }

            contentDiv.querySelectorAll('.code-block-wrapper').forEach(wrapper => {
                const btnContainer = document.createElement('div');
                btnContainer.className = 'playground-code-actions';

                const codeEl = wrapper.querySelector('code');
                const lang = [...codeEl.classList].find(c => c.startsWith('language-'))?.replace('language-', '') || '';

                const copyBtn = document.createElement('button');
                copyBtn.className = 'playground-code-btn copy-code-btn';
                copyBtn.title = 'Copy code';
                copyBtn.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"></path></svg>`;

                const saveBtn = document.createElement('button');
                saveBtn.className = 'playground-code-btn save-code-btn';
                saveBtn.title = 'Save as file';
                saveBtn.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"></path></svg>`;

                btnContainer.appendChild(copyBtn);
                btnContainer.appendChild(saveBtn);

                if (['html', 'svg', 'mermaid'].includes(lang)) {
                    const viewBtn = document.createElement('button');
                    viewBtn.className = 'playground-code-btn view-code-btn';
                    viewBtn.title = 'View rendered';
                    viewBtn.innerHTML = `<svg class="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"/><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M2.458 12C3.732 7.943 7.523 5 12 5c4.478 0 8.268 2.943 9.542 7-1.274 4.057-5.064 7-9.542 7-4.477 0-8.268-2.943-9.542-7z"/></svg>`;
                    btnContainer.appendChild(viewBtn);
                }
                wrapper.appendChild(btnContainer);
            });
            renderMermaid(contentDiv);
        }

        contentWrapper.appendChild(actionsDiv);
        contentWrapper.appendChild(contentDiv);
        if (role === 'assistant') contentWrapper.appendChild(statsDiv);
        messageWrapper.appendChild(contentWrapper);
        chatContainer.appendChild(messageWrapper);
    }

    function processAnnotations() {
        chatContainer.querySelectorAll('.annotated-image-container').forEach(container => {
            if (container.dataset.processed) return;
            
            const annotationsJson = container.dataset.annotations.replace(/&apos;/g, "'");
            try {
                const annotations = JSON.parse(annotationsJson);
                const overlay = container.querySelector('.annotation-overlay');
                drawAnnotations(annotations, overlay);
                container.dataset.processed = "true";
            } catch (e) {
                console.error("Failed to parse or draw annotation JSON:", e);
                const errorP = document.createElement('p');
                errorP.className = 'text-red-400 text-xs italic mt-1';
                errorP.textContent = 'Error rendering annotations.';
                container.appendChild(errorP);
            }
        });
    }
    
    function drawAnnotations(annotations, overlay) {
        if (!overlay) return;
        overlay.innerHTML = '';

        // --- BACKWARD COMPATIBILITY SHIM ---
        // Transform old format {"box": [...], "class": "..."} to new format
        const transformedAnnotations = annotations.map(ann => {
            if (ann.box && ann.class) {
                return { type: 'bbox', bbox: ann.box, label: ann.class };
            }
            // If it has 'bbox' but no 'type', assume it's a bbox.
            if (ann.bbox && !ann.type) {
                ann.type = 'bbox';
            }
            return ann;
        });

        transformedAnnotations.forEach(ann => {
            const el = document.createElement('div');
            el.className = 'annotation-element';
            
            const display = ann.display || {};
            const borderColor = display.border_color || '#ff00ff';
            
            // Shared label logic
            let labelEl = '';
            if (ann.label) {
                labelEl = `<div class="annotation-label" style="background-color: ${borderColor};">${ann.label}</div>`;
            }

            switch (ann.type) {
                case 'bbox':
                case 'annotation': // Treat as bbox
                    const [x1, y1, x2, y2] = ann.bbox;
                    el.style.left = `${x1 * 100}%`;
                    el.style.top = `${y1 * 100}%`;
                    el.style.width = `${(x2 - x1) * 100}%`;
                    el.style.height = `${(y2 - y1) * 100}%`;
                    if (display.show_border !== false) {
                        el.style.border = `${display.border_width || 2}px solid ${borderColor}`;
                    }
                    el.style.backgroundColor = `rgba(255, 255, 255, ${display.fill_opacity || 0.1})`;
                    el.innerHTML = labelEl;
                    break;
                case 'point':
                    const [px, py] = ann.coords;
                    el.classList.add('point');
                    el.style.left = `${px * 100}%`;
                    el.style.top = `${py * 100}%`;
                    el.style.backgroundColor = borderColor;
                    el.innerHTML = labelEl;
                    break;
                case 'polygon':
                    // For simplicity, we'll draw a bounding box around the polygon points.
                    // True polygon drawing would require SVG.
                    const allX = ann.points.map(p => p[0]);
                    const allY = ann.points.map(p => p[1]);
                    const minX = Math.min(...allX);
                    const minY = Math.min(...allY);
                    const maxX = Math.max(...allX);
                    const maxY = Math.max(...allY);
                    el.style.left = `${minX * 100}%`;
                    el.style.top = `${minY * 100}%`;
                    el.style.width = `${(maxX - minX) * 100}%`;
                    el.style.height = `${(maxY - minY) * 100}%`;
                    el.style.border = `2px dashed ${borderColor}`;
                    el.innerHTML = labelEl;
                    break;
                case 'redaction':
                case 'replace':
                    const [rx1, ry1, rx2, ry2] = ann.bbox;
                    el.style.left = `${rx1 * 100}%`;
                    el.style.top = `${ry1 * 100}%`;
                    el.style.width = `${(rx2 - rx1) * 100}%`;
                    el.style.height = `${(ry2 - ry1) * 100}%`;
                    if (ann.action === 'blur') {
                        el.style.backdropFilter = `blur(${ann.strength || 5}px)`;
                    } else if (ann.type === 'replace') {
                        el.classList.add('text-overlay');
                        el.style.backgroundColor = ann.color || '#000000';
                        el.style.color = 'white';
                        el.textContent = ann.replacement_text;
                    }
                    break;
                 case 'translated_text':
                    const [tx1, ty1, tx2, ty2] = ann.bbox;
                    el.classList.add('text-overlay');
                    el.style.left = `${tx1 * 100}%`;
                    el.style.top = `${ty1 * 100}%`;
                    el.style.width = `${(tx2 - tx1) * 100}%`;
                    el.style.height = `${(ty2 - ty1) * 100}%`;
                    el.style.backgroundColor = 'white';
                    el.style.color = 'black';
                    el.textContent = ann.translated_text;
                    break;
                case 'text':
                    const [tx, ty] = ann.position;
                    el.style.left = `${tx * 100}%`;
                    el.style.top = `${ty * 100}%`;
                    el.textContent = ann.text;

                    if (ann.style) {
                        const allowedStyles = ['fontSize', 'fontWeight', 'color', 'backgroundColor', 'padding', 'borderRadius', 'textAlign', 'fontFamily'];
                        for (const [key, value] of Object.entries(ann.style)) {
                            const camelKey = key.replace(/_([a-z])/g, g => g[1].toUpperCase());
                            if (allowedStyles.includes(camelKey)) {
                                el.style[camelKey] = value;
                            }
                        }
                    }
                    if (!el.style.color) el.style.color = 'white';
                    if (!el.style.backgroundColor) el.style.backgroundColor = 'rgba(0,0,0,0.5)';
                    if (!el.style.padding) el.style.padding = '2px 4px';
                    if (!el.style.borderRadius) el.style.borderRadius = '3px';
                    el.style.whiteSpace = 'nowrap';
                    break;
            }
            overlay.appendChild(el);
        });
    }

    async function saveAnnotatedImage(imageUrl, container) {
        const image = container.querySelector('img');
        const annotationsJson = container.dataset.annotations.replace(/&apos;/g, "'");
        const annotations = JSON.parse(annotationsJson);
        
        const canvas = document.createElement('canvas');
        const ctx = canvas.getContext('2d');
        canvas.width = image.naturalWidth;
        canvas.height = image.naturalHeight;
        ctx.drawImage(image, 0, 0);

        // Draw annotations onto canvas
        annotations.forEach(ann => {
            // Simplified drawing logic for saving. Add more detail if needed.
            const display = ann.display || {};
            const color = display.border_color || '#ff00ff';
            ctx.strokeStyle = color;
            ctx.fillStyle = color;
            ctx.lineWidth = 2;
            
            if (ann.type === 'bbox' || ann.type === 'annotation') {
                const [x1, y1, x2, y2] = ann.bbox;
                ctx.strokeRect(x1 * canvas.width, y1 * canvas.height, (x2 - x1) * canvas.width, (y2 - y1) * canvas.height);
            } else if (ann.type === 'redaction' || ann.type === 'replace') {
                const [x1, y1, x2, y2] = ann.bbox;
                ctx.fillStyle = ann.color || '#000000';
                ctx.fillRect(x1 * canvas.width, y1 * canvas.height, (x2 - x1) * canvas.width, (y2 - y1) * canvas.height);
            }
        });

        const link = document.createElement('a');
        link.download = `annotated-image-${Date.now()}.png`;
        link.href = canvas.toDataURL('image/png');
        link.click();
    }
    
    function findLastBotMessageWrapper() {
        const wrappers = chatContainer.querySelectorAll('.message-wrapper');
        return wrappers[wrappers.length - 1];
    }

    function toggleLoading(state) {
        isLoading = state;
        messageInput.disabled = state;
        if (state) {
            sendBtn.innerHTML = `<svg class="w-5 h-5 mr-1" fill="currentColor" viewBox="0 0 20 20"><path d="M5 3a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2V5a2 2 0 00-2-2H5z"></path></svg> Stop`;
            sendBtn.classList.remove('bg-[var(--color-primary-600)]', 'hover:bg-[var(--color-primary-700)]');
            sendBtn.classList.add('bg-red-600', 'hover:bg-red-700');
            sendBtn.type = 'button';
        } else {
            sendBtn.innerHTML = 'Send';
            sendBtn.classList.add('bg-[var(--color-primary-600)]', 'hover:bg-[var(--color-primary-700)]');
            sendBtn.classList.remove('bg-red-600', 'hover:bg-red-700');
            sendBtn.type = 'submit';
            abortController = null;
        }
    }
    
    function updateGlobalStats() {
        const stats = messages.filter(m => m.role === 'assistant' && m.stats).map(m => m.stats);
        const meanTTFTEl = document.getElementById('mean-ttft');
        const meanTPSEl = document.getElementById('mean-tps');

        if (stats.length === 0) {
            meanTTFTEl.textContent = 'N/A';
            meanTPSEl.textContent = 'N/A';
            return;
        }

        const totalTTFT = stats.reduce((sum, s) => sum + s.ttft, 0);
        const totalTPS = stats.reduce((sum, s) => sum + s.tps, 0);

        const meanTTFT = totalTTFT / stats.length;
        const meanTPS = totalTPS / stats.length;

        meanTTFTEl.textContent = `${Math.round(meanTTFT)} ms`;
        meanTPSEl.textContent = `${meanTPS.toFixed(2)} t/s`;
    }

    function exportChat() {
        if (messages.length === 0) return;
        const blob = new Blob([JSON.stringify(messages, null, 2)], { type: 'application/json' });
        const url = URL.createObjectURL(blob);
        const a = document.createElement('a');
        a.href = url;
        a.download = `chat-history-${new Date().toISOString()}.json`;
        a.click();
        URL.revokeObjectURL(url);
    }

    function importChat(event) {
        const file = event.target.files[0];
        if (!file) return;
        const reader = new FileReader();
        reader.onload = (e) => {
            try {
                const imported = JSON.parse(e.target.result);
                if (Array.isArray(imported) && imported.every(m => m.role && m.content !== undefined)) {
                    messages = imported;
                    renderChat();
                    updateGlobalStats();
                } else { alert('Invalid chat file format.'); }
            } catch (error) { alert('Failed to read or parse chat file.'); }
        };
        reader.readAsText(file);
        event.target.value = null;
    }

    function copyToClipboard(text, button) {
        const originalContent = button.innerHTML;
        navigator.clipboard.writeText(text).then(() => {
            button.innerHTML = 'Copied!';
            setTimeout(() => { button.innerHTML = originalContent; }, 2000);
        }).catch(() => { // Fallback for http or other issues
            const textArea = document.createElement('textarea');
            textArea.value = text;
            textArea.style.position = 'fixed'; textArea.style.opacity = '0';
            document.body.appendChild(textArea);
            textArea.focus(); textArea.select();
            try {
                document.execCommand('copy');
                button.innerHTML = 'Copied!';
                setTimeout(() => { button.innerHTML = originalContent; }, 2000);
            } catch (err) { alert('Failed to copy'); }
            document.body.removeChild(textArea);
        });
    }
    
    function enterEditMode(messageWrapper, messageIndex) {
        const contentWrapper = messageWrapper.querySelector('.rounded-2xl');
        const contentDiv = contentWrapper.querySelector('.message-content');
        contentDiv.style.display = 'none';
        
        const actionsDiv = contentWrapper.querySelector('.message-actions');
        if (actionsDiv) actionsDiv.style.display = 'none';
        
        const statsDiv = contentWrapper.querySelector('.message-stats');
        if(statsDiv) statsDiv.style.display = 'none';

        const editContainer = document.createElement('div');
        editContainer.className = 'edit-container mt-2 w-full';
        const textarea = document.createElement('textarea');
        textarea.className = 'block w-full px-4 py-3 rounded-xl shadow-sm sm:text-sm bg-black/40 border border-white/20 text-white focus:ring-indigo-500 focus:border-indigo-500 custom-scrollbar outline-none';
        
        let rawTxt = contentDiv.dataset.rawContent || '';
        textarea.value = rawTxt;
        // Automatically calculate rows based on content
        const lineCount = (rawTxt.match(/\n/g) || []).length + 1;
        textarea.rows = Math.min(25, Math.max(5, lineCount + 2));
        
        const buttonGroup = document.createElement('div');
        buttonGroup.className = 'flex justify-end gap-3 mt-3';
        
        const saveAndSubmitBtn = document.createElement('button');
        saveAndSubmitBtn.textContent = 'Save & Submit';
        saveAndSubmitBtn.className = 'py-1.5 px-4 bg-indigo-600 text-white rounded-lg hover:bg-indigo-500 font-bold text-xs uppercase tracking-wider shadow-lg';
        saveAndSubmitBtn.onclick = () => saveAndSubmitEdit(messageIndex, textarea.value);
        
        const saveBtn = document.createElement('button');
        saveBtn.textContent = 'Save Only';
        saveBtn.className = 'py-1.5 px-4 bg-white/10 text-white rounded-lg hover:bg-white/20 font-bold text-xs uppercase tracking-wider';
        saveBtn.onclick = () => saveEdit(messageIndex, textarea.value);
        
        const cancelBtn = document.createElement('button');
        cancelBtn.textContent = 'Cancel';
        cancelBtn.className = 'py-1.5 px-4 bg-transparent text-gray-400 hover:text-white rounded-lg font-bold text-xs uppercase tracking-wider';
        cancelBtn.onclick = () => {
            editContainer.remove();
            contentDiv.style.display = 'block';
            if (actionsDiv) actionsDiv.style.display = 'flex';
            if (statsDiv && statsDiv.innerHTML.trim() !== '') statsDiv.style.display = 'block';
        };

        buttonGroup.appendChild(cancelBtn);
        buttonGroup.appendChild(saveBtn);
        buttonGroup.appendChild(saveAndSubmitBtn);
        editContainer.appendChild(textarea);
        editContainer.appendChild(buttonGroup);
        contentWrapper.appendChild(editContainer);
        textarea.focus();
    }
    
    function saveEdit(messageIndex, newContent) {
        const oldMessage = messages[messageIndex];
        
        // 1. Update the in-memory array
        if (Array.isArray(oldMessage.content)) {
            const textPart = oldMessage.content.find(p => p.type === 'text');
            if (textPart) {
                textPart.text = newContent;
            } else {
                oldMessage.content.unshift({type: 'text', text: newContent});
            }
        } else {
            oldMessage.content = newContent;
        }

        // 2. CRITICAL: Persist to LocalStorage
        if (window.updateCurrentChat) {
            window.updateCurrentChat();
        }

        // 3. Refresh UI
        renderChat();
    }
    
    function saveAndSubmitEdit(messageIndex, newContent) {
        saveEdit(messageIndex, newContent);
        messages.splice(messageIndex + 1); // Remove all subsequent messages
        updateCurrentChat();
        renderChat();
        fetchStreamedResponse();
    }
    
    function insertEditableMessage(role) {
        messages.push({ role: role, content: 'New message...' });
        renderChat();
        const newMessageWrapper = chatContainer.lastElementChild;
        if (newMessageWrapper) {
            enterEditMode(newMessageWrapper, messages.length - 1);
        }
    }
    
    function updateThinkingControls() {
        // Expand the list to match the backend updates
        const simpleThinkModels =['qwen', 'deepseek', 'glm', 'qwq', 'phi-4'];
        const gptOssKeyword = 'gpt-oss';
        const selectedModel = modelSelector.value.toLowerCase();

        const isGptOss = selectedModel.includes(gptOssKeyword);
        const isSimpleThink = simpleThinkModels.some(keyword => selectedModel.includes(keyword));

        thinkingToggleContainer.classList.add('hidden');
        thinkingLevelContainer.classList.add('hidden');

        if (isGptOss || isSimpleThink) {
            thinkingToggleContainer.classList.remove('hidden');
        }

        if (isGptOss && isThinkingMode) {
            thinkingLevelContainer.classList.remove('hidden');
        }
    }


    // --- Playground Paint Bridge ---
    window.openPlaygroundPaint = (url = null, idx = null) => {
        // Register the Playground-specific save behavior
        window.customPaintSaveHandler = (dataUrl, msgIndex) => {
            if (msgIndex !== null) {
                // Edit existing message in history
                const msg = messages[msgIndex];
                if (Array.isArray(msg.content)) {
                    const imgPart = msg.content.find(p => p.type === 'image_url');
                    if (imgPart) imgPart.image_url.url = dataUrl;
                }
                updateCurrentChat();
                renderChat();
            } else {
                // New attachment to draft
                attachedImages.push(dataUrl);
                renderImagePreviews();
            }
            // Reset handler after use
            window.customPaintSaveHandler = null;
        };

        // Open the global paint tool
        if (window.openCorePaint) window.openCorePaint(url, idx);
    };

    // --- Event Listeners & Init ---
    messageInput.addEventListener('input', () => { messageInput.style.height = 'auto'; messageInput.style.height = `${Math.min(messageInput.scrollHeight, 200)}px`; });
    messageInput.addEventListener('keydown', (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); chatForm.dispatchEvent(new Event('submit', { cancelable: true })); } });
    messageInput.addEventListener('paste', handlePaste);
    clearChatBtn.addEventListener('click', () => { 
        createNewChat();
    });
    document.getElementById('new-chat-btn').addEventListener('click', createNewChat);
    exportChatBtn.addEventListener('click', exportChat);
    importChatBtn.addEventListener('click', () => importChatInput.click());
    importChatInput.addEventListener('change', importChat);
    chatForm.addEventListener('submit', handleChatSubmit);
    testPromptsBtn.addEventListener('click', showTestPrompts);
    systemPromptBtn.addEventListener('click', showSystemPromptModal);
    systemPromptIndicator.addEventListener('click', showSystemPromptModal);
    attachImgBtn.addEventListener('click', () => imageInput.click());
    imageInput.addEventListener('change', (e) => handleImageFiles(e.target.files));
    attachDocBtn.addEventListener('click', () => docInput.click());
    document.getElementById('global-import-btn').addEventListener('click', () => {
        window.openKnowledgeImporter((title, content) => {
            attachedDocs.push({ name: title, content: content });
            renderDocPreviews();
        });
    });
    docInput.addEventListener('change', (e) => handleDocFiles(e.target.files));
    addMessageBtn.addEventListener('click', (e) => { e.stopPropagation(); addMessageDropdown.classList.toggle('hidden'); });
    addUserMsgBtn.addEventListener('click', (e) => { e.preventDefault(); insertEditableMessage('user'); addMessageDropdown.classList.add('hidden'); });
    addAiMsgBtn.addEventListener('click', (e) => { e.preventDefault(); insertEditableMessage('assistant'); addMessageDropdown.classList.add('hidden'); });
    document.addEventListener('click', (e) => { if (!addMessageMenuContainer.contains(e.target)) addMessageDropdown.classList.add('hidden'); });

    sendBtn.addEventListener('click', () => {
        if (isLoading && abortController) {
            abortController.abort();
        }
    });

    modelSelector.addEventListener('change', () => {
        localStorage.setItem(LAST_MODEL_KEY, modelSelector.value);
        updateThinkingControls();
    });
    
    annotationToggle.addEventListener('click', () => {
        isAnnotationMode = !isAnnotationMode;
        localStorage.setItem(ANNOTATION_MODE_KEY, isAnnotationMode);
        updateAnnotationToggleUI();
        updateSystemPromptIndicator();
    });
    
    thinkingToggle.addEventListener('click', () => {
        isThinkingMode = !isThinkingMode;
        localStorage.setItem(THINKING_MODE_KEY, isThinkingMode);
        updateThinkingToggleUI();
        updateThinkingControls();
    });

    document.addEventListener('click', async e => {
        const copyBtn = e.target.closest('.copy-btn');
        if (copyBtn) {
            const contentDiv = copyBtn.closest('.rounded-lg').querySelector('.message-content');
            copyToClipboard(contentDiv.dataset.rawContent || '', copyBtn); return;
        }
        
        const codeCopyBtn = e.target.closest('.copy-code-btn');
        if (codeCopyBtn) {
            const codeEl = codeCopyBtn.closest('.code-block-wrapper').querySelector('code');
            copyToClipboard(codeEl.innerText, codeCopyBtn); return;
        }
        
        const saveBtn = e.target.closest('.save-code-btn');
        const viewBtn = e.target.closest('.view-code-btn');
        if (saveBtn || viewBtn) {
            const wrapper = (saveBtn || viewBtn).closest('.code-block-wrapper');
            const codeEl = wrapper.querySelector('code');
            const codeContent = codeEl.innerText;
            const lang = [...codeEl.classList].find(c => c.startsWith('language-'))?.replace('language-', '') || 'txt';
            
            if (saveBtn) {
                const extension = langExtMap[lang.toLowerCase()] || lang || 'txt';
                const blob = new Blob([codeContent], { type: 'text/plain;charset=utf-8' });
                const url = URL.createObjectURL(blob);
                const a = document.createElement('a');
                a.href = url;
                a.download = `snippet.${extension}`;
                document.body.appendChild(a);
                a.click();
                document.body.removeChild(a);
                URL.revokeObjectURL(url);
            } else if (viewBtn) {
                if (lang === 'html' || lang === 'svg') {
                    const blob = new Blob([codeContent], { type: lang === 'html' ? 'text/html' : 'image/svg+xml' });
                    window.open(URL.createObjectURL(blob), '_blank');
                } else if (lang === 'mermaid') {
                    window.showModal("Rendered Diagram", `<div class="mermaid">${codeContent}</div>`);
                    mermaid.run();
                }
            }
            return;
        }

        const editBtn = e.target.closest('.edit-btn');
        if (editBtn) {
            const messageWrapper = editBtn.closest('.message-wrapper');
            enterEditMode(messageWrapper, parseInt(messageWrapper.dataset.messageIndex, 10)); return;
        }
        const deleteBtn = e.target.closest('.delete-btn');
        if (deleteBtn) {
            if (isLoading) return;
            const messageWrapper = deleteBtn.closest('.message-wrapper');
            const messageIndex = parseInt(messageWrapper.dataset.messageIndex, 10);
            if (!isNaN(messageIndex) && confirm('Are you sure you want to delete this message and all subsequent messages?')) {
                messages.splice(messageIndex); 
                updateCurrentChat();
                renderChat(); 
                updateGlobalStats();
            } return;
        }
        const regenerateBtn = e.target.closest('.regenerate-btn');
        if (regenerateBtn) {
            if (isLoading) return;
            const messageWrapper = regenerateBtn.closest('.message-wrapper');
            const messageIndex = parseInt(messageWrapper.dataset.messageIndex, 10);
            if (!isNaN(messageIndex)) {
                messages.splice(messageIndex); 
                updateCurrentChat();
                fetchStreamedResponse();
            } return;
        }
        
        const thumbUpBtn = e.target.closest('.thumb-up-btn');
        if (thumbUpBtn) {
            const messageWrapper = thumbUpBtn.closest('.message-wrapper');
            const idx = parseInt(messageWrapper.dataset.messageIndex, 10);
            messages[idx].rating = messages[idx].rating === 1 ? 0 : 1;
            updateCurrentChat();
            renderChat();
            return;
        }

        const thumbDownBtn = e.target.closest('.thumb-down-btn');
        if (thumbDownBtn) {
            const messageWrapper = thumbDownBtn.closest('.message-wrapper');
            const idx = parseInt(messageWrapper.dataset.messageIndex, 10);
            messages[idx].rating = messages[idx].rating === -1 ? 0 : -1;
            updateCurrentChat();
            renderChat();
            return;
        }
        const saveAnnotationBtn = e.target.closest('.save-annotation-btn');
        if (saveAnnotationBtn) {
            const container = saveAnnotationBtn.closest('.annotated-image-container');
            const imageUrl = container.dataset.imageUrl;
            await saveAnnotatedImage(imageUrl, container);
            return;
        }
        
        const removeDocBtn = e.target.closest('.remove-doc-btn');
        if (removeDocBtn) {
            const index = parseInt(removeDocBtn.dataset.index, 10);
            attachedDocs.splice(index, 1);
            renderDocPreviews();
            return;
        }

        if (e.target.classList.contains('use-prompt-btn')) {
            messageInput.value = unescape(e.target.dataset.prompt);
            document.getElementById('modal-close-btn').click();
            messageInput.focus(); messageInput.dispatchEvent(new Event('input'));
        }
        if (e.target.id === 'save-new-prompt-btn') {
            const [category, title, prompt, expected] = ['#new-prompt-category', '#new-prompt-title', '#new-prompt-prompt', '#new-prompt-expected'].map(id => document.getElementById(id.substring(1)).value.trim());
            if (category && title && prompt && expected) {
                const customPrompts = JSON.parse(localStorage.getItem(CUSTOM_PROMPTS_KEY) || '[]');
                customPrompts.push({ category, title, prompt, expected_outcome: expected });
                localStorage.setItem(CUSTOM_PROMPTS_KEY, JSON.stringify(customPrompts));
                showTestPrompts();
            } else { alert('All fields are required.'); }
        }
        if (e.target.classList.contains('delete-prompt-btn')) {
            const titleToDelete = unescape(e.target.dataset.promptTitle);
            if (confirm(`Delete custom prompt "${titleToDelete}"?`)) {
                let customPrompts = JSON.parse(localStorage.getItem(CUSTOM_PROMPTS_KEY) || '[]');
                customPrompts = customPrompts.filter(p => p.title !== titleToDelete);
                localStorage.setItem(CUSTOM_PROMPTS_KEY, JSON.stringify(customPrompts));
                showTestPrompts();
            }
        }
        if (e.target.id === 'load-system-prompt-btn') {
            document.getElementById('system-prompt-dropdown').classList.toggle('hidden');
        }
        if (e.target.classList.contains('system-prompt-item')) {
            e.preventDefault();
            document.getElementById('system-prompt-textarea').value = unescape(e.target.dataset.prompt);
            document.getElementById('system-prompt-dropdown').classList.add('hidden');
        }
        if (e.target.id === 'apply-system-prompt-btn') {
            userSystemPrompt = document.getElementById('system-prompt-textarea').value.trim();
            localStorage.setItem(LAST_SYSTEM_PROMPT_KEY, userSystemPrompt);
            updateSystemPromptIndicator();
            document.getElementById('modal-close-btn').click();
        }
        if (e.target.id === 'clear-system-prompt-btn') {
            userSystemPrompt = '';
            localStorage.removeItem(LAST_SYSTEM_PROMPT_KEY);
            updateSystemPromptIndicator();
            document.getElementById('modal-close-btn').click();
        }
        if (e.target.id === 'save-custom-system-prompt-btn') {
            const name = document.getElementById('new-system-prompt-name').value.trim();
            const prompt = document.getElementById('system-prompt-textarea').value.trim();
            if (name && prompt) {
                const custom = JSON.parse(localStorage.getItem(CUSTOM_SYSTEM_PROMPTS_KEY) || '[]');
                if (custom.some(p => p.name === name)) { alert('A preset with this name already exists.'); return; }
                custom.push({ name, prompt });
                localStorage.setItem(CUSTOM_SYSTEM_PROMPTS_KEY, JSON.stringify(custom));
                showSystemPromptModal();
            } else { alert('Please provide a name and a prompt to save.'); }
        }
        if (e.target.classList.contains('delete-system-prompt-btn')) {
            const nameToDelete = unescape(e.target.dataset.name);
            if (confirm(`Delete system prompt preset "${nameToDelete}"?`)) {
                let custom = JSON.parse(localStorage.getItem(CUSTOM_SYSTEM_PROMPTS_KEY) || '[]');
                custom = custom.filter(p => p.name !== nameToDelete);
                localStorage.setItem(CUSTOM_SYSTEM_PROMPTS_KEY, JSON.stringify(custom));
                showSystemPromptModal();
            }
        }
    });

    // --- Init ---
    // PRIORITY: 1. URL Param, 2. LocalStorage, 3. Default
    // --- Mermaid & Syntax Init ---
    mermaid.initialize({ startOnLoad: false, theme: 'dark' });

    // --- Init ---
    const urlParams = new URLSearchParams(window.location.search);
    const modelFromUrl = urlParams.get('model');
    
    if (modelFromUrl && modelSelector.querySelector(`option[value="${modelFromUrl}"]`)) {
        modelSelector.value = modelFromUrl;
        localStorage.setItem(LAST_MODEL_KEY, modelFromUrl);
        // FORCE: If coming from Architect, create a fresh chat for this model immediately
        createNewChat();
    } else {
        const lastModel = localStorage.getItem(LAST_MODEL_KEY);
        if (lastModel && modelSelector.querySelector(`option[value="${lastModel}"]`)) {
            modelSelector.value = lastModel;
        }
    }

    userSystemPrompt = localStorage.getItem(LAST_SYSTEM_PROMPT_KEY) || '';
    isAnnotationMode = localStorage.getItem(ANNOTATION_MODE_KEY) === 'true';
    isThinkingMode = localStorage.getItem(THINKING_MODE_KEY) === 'true';
    updateAnnotationToggleUI();
    updateThinkingToggleUI();
    
    // Only perform default loading if we DIDN'T just force a new chat via URL
    if (!modelFromUrl) {
        renderHistoryList();
        if (chatHistory.length > 0) {
            loadChat(chatHistory[0].id);
        } else {
            createNewChat();
        }
    } else {
        renderHistoryList();
    }
    
    updateSystemPromptIndicator();
    updateGlobalStats();
    updateThinkingControls(); // Initial check for think controls
});
</script>
{% endblock %}