frontend/workspace.js

// ============================================
// Workspace State Persistence
// ============================================

async function loadWorkspace() {
    try {
        const response = await apiFetch('/api/workspace');
        if (response.ok) {
            const workspace = await response.json();
            console.log('Workspace loaded:', workspace);
            restoreWorkspace(workspace);
        }
    } catch (e) {
        console.log('Could not load workspace from backend:', e);
    }
}

function restoreWorkspace(workspace) {
    // Restore counters
    tabCounter = workspace.tabCounter || 1;
    agentCounters = workspace.agentCounters || workspace.notebookCounters || getDefaultCounters();

    // Restore timeline data before tabs so renderTimeline works
    if (workspace.timelineData) {
        Object.keys(timelineData).forEach(k => delete timelineData[k]);
        for (const [tabId, data] of Object.entries(workspace.timelineData)) {
            timelineData[tabId] = data;
        }
    }

    // Restore debug history
    if (workspace.debugHistory) {
        Object.keys(debugHistory).forEach(k => delete debugHistory[k]);
        for (const [tabId, calls] of Object.entries(workspace.debugHistory)) {
            debugHistory[tabId] = calls;
        }
    }

    // Restore tabs (skip command center as it already exists in HTML)
    const tabs = workspace.tabs || [];
    for (const tabData of tabs) {
        if (tabData.id === 0) {
            // Restore command center messages
            restoreTabMessages(tabData);
        } else {
            // Create and restore other tabs
            restoreTab(tabData);
        }
    }

    // Restore globalFigureRegistry from saved code-cell images
    // so that <figure_T1_1> tags in results resolve after reload
    for (const tabData of tabs) {
        for (const msg of (tabData.messages || [])) {
            if (msg.type === 'code-cell' && msg.images) {
                for (const img of msg.images) {
                    if (img.name && img.src) {
                        // Parse data URL: "data:image/png;base64,..." -> {type: "png", data: "..."}
                        const m = img.src.match(/^data:image\/(\w+);base64,(.+)$/);
                        if (m) {
                            globalFigureRegistry[img.name] = { type: m[1], data: m[2] };
                        }
                    }
                }
            }
        }
    }

    // Re-resolve any figure refs in rendered HTML now that the registry is populated
    if (Object.keys(globalFigureRegistry).length > 0) {
        document.querySelectorAll('.action-widget-result-section .section-content, .result-preview-content, .result-content').forEach(el => {
            const resolved = resolveGlobalFigureRefs(el.innerHTML);
            if (resolved !== el.innerHTML) {
                el.innerHTML = resolved;
            }
        });
    }

    // Switch to the active tab
    if (workspace.activeTabId !== undefined) {
        switchToTab(workspace.activeTabId);
    }

    // Render timeline after everything is restored
    renderTimeline();
}

function restoreTab(tabData) {
    // Create tab element
    const tab = document.createElement('div');
    tab.className = 'tab';
    tab.dataset.tabId = tabData.id;
    tab.innerHTML = `
        <span class="tab-title">${tabData.title || getTypeLabel(tabData.type) || 'TAB'}</span>
        <span class="tab-status" style="display: none;"><span></span><span></span><span></span></span>
        <span class="tab-close">×</span>
    `;

    // Insert into dynamic tabs container
    const dynamicTabs = document.getElementById('dynamicTabs');
    dynamicTabs.appendChild(tab);

    // Create content element
    const content = document.createElement('div');
    content.className = 'tab-content';
    content.dataset.contentId = tabData.id;
    content.innerHTML = createAgentContent(tabData.type, tabData.id);
    document.querySelector('.main-content').appendChild(content);

    // Add event listeners for the new content
    setupInputListeners(content, tabData.id);

    // Restore messages
    restoreTabMessages(tabData);

    // If this is a code agent, start the sandbox proactively
    if (tabData.type === 'code') {
        startSandbox(tabData.id);
    }
}

function restoreTabMessages(tabData) {
    const content = document.querySelector(`[data-content-id="${tabData.id}"]`);
    if (!content) return;

    const chatContainer = content.querySelector('.chat-container');
    if (!chatContainer) return;

    // Remove welcome message if restoring messages
    if (tabData.messages && tabData.messages.length > 0) {
        const welcomeMsg = chatContainer.querySelector('.welcome-message');
        if (welcomeMsg) {
            welcomeMsg.remove();
        }
    }

    // Restore each message - messages are now flat, each with its own type
    for (const msg of (tabData.messages || [])) {
        if (msg.role === 'user') {
            const userMsg = document.createElement('div');
            userMsg.className = 'message user';
            userMsg.innerHTML = `<div class="message-content">${parseMarkdown(msg.content || '')}</div>`;
            linkifyFilePaths(userMsg);
            chatContainer.appendChild(userMsg);
        } else if (msg.role === 'assistant') {
            // New flat structure: each message has a type field
            switch (msg.type) {
                case 'text':
                    const assistantMsg = document.createElement('div');
                    assistantMsg.className = 'message assistant';
                    const textDiv = document.createElement('div');
                    textDiv.className = 'message-content';
                    textDiv.innerHTML = msg.html || parseMarkdown(msg.content || '');
                    assistantMsg.appendChild(textDiv);
                    chatContainer.appendChild(assistantMsg);
                    break;

                case 'action-widget':
                    const widget = renderActionWidget(msg);
                    chatContainer.appendChild(widget);
                    break;

                case 'code-cell':
                    const codeCell = renderCodeCell(msg);
                    chatContainer.appendChild(codeCell);
                    break;

                case 'tool-call':
                    // Restore hidden tool-call element for getConversationHistory()
                    const tcMsg = document.createElement('div');
                    tcMsg.className = 'message assistant';
                    tcMsg.style.display = 'none';
                    tcMsg.setAttribute('data-tool-call', JSON.stringify(msg.data));
                    chatContainer.appendChild(tcMsg);
                    break;

                case 'tool-cell':
                    const toolCell = document.createElement('div');
                    toolCell.className = 'tool-cell';
                    toolCell.setAttribute('data-tool-name', msg.toolName || '');
                    const toolDesc = msg.input ? `<span class="tool-cell-desc">${escapeHtml(msg.input)}</span>` : '';
                    let toolCellHtml = `<div class="tool-cell-label"><div class="widget-collapse-toggle"></div><span>${escapeHtml(msg.label || '')}</span>${toolDesc}</div>`;
                    toolCellHtml += `<div class="tool-cell-input">${escapeHtml(msg.input || '')}</div>`;
                    if (msg.outputHtml) {
                        toolCellHtml += `<div class="tool-cell-output">${msg.outputHtml}</div>`;
                    }
                    toolCell.innerHTML = toolCellHtml;
                    setupCollapseToggle(toolCell, '.tool-cell-label');
                    chatContainer.appendChild(toolCell);
                    break;

                case 'result-preview':
                    const preview = renderResultPreview(msg);
                    chatContainer.appendChild(preview);
                    break;

                case 'agent-result':
                    const report = renderResearchReport(msg);
                    chatContainer.appendChild(report);
                    break;

                case 'research-container':
                    const researchWidget = renderResearchContainer(msg);
                    chatContainer.appendChild(researchWidget);
                    break;

                default:
                    // Legacy fallback: messages without type or with parts array
                    if (msg.parts && msg.parts.length > 0) {
                        // Old format with parts array
                        for (const part of msg.parts) {
                            switch (part.type) {
                                case 'text':
                                    const oldTextDiv = document.createElement('div');
                                    oldTextDiv.className = 'message assistant';
                                    oldTextDiv.innerHTML = `<div class="message-content">${part.html || parseMarkdown(part.content || '')}</div>`;
                                    chatContainer.appendChild(oldTextDiv);
                                    break;
                                case 'action-widget':
                                    chatContainer.appendChild(renderActionWidget(part));
                                    break;
                                case 'code-cell':
                                    chatContainer.appendChild(renderCodeCell(part));
                                    break;
                                case 'result-preview':
                                    chatContainer.appendChild(renderResultPreview(part));
                                    break;
                                case 'agent-result':
                                    chatContainer.appendChild(renderResearchReport(part));
                                    break;
                            }
                        }
                    } else if (msg.contentHtml) {
                        const legacyMsg = document.createElement('div');
                        legacyMsg.className = 'message assistant';
                        legacyMsg.innerHTML = msg.contentHtml;
                        chatContainer.appendChild(legacyMsg);
                    } else if (msg.content) {
                        const contentMsg = document.createElement('div');
                        contentMsg.className = 'message assistant';
                        contentMsg.innerHTML = `<div class="message-content">${parseMarkdown(msg.content)}</div>`;
                        chatContainer.appendChild(contentMsg);
                    }
                    break;
            }
        } else if (msg.role === 'tool' && msg.type === 'tool-response') {
            // Restore hidden tool-response element for getConversationHistory()
            const trMsg = document.createElement('div');
            trMsg.className = 'message tool';
            trMsg.style.display = 'none';
            trMsg.setAttribute('data-tool-response', JSON.stringify(msg.data));
            chatContainer.appendChild(trMsg);
        }
    }
}

// Helper functions to render saved message parts

function renderActionWidget(data) {
    const widget = document.createElement('div');
    widget.className = 'action-widget';
    if (data.targetTabId) {
        widget.dataset.targetTabId = data.targetTabId;
    }

    const indicatorHtml = data.isDone
        ? '<div class="done-indicator"></div>'
        : '<div class="orbit-indicator"><span></span><span></span><span></span></div>';

    let bodyHtml = `
        <div class="section-label">QUERY</div>
        <div class="section-content">${parseMarkdown(data.query || '')}</div>
    `;

    // Add result section if present
    if (data.resultHtml || data.result) {
        bodyHtml += `
            <div class="action-widget-result-section">
                <div class="section-label" style="margin-top: 10px; padding-top: 10px; border-top: 1px solid var(--border-primary);">RESULT</div>
                <div class="section-content">${data.resultHtml || escapeHtml(data.result || '')}</div>
            </div>
        `;
    }

    widget.innerHTML = `
        <div class="action-widget-clickable">
            <div class="action-widget-header">
                <div class="widget-collapse-toggle"></div>
                <span class="action-widget-type">TASK: ${(data.actionType || '').toUpperCase()}</span>
                <div class="action-widget-bar-right">
                    ${indicatorHtml}
                </div>
            </div>
        </div>
        <div class="action-widget-body">
            ${bodyHtml}
        </div>
    `;

    // Collapse toggle
    const collapseToggle = widget.querySelector('.widget-collapse-toggle');
    collapseToggle.addEventListener('click', (e) => {
        e.stopPropagation();
        widget.classList.toggle('collapsed');
        collapseToggle.classList.toggle('collapsed');
    });

    // Make clickable if we have a target tab
    if (data.targetTabId) {
        const clickableArea = widget.querySelector('.action-widget-clickable');
        clickableArea.addEventListener('click', () => {
            const tabId = parseInt(data.targetTabId);
            if (!isNaN(tabId)) {
                // Check if the tab still exists (use .tab to avoid matching timeline elements)
                const tab = document.querySelector(`.tab[data-tab-id="${tabId}"]`);
                if (tab) {
                    // Tab exists, just switch to it
                    switchToTab(tabId);
                } else {
                    // Tab was closed - restore from timeline data
                    const notebook = timelineData[tabId];
                    if (notebook) {
                        reopenClosedTab(tabId, notebook);
                    }
                }
            }
        });
    }

    return widget;
}

function renderCodeCell(data) {
    const cell = document.createElement('div');
    cell.className = 'code-cell';

    let outputHtml = '';
    const cleanedData = cleanCodeOutput(data.output);
    if (data.outputHtml || cleanedData) {
        const errorClass = data.isError ? ' error' : '';
        outputHtml = `<div class="code-cell-output${errorClass}">${data.outputHtml || escapeHtml(cleanedData)}</div>`;
    }

    // Build images HTML
    let imagesHtml = '';
    if (data.images && data.images.length > 0) {
        for (const img of data.images) {
            let labelHtml = '';
            if (img.name) {
                labelHtml = `<div class="figure-label">${escapeHtml(img.name)}</div>`;
            }
            imagesHtml += `<div class="code-cell-image">${labelHtml}<img src="${img.src}" onclick="openImageModal(this.src)"></div>`;
        }
    }

    cell.innerHTML = `
        <div class="code-cell-label"><div class="widget-collapse-toggle"></div><span>Code</span></div>
        <div class="code-cell-code">
            <pre><code class="language-python">${escapeHtml(data.code || '')}</code></pre>
        </div>
        ${outputHtml}
        ${imagesHtml}
    `;
    setupCollapseToggle(cell, '.code-cell-label');

    return cell;
}

function renderResultPreview(data) {
    const preview = document.createElement('div');
    preview.className = 'result-preview';

    preview.innerHTML = `
        <div class="result-preview-label">Result</div>
        <div class="result-preview-content">${data.html || escapeHtml(data.content || '')}</div>
    `;
    linkifyFilePaths(preview);

    return preview;
}

function renderResearchReport(data) {
    const report = document.createElement('div');
    report.className = 'agent-result';

    report.innerHTML = `
        <div class="result-header">Result</div>
        <div class="result-content">${data.html || parseMarkdown(data.content || '')}</div>
    `;
    linkifyFilePaths(report);

    return report;
}

function renderResearchContainer(data) {
    const container = document.createElement('div');
    container.className = 'research-container message assistant';

    container.innerHTML = `
        <div class="research-header">
            <span class="research-header-title">RESEARCH</span>
        </div>
        <div class="research-body">${data.html || ''}</div>
    `;

    return container;
}

function serializeWorkspace() {
    const workspace = {
        version: 1,
        tabCounter: tabCounter,
        activeTabId: activeTabId,
        agentCounters: agentCounters,
        tabs: [],
        timelineData: serializeTimelineData(),
        debugHistory: debugHistory
    };

    // Serialize command center (tab 0)
    workspace.tabs.push(serializeTab(0, 'command-center'));

    // Serialize all other tabs
    const tabElements = document.querySelectorAll('#dynamicTabs .tab');
    for (const tabEl of tabElements) {
        const tabId = parseInt(tabEl.dataset.tabId);
        const content = document.querySelector(`[data-content-id="${tabId}"]`);
        if (content) {
            const chatContainer = content.querySelector('.chat-container');
            const agentType = chatContainer?.dataset.agentType || 'chat';
            workspace.tabs.push(serializeTab(tabId, agentType));
        }
    }

    return workspace;
}

function serializeTimelineData() {
    const serialized = {};
    for (const [tabId, data] of Object.entries(timelineData)) {
        // Clone but strip savedContent (large HTML, not needed for timeline)
        serialized[tabId] = {
            type: data.type,
            title: data.title,
            events: data.events,
            parentTabId: data.parentTabId,
            isGenerating: false,
            isClosed: data.isClosed || false
        };
    }
    return serialized;
}

function serializeTab(tabId, type) {
    const tabEl = document.querySelector(`[data-tab-id="${tabId}"]`);
    const content = document.querySelector(`[data-content-id="${tabId}"]`);

    const tabData = {
        id: tabId,
        type: type,
        title: tabEl?.querySelector('.tab-title')?.textContent || getTypeLabel(type) || 'TAB',
        messages: []
    };

    if (!content) return tabData;

    const chatContainer = content.querySelector('.chat-container');
    if (!chatContainer) return tabData;

    // Iterate over ALL direct children of chatContainer in order
    // Elements can be: .message.user, .message.assistant, .action-widget, .code-cell, .agent-result, .system-message, .welcome-message
    for (const child of chatContainer.children) {
        // Skip welcome message and system messages
        if (child.classList.contains('welcome-message') || child.classList.contains('system-message')) {
            continue;
        }

        // User message
        if (child.classList.contains('message') && child.classList.contains('user')) {
            tabData.messages.push({
                role: 'user',
                content: child.textContent.trim()
            });
            continue;
        }

        // Research container (the search tree widget) - check BEFORE generic assistant message
        // because research-container also has .message.assistant classes
        if (child.classList.contains('research-container')) {
            const bodyEl = child.querySelector('.research-body');
            tabData.messages.push({
                role: 'assistant',
                type: 'research-container',
                html: bodyEl?.innerHTML || ''
            });
            continue;
        }

        // Tool-call message (hidden, used for LLM conversation history)
        if (child.classList.contains('message') && child.hasAttribute('data-tool-call')) {
            try {
                tabData.messages.push({
                    role: 'assistant',
                    type: 'tool-call',
                    data: JSON.parse(child.getAttribute('data-tool-call'))
                });
            } catch (e) { /* skip unparseable */ }
            continue;
        }

        // Tool-response message (hidden, used for LLM conversation history)
        if (child.classList.contains('message') && (child.hasAttribute('data-tool-response') || child.classList.contains('tool'))) {
            try {
                tabData.messages.push({
                    role: 'tool',
                    type: 'tool-response',
                    data: JSON.parse(child.getAttribute('data-tool-response'))
                });
            } catch (e) { /* skip unparseable */ }
            continue;
        }

        // Assistant message (with .message-content inside)
        if (child.classList.contains('message') && child.classList.contains('assistant')) {
            const messageContent = child.querySelector('.message-content');
            if (messageContent) {
                tabData.messages.push({
                    role: 'assistant',
                    type: 'text',
                    content: messageContent.textContent.trim(),
                    html: messageContent.innerHTML
                });
            } else {
                // Fallback to raw HTML
                tabData.messages.push({
                    role: 'assistant',
                    type: 'text',
                    html: child.innerHTML
                });
            }
            continue;
        }

        // Action widget (appended directly to chatContainer)
        if (child.classList.contains('action-widget')) {
            const widgetData = {
                role: 'assistant',
                type: 'action-widget',
                targetTabId: child.dataset.targetTabId,
                actionType: child.querySelector('.action-widget-type')?.textContent?.replace('TASK: ', '') || '',
                query: child.querySelector('.action-widget-body .section-content')?.textContent?.trim() || '',
                isDone: !!child.querySelector('.done-indicator'),
                result: null,
                resultHtml: null
            };

            // Extract result if present
            const resultSection = child.querySelector('.action-widget-result-section');
            if (resultSection) {
                const resultContent = resultSection.querySelector('.section-content');
                if (resultContent) {
                    widgetData.result = resultContent.textContent.trim();
                    widgetData.resultHtml = resultContent.innerHTML;
                }
            }

            tabData.messages.push(widgetData);
            continue;
        }

        // Code cell (appended directly to chatContainer)
        if (child.classList.contains('code-cell')) {
            const codeEl = child.querySelector('pre code');
            const outputEl = child.querySelector('.code-cell-output');

            // Collect images from the code cell
            const images = [];
            const imageEls = child.querySelectorAll('.code-cell-image');
            for (const imgDiv of imageEls) {
                const imgEl = imgDiv.querySelector('img');
                const labelEl = imgDiv.querySelector('.figure-label');
                if (imgEl) {
                    images.push({
                        src: imgEl.src,  // data URL with base64
                        name: labelEl?.textContent || ''
                    });
                }
            }

            tabData.messages.push({
                role: 'assistant',
                type: 'code-cell',
                code: codeEl?.textContent || '',
                output: outputEl?.textContent || '',
                outputHtml: outputEl?.innerHTML || '',
                isError: outputEl?.classList.contains('error') || false,
                images: images
            });
            continue;
        }

        // Tool cell (web agent / image agent tool calls)
        if (child.classList.contains('tool-cell')) {
            const labelEl = child.querySelector('.tool-cell-label span');
            const inputEl = child.querySelector('.tool-cell-input');
            const outputEl = child.querySelector('.tool-cell-output');

            tabData.messages.push({
                role: 'assistant',
                type: 'tool-cell',
                toolName: child.getAttribute('data-tool-name') || '',
                label: labelEl?.textContent || '',
                input: inputEl?.textContent || '',
                outputHtml: outputEl?.innerHTML || '',
            });
            continue;
        }

        // Research report (appended directly to chatContainer)
        if (child.classList.contains('agent-result')) {
            const contentEl = child.querySelector('.result-content');
            tabData.messages.push({
                role: 'assistant',
                type: 'agent-result',
                content: contentEl?.textContent || '',
                html: contentEl?.innerHTML || ''
            });
            continue;
        }

        // Result preview
        if (child.classList.contains('result-preview')) {
            const contentEl = child.querySelector('.result-preview-content');
            tabData.messages.push({
                role: 'assistant',
                type: 'result-preview',
                content: contentEl?.textContent || '',
                html: contentEl?.innerHTML || ''
            });
            continue;
        }
    }

    return tabData;
}

function saveWorkspaceDebounced() {
    // Clear any pending save
    if (saveWorkspaceTimer) {
        clearTimeout(saveWorkspaceTimer);
    }

    // Schedule save in 500ms
    saveWorkspaceTimer = setTimeout(async () => {
        try {
            const workspace = serializeWorkspace();
            const response = await apiFetch('/api/workspace', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(workspace)
            });
            if (response.ok) {
                console.log('Workspace saved');
            }
        } catch (e) {
            console.error('Failed to save workspace:', e);
        }
    }, 500);
}