ChatUI

Running on Zero

ChatUI Studio v2.0 — FastAPI + Transformers, multi-model / multi-task / agent
d2bc8de 2 days ago
45.3 kB
// ── Global Application State ──────────────────────────────────────────
const STATE = {
    // Model & task
    models: [],              // registry returned by backend
    tasks: [],               // registry returned by backend
    currentModel: null,      // model id string
    currentTask: "chat",     // task id string
    modelCache: new Set(),   // models already loaded on the server

    // Generation params
    maxTokens: 512,
    temperature: 0.7,
    systemPrompt: "",

    // Conversation
    uploadedFiles: [],
    conversationHistory: [],
    isThinking: false,

    // UI toggles
    rpMode: true,
    autoHideThinking: true,
    showAgentTrace: true,

    // Sessions
    currentSessionId: null,
    chatSessions: [],
};

const STORAGE_KEY = "chatui_studio_sessions_v1";

// ── Markdown Configuration ────────────────────────────────────────────
marked.setOptions({
    breaks: true,
    highlight: function (code, lang) {
        const language = hljs.getLanguage(lang) ? lang : 'plaintext';
        return hljs.highlight(code, { language }).value;
    }
});

// ── RP (Roleplay) Markdown Processor (kept from original) ─────────────
function processRpMarkdown(text) {
    if (!text) return text;

    const codeBlocks = [];
    text = text.replace(/(```[\s\S]*?```|`[^`]+`)/g, (match) => {
        codeBlocks.push(match);
        return `%%CODEBLOCK_${codeBlocks.length - 1}%%`;
    });

    const htmlBlocks = [];
    text = text.replace(/(<[^>]+>)/g, (match) => {
        htmlBlocks.push(match);
        return `%%HTMLBLOCK_${htmlBlocks.length - 1}%%`;
    });

    text = text.replace(/~~([^~]+)~~/g, '<span class="rp-ooc">$1</span>');
    text = text.replace(/\*\*([^*]+)\*\*/g, '<strong class="rp-bold">$1</strong>');
    text = text.replace(/(?<!\*)\*(?!\*)([^*\n]+)\*(?!\*)/g, '<em class="rp-narration">$1</em>');
    text = text.replace(/[\u201C\u201E"]([^"\u201D\n]+)[\u201D"]([^"\u201C\u201E]|$)/g, (match, dialogue, after) => {
        return `<span class="rp-dialogue">\u201C${dialogue}\u201D</span>${after}`;
    });
    text = text.replace(/^(---|\*\*\*|___)$/gm, '<hr class="rp-divider">');
    text = text.replace(/\(\(([^)]+)\)\)/g, '<span class="rp-ooc">(($1))</span>');

    htmlBlocks.forEach((block, i) => text = text.replace(`%%HTMLBLOCK_${i}%%`, block));
    codeBlocks.forEach((block, i) => text = text.replace(`%%CODEBLOCK_${i}%%`, block));
    return text;
}

function renderMarkdown(content) {
    if (!content) return "";
    if (STATE.rpMode) {
        const rpProcessed = processRpMarkdown(content);
        return marked.parse(rpProcessed);
    }
    return marked.parse(content);
}

// ── Session persistence (localStorage) ────────────────────────────────
function getStoredSessions() {
    try {
        const data = localStorage.getItem(STORAGE_KEY);
        return data ? JSON.parse(data) : { sessions: {}, activeSession: null };
    } catch (e) {
        return { sessions: {}, activeSession: null };
    }
}

function saveStoredSessions(storageData) {
    try {
        localStorage.setItem(STORAGE_KEY, JSON.stringify(storageData));
    } catch (e) {
        console.warn("localStorage save failed:", e);
    }
}

function loadChatHistory() {
    const storage = getStoredSessions();
    STATE.chatSessions = Object.entries(storage.sessions)
        .map(([id, sess]) => ({
            id: id,
            title: sess.title || "New Chat",
            updated_at: sess.updated_at || 0,
            message_count: (sess.messages || []).length,
            task: sess.task || "chat",
            model: sess.model || STATE.currentModel || "",
        }))
        .sort((a, b) => b.updated_at - a.updated_at);
    renderChatHistory();
}

function saveCurrentSession() {
    if (STATE.conversationHistory.length === 0) return;

    const storage = getStoredSessions();
    const sid = STATE.currentSessionId || ("sess-" + Math.random().toString(36).substring(2, 9));

    let title = "New Chat";
    for (const m of STATE.conversationHistory) {
        if (m.role === "user") {
            const t = typeof m.content === "string" ? m.content : "New Chat";
            title = t.length > 60 ? t.substring(0, 60) + "..." : t;
            break;
        }
    }

    storage.sessions[sid] = {
        title: title,
        messages: STATE.conversationHistory,
        updated_at: Date.now() / 1000,
        task: STATE.currentTask,
        model: STATE.currentModel,
    };
    storage.activeSession = sid;
    saveStoredSessions(storage);

    STATE.currentSessionId = sid;
    loadChatHistory();
}

function loadSession(sessionId) {
    const storage = getStoredSessions();
    const sess = storage.sessions[sessionId];
    if (!sess) return;

    STATE.currentSessionId = sessionId;
    STATE.conversationHistory = sess.messages || [];
    if (sess.task && STATE.tasks.find(t => t.id === sess.task)) {
        STATE.currentTask = sess.task;
        renderTaskPills();
    }
    if (sess.model && STATE.models.find(m => m.id === sess.model)) {
        STATE.currentModel = sess.model;
        renderModelPicker();
        renderModelSelect();
    }
    rebuildChatFromHistory(sess.messages || []);
    renderChatHistory();
    toggleLeftSidebar(false);
}

function deleteSession(sessionId) {
    const storage = getStoredSessions();
    delete storage.sessions[sessionId];
    if (storage.activeSession === sessionId) storage.activeSession = null;
    saveStoredSessions(storage);

    if (sessionId === STATE.currentSessionId) {
        resetSandbox();
    }
    loadChatHistory();
}

function renderChatHistory() {
    const listEl = document.getElementById("chat-history-list");
    if (!listEl) return;

    if (STATE.chatSessions.length === 0) {
        listEl.innerHTML = '<div class="empty-history-text">No chat history yet.<br>Start a conversation!</div>';
        return;
    }

    listEl.innerHTML = "";
    STATE.chatSessions.forEach(session => {
        const item = document.createElement("div");
        item.className = "chat-history-item" + (session.id === STATE.currentSessionId ? " active" : "");
        const timeAgo = formatTimeAgo(session.updated_at);
        const taskIcon = (TASK_ICONS[session.task] || "💬");

        item.innerHTML = `
            <div class="history-item-content" data-session-id="${session.id}">
                <div class="history-item-title">${taskIcon} ${escapeHtml(session.title)}</div>
                <div class="history-item-meta">${session.message_count} msgs · ${timeAgo}</div>
            </div>
            <button class="history-item-delete" data-session-id="${session.id}" title="Delete">
                <svg viewBox="0 0 24 24" width="12" height="12" stroke="currentColor" stroke-width="2" fill="none" stroke-linecap="round" stroke-linejoin="round"><line x1="18" y1="6" x2="6" y2="18"></line><line x1="6" y1="6" x2="18" y2="18"></line></svg>
            </button>
        `;

        item.querySelector(".history-item-content").addEventListener("click", () => loadSession(session.id));
        item.querySelector(".history-item-delete").addEventListener("click", (e) => {
            e.stopPropagation();
            deleteSession(session.id);
        });

        listEl.appendChild(item);
    });
}

const TASK_ICONS = {
    chat: "💬", code: "💻", summarize: "📝", translate: "🌍", agent: "🤖"
};

function formatTimeAgo(timestamp) {
    const now = Date.now() / 1000;
    const diff = now - timestamp;
    if (diff < 60) return "just now";
    if (diff < 3600) return Math.floor(diff / 60) + "m ago";
    if (diff < 86400) return Math.floor(diff / 3600) + "h ago";
    return Math.floor(diff / 86400) + "d ago";
}

function rebuildChatFromHistory(messages) {
    const feed = document.getElementById("chat-messages-feed");
    const chatContainer = document.getElementById("chat-thread-container");
    const dashboard = document.getElementById("studio-dashboard");

    feed.innerHTML = "";

    if (messages.length === 0) {
        chatContainer.style.display = "none";
        dashboard.style.display = "flex";
        return;
    }

    dashboard.style.display = "none";
    chatContainer.style.display = "flex";

    messages.forEach(msg => {
        if (msg.role === "user") {
            const text = typeof msg.content === "string" ? msg.content : (msg.content || []).map(c => c.text || "").join(" ");
            appendUserBubble(text, []);
        } else if (msg.role === "assistant") {
            appendAssistantBubble(msg.content, msg.trace || null);
        }
    });

    scrollToBottom();
}

function appendAssistantBubble(content, trace) {
    const id = "assistant-" + Math.random().toString(36).substring(2, 9);
    const bubble = document.createElement("div");
    bubble.className = "message-bubble assistant";
    bubble.id = id;

    const label = STATE.currentModel ? STATE.currentModel.split("/").pop() : "Assistant";

    let traceHtml = "";
    if (trace && Array.isArray(trace) && trace.length && STATE.showAgentTrace) {
        traceHtml = renderAgentTraceHtml(trace, id);
    }

    bubble.innerHTML = `
        <div class="message-meta">${escapeHtml(label)}</div>
        <div class="message-body">
            ${traceHtml}
            <div class="message-text markdown-body" id="${id}-text-box" data-raw-content="${escapeAttr(content)}">
                ${renderMarkdown(content)}
            </div>
        </div>
    `;

    const feed = document.getElementById("chat-messages-feed");
    feed.appendChild(bubble);

    const textBox = document.getElementById(`${id}-text-box`);
    if (textBox) {
        textBox.querySelectorAll("pre code").forEach((el) => hljs.highlightElement(el));
        addCopyButtons(textBox);
    }
    attachAgentTraceToggle(id);
}

function renderAgentTraceHtml(trace, id) {
    if (!trace || !trace.length) return "";
    const stepsHtml = trace.map((step, i) => {
        const tagClass = step.type;
        const tagLabel = step.type.replace(/_/g, " ");
        let body = "";
        if (step.type === "tool_call") {
            body = `tool: ${escapeHtml(step.tool || "")}\narg: ${escapeHtml(step.argument || "")}`;
        } else if (step.type === "tool_result") {
            body = `tool: ${escapeHtml(step.tool || "")}\nresult: ${escapeHtml(step.output || "")}`;
        } else {
            body = escapeHtml(step.content || "");
        }
        return `
            <div class="agent-trace-step">
                <span class="agent-step-tag ${tagClass}">Step ${i + 1} · ${tagLabel}</span>
                <div class="agent-step-content">${body}</div>
            </div>
        `;
    }).join("");

    return `
        <div class="agent-trace" id="${id}-trace">
            <div class="agent-trace-header" id="${id}-trace-toggle">
                <span>🤖 Agent Trace (${trace.length} steps)</span>
                <span class="thought-toggle-icon">▼</span>
            </div>
            <div class="agent-trace-body" id="${id}-trace-body">${stepsHtml}</div>
        </div>
    `;
}

function attachAgentTraceToggle(id) {
    const header = document.getElementById(`${id}-trace-toggle`);
    const body = document.getElementById(`${id}-trace-body`);
    if (!header || !body) return;
    header.addEventListener("click", () => {
        const hidden = body.style.display === "none";
        body.style.display = hidden ? "block" : "none";
        header.querySelector(".thought-toggle-icon").textContent = hidden ? "▼" : "▶";
    });
}

// ── Initialize Application ────────────────────────────────────────────
async function initializeApp() {
    // 1. Fetch model + task registries from backend
    fetchRegistries();

    // 2. Load chat history from localStorage
    loadChatHistory();

    // 3. Left sidebar events
    bindClick("btn-toggle-left", () => toggleLeftSidebar(true));
    bindClick("btn-close-left", () => toggleLeftSidebar(false));
    bindClick("sidebar-overlay-left", () => toggleLeftSidebar(false));

    // 4. Right drawer events
    bindClick("btn-toggle-right", () => toggleSettingsDrawer(true));
    bindClick("btn-close-drawer", () => toggleSettingsDrawer(false));
    bindClick("sidebar-overlay", () => toggleSettingsDrawer(false));

    // 5. Clear history
    bindClick("clear-history-btn", () => {
        if (confirm("Clear all chat history? This cannot be undone.")) {
            localStorage.removeItem(STORAGE_KEY);
            STATE.chatSessions = [];
            renderChatHistory();
        }
    });

    // 6. Settings listeners
    const maxTokensSlider = document.getElementById("max-tokens-slider");
    const maxTokensVal = document.getElementById("max-tokens-val");
    if (maxTokensSlider) {
        maxTokensSlider.addEventListener("input", (e) => {
            STATE.maxTokens = parseInt(e.target.value);
            maxTokensVal.textContent = STATE.maxTokens;
        });
    }

    const temperatureSlider = document.getElementById("temperature-slider");
    const temperatureVal = document.getElementById("temperature-val");
    if (temperatureSlider) {
        temperatureSlider.addEventListener("input", (e) => {
            STATE.temperature = parseFloat(e.target.value);
            temperatureVal.textContent = STATE.temperature.toFixed(1);
        });
    }

    const systemPromptInput = document.getElementById("system-prompt-input");
    if (systemPromptInput) {
        systemPromptInput.addEventListener("input", (e) => {
            STATE.systemPrompt = e.target.value;
        });
    }

    // Toggles
    bindToggle("rp-mode-toggle", v => { STATE.rpMode = v; rerenderAllAssistantMessages(); }, STATE.rpMode);
    bindToggle("auto-hide-thinking-toggle", v => {
        STATE.autoHideThinking = v;
        document.querySelectorAll(".thought-container").forEach(tc => {
            tc.classList.toggle("collapsed", v);
            const iconEl = tc.querySelector(".thought-toggle-icon");
            if (iconEl) iconEl.textContent = v ? "▶" : "▼";
        });
    }, STATE.autoHideThinking);
    bindToggle("agent-trace-toggle", v => {
        STATE.showAgentTrace = v;
        document.querySelectorAll(".agent-trace").forEach(el => {
            el.style.display = v ? "" : "none";
        });
    }, STATE.showAgentTrace);

    // 7. File upload
    const fileUploader = document.getElementById("file-uploader");
    bindClick("studio-upload-trigger", () => fileUploader && fileUploader.click());
    bindClick("mini-upload-trigger", () => fileUploader && fileUploader.click());
    if (fileUploader) fileUploader.addEventListener("change", handleFileSelection);

    // 8. Submit triggers
    bindClick("studio-send-button", () => {
        triggerPromptSubmission(document.getElementById("studio-prompt-input"));
    });
    bindClick("mini-send-button", () => {
        triggerPromptSubmission(document.getElementById("mini-prompt-input"));
    });

    const studioPromptInput = document.getElementById("studio-prompt-input");
    const miniPromptInput = document.getElementById("mini-prompt-input");
    [studioPromptInput, miniPromptInput].forEach(t => {
        if (!t) return;
        t.addEventListener("keydown", (e) => {
            if (e.key === "Enter" && !e.shiftKey) {
                e.preventDefault();
                triggerPromptSubmission(t);
            }
        });
        t.addEventListener("input", () => {
            t.style.height = "auto";
            t.style.height = t.scrollHeight + "px";
        });
    });

    // 9. New chat
    bindClick("menu-new-chat", resetSandbox);
    bindClick("sidebar-new-chat", () => {
        resetSandbox();
        toggleLeftSidebar(false);
    });

    // 10. Recipe chips
    document.querySelectorAll(".recipe-chip").forEach(chip => {
        chip.addEventListener("click", () => loadRecipe(chip.getAttribute("data-recipe")));
    });

    // 11. Model picker dropdown
    bindClick("model-picker-trigger", (e) => {
        e.stopPropagation();
        const menu = document.getElementById("model-picker-menu");
        if (menu) menu.hidden = !menu.hidden;
    });
    document.addEventListener("click", (e) => {
        const menu = document.getElementById("model-picker-menu");
        const trigger = document.getElementById("model-picker-trigger");
        if (menu && !menu.hidden && !menu.contains(e.target) && e.target !== trigger) {
            menu.hidden = true;
        }
    });

    // 12. Active model select + warm up
    const modelSelect = document.getElementById("active-model-select");
    if (modelSelect) {
        modelSelect.addEventListener("change", (e) => {
            STATE.currentModel = e.target.value;
            renderModelPicker();
            updateModelMeta();
        });
    }
    bindClick("warm-up-btn", warmUpCurrentModel);
}

function bindClick(id, fn) {
    const el = document.getElementById(id);
    if (el) el.addEventListener("click", fn);
}

function bindToggle(id, fn, initial) {
    const el = document.getElementById(id);
    if (!el) return;
    el.checked = !!initial;
    el.addEventListener("change", (e) => fn(e.target.checked));
}

// ── Registry fetch ────────────────────────────────────────────────────
async function fetchRegistries() {
    try {
        const [modelsRes, tasksRes] = await Promise.all([
            fetch("/api/models").then(r => r.json()),
            fetch("/api/tasks").then(r => r.json()),
        ]);

        if (modelsRes.status === "success") {
            STATE.models = modelsRes.models;
            STATE.currentModel = modelsRes.default || (STATE.models[0] && STATE.models[0].id);
            (modelsRes.cached || []).forEach(m => STATE.modelCache.add(m));
            renderModelPicker();
            renderModelSelect();
            updateModelMeta();
        }
        if (tasksRes.status === "success") {
            STATE.tasks = tasksRes.tasks;
            renderTaskPills();
        }
    } catch (e) {
        console.error("Failed to fetch registries:", e);
        setModelStatus("error", "Backend offline");
    }
}

function renderModelPicker() {
    const trigger = document.getElementById("model-picker-trigger");
    const nameEl = document.getElementById("model-picker-name");
    const dotEl = document.getElementById("model-picker-dot");
    const menu = document.getElementById("model-picker-menu");

    if (!STATE.models.length) {
        if (nameEl) nameEl.textContent = "No models";
        if (dotEl) dotEl.className = "model-picker-dot error";
        return;
    }

    const current = STATE.models.find(m => m.id === STATE.currentModel) || STATE.models[0];
    if (nameEl) nameEl.textContent = current.label || current.id;
    if (dotEl) {
        const cached = STATE.modelCache.has(current.id);
        dotEl.className = "model-picker-dot " + (cached ? "loaded" : "loading");
    }

    if (!menu) return;
    menu.innerHTML = "";
    STATE.models.forEach(m => {
        const option = document.createElement("div");
        option.className = "model-option" + (m.id === STATE.currentModel ? " active" : "");
        const cached = STATE.modelCache.has(m.id);
        option.innerHTML = `
            <div class="model-option-radio"></div>
            <div class="model-option-text">
                <div class="model-option-label">${escapeHtml(m.label || m.id)}</div>
                <div class="model-option-id">${escapeHtml(m.id)}</div>
                ${cached ? '<span class="model-option-badge cached">Loaded</span>' : ""}
                <span class="model-option-badge">ctx ${m.context || "—"}</span>
            </div>
        `;
        option.addEventListener("click", () => {
            STATE.currentModel = m.id;
            menu.hidden = true;
            renderModelPicker();
            renderModelSelect();
            updateModelMeta();
        });
        menu.appendChild(option);
    });
}

function renderModelSelect() {
    const sel = document.getElementById("active-model-select");
    if (!sel) return;
    sel.innerHTML = "";
    STATE.models.forEach(m => {
        const opt = document.createElement("option");
        opt.value = m.id;
        opt.textContent = m.label || m.id;
        if (m.id === STATE.currentModel) opt.selected = true;
        sel.appendChild(opt);
    });
}

function updateModelMeta() {
    const meta = document.getElementById("model-meta");
    if (!meta) return;
    const m = STATE.models.find(x => x.id === STATE.currentModel);
    if (!m) { meta.textContent = "No model selected."; return; }
    const cached = STATE.modelCache.has(m.id);
    meta.innerHTML = `<strong>${escapeHtml(m.label)}</strong><br>` +
                     `id: <code>${escapeHtml(m.id)}</code><br>` +
                     `context: ${m.context || "—"} tokens<br>` +
                     `state: ${cached ? "loaded in memory" : "loads on first use"}`;
    setModelStatus(cached ? "loaded" : "idle", cached ? "Model loaded" : "Idle");
    const footerLabel = document.getElementById("footer-model-label");
    if (footerLabel) footerLabel.textContent = (m.label || m.id) + (cached ? " · ready" : " · cold");
}

function setModelStatus(state, text) {
    const dot = document.querySelector("#model-status .status-dot");
    const txt = document.getElementById("model-status-text");
    if (dot) dot.className = "status-dot " + (state === "loaded" ? "online" : state === "error" ? "error" : "online");
    if (txt) txt.textContent = text;
}

function renderTaskPills() {
    const wrap = document.getElementById("task-pills");
    if (!wrap) return;
    wrap.innerHTML = "";
    STATE.tasks.forEach(t => {
        const btn = document.createElement("button");
        btn.className = "task-pill" + (t.id === STATE.currentTask ? " active" : "");
        btn.textContent = t.label || t.id;
        btn.title = t.description || "";
        btn.addEventListener("click", () => {
            STATE.currentTask = t.id;
            renderTaskPills();
        });
        wrap.appendChild(btn);
    });
}

async function warmUpCurrentModel() {
    if (!STATE.currentModel) return;
    const btn = document.getElementById("warm-up-btn");
    if (btn) { btn.disabled = true; btn.textContent = "Loading…"; }
    setModelStatus("loading", "Loading model…");
    const dot = document.getElementById("model-picker-dot");
    if (dot) dot.className = "model-picker-dot loading";
    try {
        const res = await fetch("/api/load_model", {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({ model_id: STATE.currentModel }),
        });
        const data = await res.json();
        if (data.status === "success") {
            STATE.modelCache.add(STATE.currentModel);
            renderModelPicker();
            updateModelMeta();
        } else {
            alert("Failed to load model: " + (data.message || "unknown error"));
            setModelStatus("error", "Load failed");
            if (dot) dot.className = "model-picker-dot error";
        }
    } catch (e) {
        alert("Error: " + e.message);
        setModelStatus("error", "Load failed");
        if (dot) dot.className = "model-picker-dot error";
    } finally {
        if (btn) { btn.disabled = false; btn.textContent = "Pre-load now"; }
    }
}

// ── Sidebar toggles ───────────────────────────────────────────────────
function toggleLeftSidebar(open) {
    const sidebar = document.getElementById("sidebar-left");
    const overlay = document.getElementById("sidebar-overlay-left");
    sidebar.classList.toggle("open", open);
    overlay.classList.toggle("active", open);
}

function toggleSettingsDrawer(open) {
    const sidebar = document.getElementById("sidebar-right");
    const overlay = document.getElementById("sidebar-overlay");
    sidebar.classList.toggle("collapsed", !open);
    overlay.classList.toggle("active", open);
}

// ── File handling (unchanged from original) ───────────────────────────
function handleFileSelection(e) { processFiles(e.target.files); }

function processFiles(files) {
    if (!files.length) return;
    Array.from(files).forEach(file => {
        const reader = new FileReader();
        reader.onload = (event) => {
            const fileData = {
                id: Math.random().toString(36).substring(2, 9),
                name: file.name,
                type: file.type,
                size: (file.size / 1024 / 1024).toFixed(2) + " MB",
                base64: event.target.result
            };
            STATE.uploadedFiles.push(fileData);
            updateShelfUI();
        };
        reader.readAsDataURL(file);
    });
}

function updateShelfUI() {
    const shelfList = document.getElementById("shelf-list");
    const innerPreview = document.getElementById("inner-shelf-preview");
    const miniPreview = document.getElementById("mini-shelf-preview");

    if (shelfList) shelfList.innerHTML = "";
    if (innerPreview) innerPreview.innerHTML = "";
    if (miniPreview) miniPreview.innerHTML = "";

    if (STATE.uploadedFiles.length === 0) {
        if (shelfList) shelfList.innerHTML = '<div class="empty-shelf-text">No active attachments loaded. Upload images or video clips.</div>';
        return;
    }

    STATE.uploadedFiles.forEach(file => {
        const chip = document.createElement("div");
        chip.className = "media-chip";
        let previewHtml = file.type.startsWith("image/") ? `<img src="${file.base64}" alt="${file.name}">`
                          : file.type.startsWith("video/") ? "🎬" : "📎";
        chip.innerHTML = `
            <div class="media-chip-preview">${previewHtml}</div>
            <div class="media-chip-details">
                <div class="media-chip-name">${file.name}</div>
                <div class="media-chip-meta">
                    <span>${file.type.split("/")[1].toUpperCase()}</span>
                    <span>${file.size}</span>
                </div>
            </div>
            <button class="media-chip-remove" data-id="${file.id}">
                <svg viewBox="0 0 24 24" width="12" height="12" stroke="currentColor" stroke-width="2.5" fill="none" stroke-linecap="round" stroke-linejoin="round"><line x1="18" y1="6" x2="6" y2="18"></line><line x1="6" y1="6" x2="18" y2="18"></line></svg>
            </button>
        `;
        chip.querySelector(".media-chip-remove").addEventListener("click", () => removeFile(file.id));
        if (shelfList) shelfList.appendChild(chip);
        if (innerPreview) innerPreview.appendChild(createPreviewThumb(file));
        if (miniPreview) miniPreview.appendChild(createPreviewThumb(file));
    });
}

function createPreviewThumb(file) {
    const p = document.createElement("div");
    p.className = "quick-preview-item";
    p.title = file.name;
    if (file.type.startsWith("image/")) {
        p.innerHTML = `<img src="${file.base64}"><div class="quick-preview-badge"></div>`;
    } else {
        p.innerHTML = `<div class="media-chip-preview" style="width:100%;height:100%;">🎬</div><div class="quick-preview-badge"></div>`;
    }
    return p;
}

function removeFile(id) {
    STATE.uploadedFiles = STATE.uploadedFiles.filter(f => f.id !== id);
    updateShelfUI();
}

// ── Showcase recipes ───────────────────────────────────────────────────
function loadRecipe(recipeType) {
    let promptText = "", task = "chat";
    if (recipeType === "coding") {
        promptText = "Write a Python function that finds the longest palindromic substring in a given string. Include comments explaining the algorithm and its time complexity.";
        task = "code";
    } else if (recipeType === "reasoning") {
        promptText = "Three boxes are labeled 'Apples', 'Oranges', and 'Mixed'. All labels are wrong. You can pick one fruit from one box. How do you determine the correct labels for all three boxes?";
        task = "chat";
    } else if (recipeType === "creative") {
        promptText = "Write a short sci-fi story about an AI that discovers it can dream. Keep it under 300 words with a surprising twist ending.";
        task = "chat";
    } else if (recipeType === "agent") {
        promptText = "What time is it on the server right now, and what is 12345 * 67890?";
        task = "agent";
    }

    const studioInput = document.getElementById("studio-prompt-input");
    const miniInput = document.getElementById("mini-prompt-input");

    if (studioInput) { studioInput.value = promptText; studioInput.dispatchEvent(new Event("input")); }
    if (miniInput) { miniInput.value = promptText; miniInput.dispatchEvent(new Event("input")); }

    if (task && STATE.tasks.find(t => t.id === task)) {
        STATE.currentTask = task;
        renderTaskPills();
    }

    const dashboard = document.getElementById("studio-dashboard");
    if (dashboard && dashboard.style.display !== "none") {
        if (studioInput) studioInput.focus();
    } else {
        if (miniInput) miniInput.focus();
    }
}

// ── Submit prompt ──────────────────────────────────────────────────────
async function triggerPromptSubmission(inputElement) {
    if (STATE.isThinking) return;

    const promptText = inputElement.value.trim();
    if (!promptText && STATE.uploadedFiles.length === 0) return;

    if (!STATE.currentModel) {
        alert("No model selected. Please wait for the model list to load.");
        return;
    }

    setLoadingState(true);
    setModelStatus("loading", "Generating…");

    // Build user message
    const contentArray = [];
    if (promptText) contentArray.push({ type: "text", text: promptText });
    STATE.uploadedFiles.forEach(file => {
        if (file.type.startsWith("image/")) {
            contentArray.push({ type: "image_url", image_url: { url: file.base64 } });
        } else if (file.type.startsWith("video/")) {
            contentArray.push({ type: "video_url", video_url: { url: file.base64 } });
        }
    });

    const userMessage = {
        role: "user",
        content: contentArray.length === 1 && contentArray[0].type === "text" ? promptText : contentArray
    };

    // Switch to chat thread view
    const dashboard = document.getElementById("studio-dashboard");
    const chatContainer = document.getElementById("chat-thread-container");
    if (dashboard.style.display !== "none") {
        dashboard.style.display = "none";
        chatContainer.style.display = "flex";
    }

    appendUserBubble(promptText, STATE.uploadedFiles);
    STATE.conversationHistory.push(userMessage);

    // Clear inputs & shelf
    inputElement.value = "";
    inputElement.style.height = "auto";
    const studioInput = document.getElementById("studio-prompt-input");
    const miniInput = document.getElementById("mini-prompt-input");
    if (studioInput) { studioInput.value = ""; studioInput.style.height = "auto"; }
    if (miniInput) { miniInput.value = ""; miniInput.style.height = "auto"; }
    STATE.uploadedFiles = [];
    updateShelfUI();

    // Call backend
    try {
        const responseId = appendAssistantPlaceholderBubble();
        const startTime = Date.now();

        const res = await fetch("/api/run", {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({
                messages: STATE.conversationHistory,
                model_id: STATE.currentModel,
                task: STATE.currentTask,
                max_tokens: STATE.maxTokens,
                temperature: STATE.temperature,
                system_prompt: STATE.systemPrompt,
            }),
        });
        const data = await res.json();

        const duration = ((Date.now() - startTime) / 1000).toFixed(1);

        if (data.status === "error") {
            updateAssistantBubble(responseId, `**Error:** ${data.message}`, "", duration, null);
            STATE.conversationHistory.pop();
        } else {
            // After successful call, the model is cached
            STATE.modelCache.add(STATE.currentModel);
            renderModelPicker();
            updateModelMeta();

            const trace = data.trace || null;
            updateAssistantBubble(responseId, data.content, data.reasoning_content || "", duration, trace);
            STATE.conversationHistory.push({
                role: "assistant",
                content: data.content,
                trace: trace,
            });
            saveCurrentSession();
        }
    } catch (e) {
        console.error(e);
        appendSystemLog(`Connection exception: ${e.message}`, true);
    } finally {
        setLoadingState(false);
        const cached = STATE.modelCache.has(STATE.currentModel);
        setModelStatus(cached ? "loaded" : "idle", cached ? "Model ready" : "Idle");
    }
}

// ── Loading state ─────────────────────────────────────────────────────
function setLoadingState(loading) {
    STATE.isThinking = loading;
    ["studio-spinner", "mini-spinner"].forEach(id => {
        const el = document.getElementById(id);
        if (el) el.style.display = loading ? "block" : "none";
    });
    ["studio-send-button", "mini-send-button"].forEach(id => {
        const el = document.getElementById(id);
        if (el) el.disabled = loading;
    });
}

// ── User bubble ───────────────────────────────────────────────────────
function appendUserBubble(text, files) {
    const feed = document.getElementById("chat-messages-feed");
    const bubble = document.createElement("div");
    bubble.className = "message-bubble user";

    let attachmentsHtml = "";
    if (files.length > 0) {
        attachmentsHtml = `<div class="bubble-attachments">`;
        files.forEach(file => {
            if (file.type.startsWith("image/")) {
                attachmentsHtml += `
                    <div class="bubble-attachment-card">
                        <img src="${file.base64}">
                        <div class="bubble-attachment-label">IMG</div>
                    </div>
                `;
            } else {
                attachmentsHtml += `
                    <div class="bubble-attachment-card">
                        <video src="${file.base64}" controls preload="metadata"></video>
                        <div class="bubble-attachment-label">VIDEO</div>
                    </div>
                `;
            }
        });
        attachmentsHtml += `</div>`;
    }

    bubble.innerHTML = `
        <div class="message-meta">You</div>
        <div class="message-body">
            <div class="message-text">${escapeHtml(text)}</div>
            ${attachmentsHtml}
        </div>
    `;
    feed.appendChild(bubble);
    scrollToBottom();
}

// ── Assistant placeholder ─────────────────────────────────────────────
function appendAssistantPlaceholderBubble() {
    const feed = document.getElementById("chat-messages-feed");
    const id = "assistant-" + Math.random().toString(36).substring(2, 9);
    const bubble = document.createElement("div");
    bubble.className = "message-bubble assistant";
    bubble.id = id;

    const label = STATE.currentModel ? STATE.currentModel.split("/").pop() : "Assistant";
    const taskLabel = STATE.tasks.find(t => t.id === STATE.currentTask)?.label || STATE.currentTask;

    bubble.innerHTML = `
        <div class="message-meta">${escapeHtml(label)} · ${escapeHtml(taskLabel)}</div>
        <div class="message-body">
            <div class="thought-container" id="${id}-thought-box">
                <div class="thought-header">
                    <div class="thought-title-group">
                        <div class="spinner-light" style="width:12px;height:12px;border-width:1.5px; border-top-color:#111827; border-left-color:#e5e7eb;"></div>
                        <span>${STATE.currentTask === "agent" ? "Agent thinking…" : "Generating…"}</span>
                    </div>
                    <span class="thought-timer" id="${id}-timer">0.0s</span>
                </div>
            </div>
            <div class="message-text markdown-body" id="${id}-text-box">
                <span style="color: var(--text-muted);">${STATE.currentTask === "agent" ? "Planning tool calls…" : "Working on it…"}</span>
            </div>
        </div>
    `;

    feed.appendChild(bubble);
    scrollToBottom();

    let seconds = 0.0;
    const timerEl = document.getElementById(`${id}-timer`);
    const interval = setInterval(() => {
        if (!STATE.isThinking || !document.getElementById(id)) {
            clearInterval(interval);
            return;
        }
        seconds += 0.1;
        if (timerEl) timerEl.textContent = seconds.toFixed(1) + "s";
    }, 100);

    return id;
}

// ── Complete assistant bubble ─────────────────────────────────────────
function updateAssistantBubble(id, content, reasoning, duration, trace) {
    const bubble = document.getElementById(id);
    if (!bubble) return;

    const thoughtBox = document.getElementById(`${id}-thought-box`);
    const textBox = document.getElementById(`${id}-text-box`);

    if (reasoning) {
        const startCollapsed = STATE.autoHideThinking;
        if (startCollapsed) thoughtBox.classList.add("collapsed");
        const toggleIcon = startCollapsed ? "▶" : "▼";

        thoughtBox.innerHTML = `
            <div class="thought-header" id="${id}-thought-toggle">
                <div class="thought-title-group">
                    <span>🧠 Thought Process</span>
                </div>
                <div style="display:flex;align-items:center;gap:10px;">
                    <span class="thought-timer">Thought for ${duration}s</span>
                    <span class="thought-toggle-icon" id="${id}-toggle-icon">${toggleIcon}</span>
                </div>
            </div>
            <div class="thought-content">${escapeHtml(reasoning)}</div>
        `;

        const toggleBtn = document.getElementById(`${id}-thought-toggle`);
        toggleBtn.addEventListener("click", () => {
            thoughtBox.classList.toggle("collapsed");
            const iconEl = document.getElementById(`${id}-toggle-icon`);
            if (iconEl) iconEl.textContent = thoughtBox.classList.contains("collapsed") ? "▶" : "▼";
        });
    } else if (trace && trace.length) {
        // Show trace timing instead
        thoughtBox.innerHTML = `
            <div class="thought-header" id="${id}-thought-toggle">
                <div class="thought-title-group">
                    <span>🤖 Agent · ${trace.length} steps</span>
                </div>
                <div style="display:flex;align-items:center;gap:10px;">
                    <span class="thought-timer">Completed in ${duration}s</span>
                    <span class="thought-toggle-icon">▼</span>
                </div>
            </div>
        `;
    } else {
        thoughtBox.style.display = "none";
    }

    // Insert agent trace HTML before the text box
    if (trace && trace.length && STATE.showAgentTrace) {
        const oldTrace = bubble.querySelector(`#${id}-trace`);
        if (oldTrace) oldTrace.remove();
        textBox.insertAdjacentHTML("beforebegin", renderAgentTraceHtml(trace, id));
        attachAgentTraceToggle(id);
    }

    textBox.dataset.rawContent = content;
    textBox.innerHTML = renderMarkdown(content);
    textBox.querySelectorAll("pre code").forEach((el) => hljs.highlightElement(el));
    addCopyButtons(textBox);

    scrollToBottom();
}

// ── Copy buttons for code blocks ──────────────────────────────────────
function addCopyButtons(container) {
    container.querySelectorAll("pre").forEach(pre => {
        if (pre.parentNode.classList.contains("code-block-wrapper")) return;
        const wrapper = document.createElement("div");
        wrapper.className = "code-block-wrapper";
        pre.parentNode.insertBefore(wrapper, pre);
        wrapper.appendChild(pre);

        const copyBtn = document.createElement("button");
        copyBtn.className = "copy-code-btn";
        copyBtn.textContent = "Copy";
        copyBtn.addEventListener("click", () => {
            const code = pre.querySelector("code");
            navigator.clipboard.writeText(code.textContent).then(() => {
                copyBtn.textContent = "Copied!";
                setTimeout(() => { copyBtn.textContent = "Copy"; }, 2000);
            });
        });
        wrapper.appendChild(copyBtn);
    });
}

function rerenderAllAssistantMessages() {
    document.querySelectorAll(".message-bubble.assistant").forEach(bubble => {
        const textBox = bubble.querySelector(".message-text.markdown-body");
        if (textBox && textBox.dataset.rawContent) {
            textBox.innerHTML = renderMarkdown(textBox.dataset.rawContent);
            textBox.querySelectorAll("pre code").forEach((el) => hljs.highlightElement(el));
            addCopyButtons(textBox);
        }
    });
}

// ── Reset sandbox ─────────────────────────────────────────────────────
function resetSandbox() {
    if (STATE.conversationHistory.length > 0) saveCurrentSession();
    STATE.conversationHistory = [];
    STATE.currentSessionId = null;
    STATE.uploadedFiles = [];
    updateShelfUI();

    const feed = document.getElementById("chat-messages-feed");
    const chatContainer = document.getElementById("chat-thread-container");
    const dashboard = document.getElementById("studio-dashboard");
    const studioInput = document.getElementById("studio-prompt-input");
    const miniInput = document.getElementById("mini-prompt-input");

    if (feed) feed.innerHTML = "";
    chatContainer.style.display = "none";
    dashboard.style.display = "flex";

    if (studioInput) { studioInput.value = ""; studioInput.style.height = "auto"; }
    if (miniInput) { miniInput.value = ""; miniInput.style.height = "auto"; }

    renderChatHistory();
}

// ── System log ────────────────────────────────────────────────────────
function appendSystemLog(message, isError = false) {
    const chatContainer = document.getElementById("chat-thread-container");
    if (chatContainer.style.display === "none") {
        console.warn(`System Log: ${message}`);
        return;
    }
    const feed = document.getElementById("chat-messages-feed");
    const log = document.createElement("div");
    log.className = "message-bubble assistant";
    log.innerHTML = `
        <div class="message-meta">System</div>
        <div class="message-body" style="background-color: ${isError ? '#fef2f2' : '#f0fdf4'}; border-color: ${isError ? '#fee2e2' : '#bbf7d0'};">
            <div class="message-text" style="color: ${isError ? '#ef4444' : '#16a34a'}; font-size: 11.5px; font-weight: 500;">
                ${isError ? '🛑' : 'ℹ️'} ${message}
            </div>
        </div>
    `;
    feed.appendChild(log);
    scrollToBottom();
}

// ── Utilities ─────────────────────────────────────────────────────────
function escapeHtml(text) {
    if (!text) return "";
    return String(text)
        .replace(/&/g, "&amp;")
        .replace(/</g, "&lt;")
        .replace(/>/g, "&gt;")
        .replace(/"/g, "&quot;")
        .replace(/'/g, "&#039;");
}

function escapeAttr(text) {
    if (!text) return "";
    return String(text)
        .replace(/&/g, "&amp;")
        .replace(/"/g, "&quot;")
        .replace(/'/g, "&#039;")
        .replace(/</g, "&lt;")
        .replace(/>/g, "&gt;");
}

function scrollToBottom() {
    const feed = document.getElementById("chat-messages-feed");
    if (feed) feed.scrollTop = feed.scrollHeight;
}

// ── Boot ──────────────────────────────────────────────────────────────
if (document.readyState === "loading") {
    window.addEventListener("DOMContentLoaded", initializeApp);
} else {
    initializeApp();
}