// ==UserScript==
// @name         酒馆ComfyUI插图脚本
// @namespace    http://tampermonkey.net/
// @version      7
// @license      GPL
// @description  用于酒馆SillyTavern的ai插图脚本，替换特定字符为图片，并使用特定字符内生成的prompt通过ComfyUI API生图。图片数据直接缓存，离线可用。
// @author       soulostar
// @match        *://*/*
// @grant        GM_addStyle
// @grant        GM_setValue
// @grant        GM_getValue
// @grant        GM_xmlhttpRequest
// @require      https://code.jquery.com/ui/1.13.2/jquery-ui.min.js
// @downloadURL none
// ==/UserScript==

(function() {
    'use strict';

    // --- 配置常量 ---
    const BUTTON_ID = 'comfyui-launcher-button';
    const PANEL_ID = 'comfyui-panel';
    const POLLING_TIMEOUT_MS = 60000;
    const POLLING_INTERVAL_MS = 2000;
    const STORAGE_KEY_IMAGES = 'comfyui_generated_images';

    // --- 新增：用于防止自动生图重复触发的状态锁 ---
    const autoGenerationTriggeredIds = new Set();


    // --- 注入自定义CSS样式 ---
    GM_addStyle(`
        /* ... CSS样式部分没有变化，为节省空间已折叠 ... */
        /* 控制面板主容器样式 */
        #${PANEL_ID} { display: none; position: fixed; top: 50%; left: 50%; transform: translate(-50%, -50%); width: 90vw; max-width: 500px; z-index: 9999; color: var(--SmartThemeBodyColor, #dcdcd2); background-color: var(--SmartThemeBlurTintColor, rgba(23, 23, 23, 0.9)); border: 1px solid var(--SmartThemeBorderColor, rgba(0, 0, 0, 0.5)); border-radius: 8px; box-shadow: 0 4px 15px var(--SmartThemeShadowColor, rgba(0, 0, 0, 0.5)); padding: 15px; box-sizing: border-box; backdrop-filter: blur(var(--blurStrength, 10px)); flex-direction: column; }
        #${PANEL_ID} .panel-control-bar { cursor: move; padding-bottom: 10px; margin-bottom: 15px; border-bottom: 1px solid var(--SmartThemeBorderColor, rgba(0, 0, 0, 0.5)); display: flex; align-items: center; justify-content: space-between; flex-shrink: 0; }
        #${PANEL_ID} .panel-control-bar b { font-size: 1.2em; margin-left: 10px; }
        #${PANEL_ID} .floating_panel_close { cursor: pointer; font-size: 1.5em; }
        #${PANEL_ID} .floating_panel_close:hover { opacity: 0.7; }
        #${PANEL_ID} .comfyui-panel-content { overflow-y: auto; flex-grow: 1; padding-right: 5px; }
        #${PANEL_ID} input[type="text"], #${PANEL_ID} input[type="number"], #${PANEL_ID} textarea { width: 100%; box-sizing: border-box; padding: 8px; border-radius: 4px; border: 1px solid var(--SmartThemeBorderColor, #555); background-color: rgba(0,0,0,0.2); color: var(--SmartThemeBodyColor, #dcdcd2); margin-bottom: 10px; }
        #${PANEL_ID} textarea { min-height: 150px; resize: vertical; }
        #${PANEL_ID} .workflow-info { font-size: 0.9em; color: #aaa; margin-top: -5px; margin-bottom: 10px;}
        .comfy-button { padding: 8px 12px; border: 1px solid black; border-radius: 4px; cursor: pointer; background: linear-gradient(135deg, #171717); color: white; font-weight: 600; transition: opacity 0.3s, background 0.3s; flex-shrink: 0; font-size: 14px; }
        .comfy-button:disabled { opacity: 0.5; cursor: not-allowed; }
        .comfy-button:hover:not(:disabled) { opacity: 0.85; }
        .comfy-button.testing { background: #555; }
        .comfy-button.success { background: linear-gradient(135deg, #28a745 0%, #218838 100%); }
        .comfy-button.error   { background: linear-gradient(135deg, #dc3545 0%, #c82333 100%); }
        #comfyui-test-conn, #comfyui-apply-width { position: relative; top: -5px; }
        .comfy-input-group { display: flex; gap: 10px; align-items: center; margin-bottom: 10px; }
        .comfy-input-group input { flex-grow: 1; margin-bottom: 0; }
        #${PANEL_ID} label { display: block; margin-bottom: 5px; font-weight: bold; }
        #options > .options-content > a#${BUTTON_ID} { display: flex; align-items: center; gap: 10px; }
        #${PANEL_ID} .comfy-tags-container { display: flex; gap: 10px; align-items: flex-end; margin-top: 10px; margin-bottom: 10px; }
        #${PANEL_ID} .comfy-tags-container div { flex-grow: 1; }
        #${PANEL_ID} .comfy-auto-generate-container { margin-bottom: 15px; padding-top: 5px; }
        #${PANEL_ID} .comfy-auto-generate-label { display: flex; align-items: center; gap: 8px; cursor: pointer; font-weight: bold; }
        #${PANEL_ID} .comfy-auto-generate-label input[type="checkbox"] { width: auto; margin-bottom: 0; transform: scale(1.2); }
        #${PANEL_ID} .comfy-auto-generate-label span { font-weight: normal; font-size: 0.9em; opacity: 0.9;}
        .comfy-button-group { display: inline-flex; align-items: center; gap: 5px; margin: 5px 4px; }
        .comfy-image-container { margin-top: 10px; max-width: 100%; }
        .comfy-image-container img { max-width: 100%; height: auto; border-radius: 8px; border: 1px solid var(--SmartThemeBorderColor, #555); }
        @media (max-width: 1000px) { #${PANEL_ID} { top: 20px; left: 50%; transform: translateX(-50%); max-height: calc(100vh - 40px); width: 95vw; } }
    `);

    // ... createComfyUIPanel 和 initPanelLogic 等函数没有变化 ...
    async function applyImageWidthToAll() {
        const width = await GM_getValue('comfyui_image_width', 400);
        const allImages = document.querySelectorAll('.comfy-image-container img');
        allImages.forEach(img => {
            img.style.maxWidth = `${width}px`;
        });
        if (typeof toastr !== 'undefined') toastr.success(`图片宽度已应用为 ${width}px`);
    }

    function createComfyUIPanel() {
        if (document.getElementById(PANEL_ID)) return;
        const panelHTML = `
            <div id="${PANEL_ID}">
                <div class="panel-control-bar">
                    <i class="fa-fw fa-solid fa-grip drag-grabber"></i>
                    <b>ComfyUI 生图设置</b>
                    <i class="fa-fw fa-solid fa-circle-xmark floating_panel_close"></i>
                </div>
                <div class="comfyui-panel-content">
                    <label for="comfyui-url">ComfyUI URL</label>
                    <div class="comfy-input-group">
                        <input id="comfyui-url" type="text" placeholder="例如: http://127.0.0.1:8188">
                        <button id="comfyui-test-conn" class="comfy-button">测试连接</button>
                    </div>

                    <label for="comfyui-image-width">图片显示宽度 (px)</label>
                    <div class="comfy-input-group">
                        <input id="comfyui-image-width" type="number" placeholder="例如: 400" min="50">
                        <button id="comfyui-apply-width" class="comfy-button">应用</button>
                    </div>

                    <div class="comfy-auto-generate-container">
                         <label class="comfy-auto-generate-label">
                            <input id="comfyui-auto-generate" type="checkbox">
                            自动生图 <span>(仅对最新消息的“开始生成”有效)</span>
                         </label>
                    </div>


                    <div class="comfy-tags-container">
                        <div>
                            <label for="comfyui-start-tag">开始标记</label>
                            <input id="comfyui-start-tag" type="text">
                        </div>
                        <div>
                            <label for="comfyui-end-tag">结束标记</label>
                            <input id="comfyui-end-tag" type="text">
                        </div>
                    </div>
                    <label for="comfyui-workflow">工作流 (JSON格式)</label>
                    <p class="workflow-info">请在您的工作流JSON中包含 <b>%prompt%</b> (必需) 和 <b>%seed%</b> (可选) 占位符。</p>
                    <textarea id="comfyui-workflow" placeholder="在此处粘贴您的ComfyUI工作流JSON..."></textarea>
                    <button id="comfyui-clear-cache" class="comfy-button error" style="margin-top: 15px; width: 100%;">删除所有图片缓存</button>
                </div>
            </div>
        `;
        document.body.insertAdjacentHTML('beforeend', panelHTML);
        initPanelLogic();
    }

    function initPanelLogic() {
        const panel = document.getElementById(PANEL_ID);
        const closeButton = panel.querySelector('.floating_panel_close');
        const testButton = document.getElementById('comfyui-test-conn');
        const clearCacheButton = document.getElementById('comfyui-clear-cache');
        const urlInput = document.getElementById('comfyui-url');
        const workflowInput = document.getElementById('comfyui-workflow');
        const startTagInput = document.getElementById('comfyui-start-tag');
        const endTagInput = document.getElementById('comfyui-end-tag');
        const widthInput = document.getElementById('comfyui-image-width');
        const applyWidthButton = document.getElementById('comfyui-apply-width');
        const autoGenerateCheckbox = document.getElementById('comfyui-auto-generate');


        closeButton.addEventListener('click', () => { panel.style.display = 'none'; });

        if (typeof $ !== 'undefined' && typeof $.fn.draggable !== 'undefined') {
            $(`#${PANEL_ID}`).draggable({ handle: ".panel-control-bar", containment: "window" });
        }

        testButton.addEventListener('click', () => {
            let url = urlInput.value.trim();
            if (!url) {
                if (typeof toastr !== 'undefined') toastr.warning('请输入ComfyUI的URL。');
                return;
            }

            if (!url.startsWith('http://') && !url.startsWith('https://')) { url = 'http://' + url; }
            if (url.endsWith('/')) { url = url.slice(0, -1); }
            urlInput.value = url;

            const testUrl = url + '/system_stats';
            if (typeof toastr !== 'undefined') toastr.info('正在尝试连接 ComfyUI...');

            testButton.classList.remove('success', 'error');
            testButton.classList.add('testing');
            testButton.disabled = true;

            GM_xmlhttpRequest({
                method: "GET",
                url: testUrl,
                timeout: 5000,
                onload: (res) => {
                    testButton.disabled = false;
                    testButton.classList.remove('testing');
                    if (res.status === 200) {
                        testButton.classList.add('success');
                        if (typeof toastr !== 'undefined') toastr.success('连接成功！ComfyUI服务可用。');
                    } else {
                        testButton.classList.add('error');
                        if (typeof toastr !== 'undefined') toastr.error(`连接失败！服务器响应状态: ${res.status}`);
                    }
                },
                onerror: () => {
                    testButton.disabled = false;
                    testButton.classList.remove('testing');
                    testButton.classList.add('error');
                    if (typeof toastr !== 'undefined') toastr.error('连接错误！请检查URL、网络或CORS设置。');
                },
                ontimeout: () => {
                    testButton.disabled = false;
                    testButton.classList.remove('testing');
                    testButton.classList.add('error');
                    if (typeof toastr !== 'undefined') toastr.error('连接超时！ComfyUI服务可能没有响应。');
                }
            });
        });

        clearCacheButton.addEventListener('click', () => {
            if (confirm('您确定要删除所有已生成的图片缓存吗？\n此操作不可撤销，将立即从界面移除所有已生成的图片。')) {
                GM_setValue(STORAGE_KEY_IMAGES, {});
                document.querySelectorAll('.comfy-button-group').forEach(group => {
                    const imageContainer = group.nextElementSibling;
                    if (imageContainer && imageContainer.classList.contains('comfy-image-container')) {
                        imageContainer.remove();
                    }
                    const deleteButton = group.querySelector('.comfy-delete-button');
                    if (deleteButton) deleteButton.remove();
                    const generateButton = group.querySelector('.comfy-chat-generate-button');
                    if (generateButton) {
                        generateButton.textContent = '开始生成';
                        generateButton.disabled = false;
                        generateButton.classList.remove('testing', 'success', 'error');
                    }
                });
                if (typeof toastr !== 'undefined') toastr.success('所有图片缓存已删除！');
            }
        });

        applyWidthButton.addEventListener('click', applyImageWidthToAll);

        loadSettings(urlInput, workflowInput, startTagInput, endTagInput, widthInput, autoGenerateCheckbox);

        [urlInput, workflowInput, startTagInput, endTagInput, widthInput].forEach(input => {
            input.addEventListener('input', () => {
                if(input === urlInput) testButton.classList.remove('success', 'error', 'testing');
                saveSettings(urlInput, workflowInput, startTagInput, endTagInput, widthInput, autoGenerateCheckbox);
            });
        });
        autoGenerateCheckbox.addEventListener('change', () => {
             saveSettings(urlInput, workflowInput, startTagInput, endTagInput, widthInput, autoGenerateCheckbox);
        });
    }

    async function loadSettings(urlInput, workflowInput, startTagInput, endTagInput, widthInput, autoGenerateCheckbox) {
        urlInput.value = await GM_getValue('comfyui_url', 'http://127.0.0.1:8188');
        workflowInput.value = await GM_getValue('comfyui_workflow', '');
        startTagInput.value = await GM_getValue('comfyui_start_tag', 'image###');
        endTagInput.value = await GM_getValue('comfyui_end_tag', '###');
        widthInput.value = await GM_getValue('comfyui_image_width', 400);
        autoGenerateCheckbox.checked = await GM_getValue('comfyui_auto_generate', false);
    }

    async function saveSettings(urlInput, workflowInput, startTagInput, endTagInput, widthInput, autoGenerateCheckbox) {
        await GM_setValue('comfyui_url', urlInput.value);
        await GM_setValue('comfyui_workflow', workflowInput.value);
        await GM_setValue('comfyui_start_tag', startTagInput.value);
        await GM_setValue('comfyui_end_tag', endTagInput.value);
        await GM_setValue('comfyui_image_width', parseInt(widthInput.value, 10) || 400);
        await GM_setValue('comfyui_auto_generate', autoGenerateCheckbox.checked);
    }

    function addMainButton() {
        if (document.getElementById(BUTTON_ID)) return;
        const optionsMenuContent = document.querySelector('#options .options-content');
        if (optionsMenuContent) {
             const continueButton = optionsMenuContent.querySelector('#option_continue');
             if (continueButton) {
                const comfyButton = document.createElement('a');
                comfyButton.id = BUTTON_ID;
                comfyButton.className = 'interactable';
                comfyButton.innerHTML = `<i class="fa-lg fa-solid fa-image"></i><span>ComfyUI生图</span>`;
                comfyButton.style.cursor = 'pointer';

                comfyButton.addEventListener('click', (event) => {
                    event.preventDefault();
                    const panel = document.getElementById(PANEL_ID);
                    if (panel) { panel.style.display = 'flex'; }
                    document.getElementById('options').style.display = 'none';
                });
                continueButton.parentNode.insertBefore(comfyButton, continueButton.nextSibling);
             }
        }
    }


    // --- 聊天消息处理与图片生成 ---
    // ... 辅助函数 escapeRegex, simpleHash, save/deleteImageRecord 等没有变化 ...
    function escapeRegex(string) { return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); }
    function simpleHash(str) { let hash = 0; for (let i = 0; i < str.length; i++) { const char = str.charCodeAt(i); hash = (hash << 5) - hash + char; hash |= 0; } return 'comfy-id-' + Math.abs(hash).toString(36); }
    async function saveImageRecord(generationId, imageDataUrl) { const records = await GM_getValue(STORAGE_KEY_IMAGES, {}); records[generationId] = imageDataUrl; await GM_setValue(STORAGE_KEY_IMAGES, records); }
    async function deleteImageRecord(generationId) { const records = await GM_getValue(STORAGE_KEY_IMAGES, {}); delete records[generationId]; await GM_setValue(STORAGE_KEY_IMAGES, records); }


    async function processMessageForComfyButton(messageNode) {
        const mesText = messageNode.querySelector('.mes_text');
        if (!mesText) return;

        const startTag = await GM_getValue('comfyui_start_tag', 'image###');
        const endTag = await GM_getValue('comfyui_end_tag', '###');
        if (!startTag || !endTag) return;

        const escapedStartTag = escapeRegex(startTag);
        const escapedEndTag = escapeRegex(endTag);
        const regex = new RegExp(escapedStartTag + '([\\s\\S]*?)' + escapedEndTag, 'g');
        const currentHtml = mesText.innerHTML;

        if (regex.test(currentHtml) && !mesText.querySelector('.comfy-button-group')) {
            mesText.innerHTML = currentHtml.replace(regex, (match, prompt) => {
                const cleanPrompt = prompt.trim();
                const encodedPrompt = cleanPrompt.replace(/"/g, '"');
                const generationId = simpleHash(cleanPrompt);
                return `<span class="comfy-button-group" data-generation-id="${generationId}">
                            <button class="comfy-button comfy-chat-generate-button" data-prompt="${encodedPrompt}">开始生成</button>
                        </span>`;
            });
        }

        const savedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});
        const buttonGroups = mesText.querySelectorAll('.comfy-button-group');

        const autoGenerateEnabled = await GM_getValue('comfyui_auto_generate', false);

        for (const group of buttonGroups) {
            if (group.dataset.listenerAttached) continue;

            const generationId = group.dataset.generationId;
            const generateButton = group.querySelector('.comfy-chat-generate-button');

            if (savedImages[generationId]) {
                await displayImage(group, savedImages[generationId]);
                setupGeneratedState(generateButton, generationId);
            } else {
                 generateButton.addEventListener('click', onGenerateButtonClick);

                // --- 修改：增加重复触发的检查 ---
                if (autoGenerateEnabled &&
                    messageNode.classList.contains('last_mes') &&
                    generateButton.textContent === '开始生成' &&
                    !autoGenerationTriggeredIds.has(generationId) // <-- 关键检查
                   )
                {
                    // 在触发前立即“锁定”这个ID
                    autoGenerationTriggeredIds.add(generationId);
                    console.log(`[ComfyUI] 自动为最新消息触发生成 (ID: ${generationId}): ${generateButton.dataset.prompt.substring(0, 30)}...`);
                    setTimeout(() => generateButton.click(), 100);
                }
            }
            group.dataset.listenerAttached = 'true';
        }
    }


    // ... setupGeneratedState, onGenerateButtonClick, fetchImageAsBase64 等函数没有变化 ...
    function setupGeneratedState(generateButton, generationId) {
        generateButton.textContent = '重新生成';
        generateButton.disabled = false;
        generateButton.classList.remove('testing', 'success', 'error');

        if (!generateButton.dataset.regenerateListener) {
            generateButton.addEventListener('click', onGenerateButtonClick);
            generateButton.dataset.regenerateListener = 'true';
        }

        const group = generateButton.closest('.comfy-button-group');
        let deleteButton = group.querySelector('.comfy-delete-button');

        if (!deleteButton) {
            deleteButton = document.createElement('button');
            deleteButton.textContent = '删除';
            deleteButton.className = 'comfy-button error comfy-delete-button';
            deleteButton.addEventListener('click', async () => {
                await deleteImageRecord(generationId);
                const imageContainer = group.nextElementSibling;
                if (imageContainer && imageContainer.classList.contains('comfy-image-container')) {
                    imageContainer.remove();
                }
                deleteButton.remove();
                generateButton.textContent = '开始生成';
                generateButton.disabled = false;
                generateButton.classList.remove('testing', 'success', 'error');
            });
            generateButton.insertAdjacentElement('afterend', deleteButton);
        }
    }
    function fetchImageAsBase64(url) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'GET',
                url: url,
                responseType: 'blob',
                timeout: 30000,
                onload: (response) => {
                    if (response.status === 200) {
                        const reader = new FileReader();
                        reader.onloadend = () => { resolve(reader.result); };
                        reader.onerror = () => { reject(new Error('无法将Blob转换为Base64。')); };
                        reader.readAsDataURL(response.response);
                    } else {
                        reject(new Error(`下载图片失败，状态码: ${response.status}`));
                    }
                },
                onerror: () => reject(new Error('下载图片时发生网络错误。')),
                ontimeout: () => reject(new Error('下载图片超时。')),
            });
        });
    }
    async function onGenerateButtonClick(event) {
        const button = event.target.closest('.comfy-chat-generate-button');
        const group = button.closest('.comfy-button-group');
        const prompt = button.dataset.prompt;
        const generationId = group.dataset.generationId;

        button.textContent = '生成中...';
        button.disabled = true;
        button.classList.remove('success', 'error');
        button.classList.add('testing');

        const deleteButton = group.querySelector('.comfy-delete-button');
        if (deleteButton) deleteButton.style.display = 'none';
        const oldImageContainer = group.nextElementSibling;
        if (oldImageContainer && oldImageContainer.classList.contains('comfy-image-container')) {
            oldImageContainer.remove();
        }

        try {
            const url = (await GM_getValue('comfyui_url', '')).trim();
            let workflowString = await GM_getValue('comfyui_workflow', '');

            if (!url || !workflowString) throw new Error('ComfyUI URL 或工作流未配置。');
            if (!workflowString.includes('%prompt%')) throw new Error('工作流中未找到必需的 %prompt% 占位符。');

            const seed = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER);
            workflowString = workflowString.replace(/%prompt%/g, JSON.stringify(prompt).slice(1, -1));
            workflowString = workflowString.replace(/%seed%/g, seed);

            const workflow = JSON.parse(workflowString);

            const promptResponse = await sendPromptRequest(url, workflow);
            const promptId = promptResponse.prompt_id;
            if (!promptId) throw new Error('ComfyUI 未返回有效的 Prompt ID。');

            const finalHistory = await pollForResult(url, promptId);
            const imageUrl = findImageUrlInHistory(finalHistory, promptId, url);
            if (!imageUrl) throw new Error('在ComfyUI返回结果中未找到图片。');

            await displayImage(group, imageUrl);

            try {
                const base64Data = await fetchImageAsBase64(imageUrl);
                await saveImageRecord(generationId, base64Data);
                console.log(`[ComfyUI] 图片 ${generationId} 已成功缓存为Base64数据。`);
            } catch (cacheError) {
                console.error('ComfyUI图片缓存失败:', cacheError);
                if (typeof toastr !== 'undefined') toastr.warning('图片已生成，但本地缓存失败，刷新后将无法显示。');
            }

            button.textContent = '生成成功';
            button.classList.remove('testing');
            button.classList.add('success');

            setTimeout(() => {
                setupGeneratedState(button, generationId);
                if (deleteButton) deleteButton.style.display = 'inline-flex';
            }, 2000);

        } catch (e) {
            if (typeof toastr !== 'undefined') toastr.error(e.message);
            console.error('ComfyUI生图脚本错误:', e);
            button.textContent = '生成失败';
            button.classList.remove('testing');
            button.classList.add('error');

            setTimeout(() => {
                const wasRegenerating = !!group.querySelector('.comfy-delete-button');
                if (wasRegenerating) {
                    setupGeneratedState(button, generationId);
                    if (deleteButton) deleteButton.style.display = 'inline-flex';
                } else {
                     button.textContent = '开始生成';
                     button.disabled = false;
                     button.classList.remove('error');
                }
            }, 3000);
        }
    }
    function sendPromptRequest(url, workflow) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'POST',
                url: `${url}/prompt`,
                headers: { 'Content-Type': 'application/json' },
                data: JSON.stringify({ prompt: workflow }),
                timeout: 10000,
                onload: (res) => {
                    if (res.status === 200) {
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至ComfyUI，排队中...');
                        resolve(JSON.parse(res.responseText));
                    } else {
                        reject(new Error(`ComfyUI API 错误 (Prompt): ${res.statusText || res.status}`));
                    }
                },
                onerror: () => reject(new Error('无法连接到 ComfyUI API。')),
                ontimeout: () => reject(new Error('连接 ComfyUI API 超时。')),
            });
        });
    }
    function pollForResult(url, promptId) {
        return new Promise((resolve, reject) => {
            const startTime = Date.now();
            const poller = setInterval(() => {
                if (Date.now() - startTime > POLLING_TIMEOUT_MS) {
                    clearInterval(poller);
                    reject(new Error('轮询ComfyUI结果超时。'));
                    return;
                }
                GM_xmlhttpRequest({
                    method: 'GET',
                    url: `${url}/history/${promptId}`,
                    onload: (res) => {
                        if (res.status === 200) {
                            const history = JSON.parse(res.responseText);
                            if (history[promptId]) {
                                clearInterval(poller);
                                resolve(history);
                            }
                        } else {
                            clearInterval(poller);
                            reject(new Error(`轮询ComfyUI结果时出错: ${res.statusText || res.status}`));
                        }
                    },
                    onerror: () => {
                        clearInterval(poller);
                        reject(new Error('轮询ComfyUI结果时网络错误。'));
                    }
                });
            }, POLLING_INTERVAL_MS);
        });
    }
    function findImageUrlInHistory(history, promptId, baseUrl) {
        const outputs = history[promptId]?.outputs;
        if (!outputs) return null;

        for (const nodeId in outputs) {
            if (outputs.hasOwnProperty(nodeId) && outputs[nodeId].images) {
                const image = outputs[nodeId].images[0];
                if (image) {
                    const params = new URLSearchParams({
                        filename: image.filename,
                        subfolder: image.subfolder,
                        type: image.type
                    });
                    return `${baseUrl}/view?${params.toString()}`;
                }
            }
        }
        return null;
    }
    async function displayImage(anchorElement, imageUrl) {
        let container = anchorElement.nextElementSibling;
        if (!container || !container.classList.contains('comfy-image-container')) {
            container = document.createElement('div');
            container.className = 'comfy-image-container';
            const img = document.createElement('img');
            img.alt = 'Generated by ComfyUI';
            container.appendChild(img);
            anchorElement.insertAdjacentElement('afterend', container);
        }
        const img = container.querySelector('img');
        img.src = imageUrl;
        const width = await GM_getValue('comfyui_image_width', 400);
        img.style.maxWidth = `${width}px`;
    }


    // --- 主执行逻辑 ---
    // ... 主执行逻辑没有变化 ...
    createComfyUIPanel();
    const chatObserver = new MutationObserver((mutations) => {
        const nodesToProcess = new Set();
        for (const mutation of mutations) {
            mutation.addedNodes.forEach(node => {
                if (node.nodeType === Node.ELEMENT_NODE) {
                    if (node.matches('.mes')) nodesToProcess.add(node);
                    node.querySelectorAll('.mes').forEach(mes => nodesToProcess.add(mes));
                }
            });
            if (mutation.target.closest) {
                const mesNode = mutation.target.closest('.mes');
                if (mesNode) nodesToProcess.add(mesNode);
            }
        }
        for (const node of nodesToProcess) {
            processMessageForComfyButton(node);
        }
    });
    function observeChat() {
        const chatElement = document.getElementById('chat');
        if (chatElement) {
            chatElement.querySelectorAll('.mes').forEach(processMessageForComfyButton);
            chatObserver.observe(chatElement, { childList: true, subtree: true, characterData: true });
        } else {
            setTimeout(observeChat, 500);
        }
    }
    const optionsObserver = new MutationObserver(() => {
        const optionsMenu = document.getElementById('options');
        if (optionsMenu && optionsMenu.style.display !== 'none') {
            addMainButton();
        }
    });
    window.addEventListener('load', () => {
        observeChat();
        const body = document.querySelector('body');
        if (body) {
            optionsObserver.observe(body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style'] });
        }
    });

})();