// ==UserScript==
// @name         公益酒馆ComfyUI插图脚本
// @namespace    http://tampermonkey.net/
// @version      26.2
// @license GPL
// @description  基于原作者@soulostar修改
// @author       feng zheng
// @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 https://update.greasyfork.icu/scripts/538457/%E5%85%AC%E7%9B%8A%E9%85%92%E9%A6%86ComfyUI%E6%8F%92%E5%9B%BE%E8%84%9A%E6%9C%AC.user.js
// @updateURL https://update.greasyfork.icu/scripts/538457/%E5%85%AC%E7%9B%8A%E9%85%92%E9%A6%86ComfyUI%E6%8F%92%E5%9B%BE%E8%84%9A%E6%9C%AC.meta.js
// ==/UserScript==


(function() {
    'use strict';

    // --- Configuration Constants ---
    const BUTTON_ID = 'comfyui-launcher-button';
    const PANEL_ID = 'comfyui-panel';
    const POLLING_TIMEOUT_MS = 120000; // 轮询超时时间 (2分钟), 增加超时以适应调度器异步处理
    const POLLING_INTERVAL_MS = 3000; // 轮询间隔 (3秒), 略微增加间隔
    const STORAGE_KEY_IMAGES = 'comfyui_generated_images';
    const STORAGE_KEY_PROMPT_PREFIX = 'comfyui_prompt_prefix'; // 提示词前缀的存储键
    const STORAGE_KEY_MAX_WIDTH = 'comfyui_image_max_width'; // 最大图片宽度的存储键
    const COOLDOWN_DURATION_MS = 60000; // 前端冷却时间 (60秒)，作为默认值或当调度器未返回具体秒数时使用

    // --- Global Cooldown Variable (default no cooldown) ---
    let globalCooldownEndTime = 0;

    // --- Cached User Settings Variables ---
    let cachedSettings = {
        comfyuiUrl: '',
        workflow: '',
        startTag: 'image###',
        endTag: '###',
        promptPrefix: '',
        maxWidth: 600
    };

    // --- Inject Custom CSS Styles ---
    GM_addStyle(`
        /* 控制面板主容器样式 */
        #${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} textarea,
        #${PANEL_ID} input[type="number"] { /* 包含数字输入框 */
            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;
            /* Modified: Changed button background to a gradient sky blue */
            background: linear-gradient(135deg, #87CEEB 0%, #00BFFF 100%); /* 天蓝色到深天蓝色渐变 */
            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 { position: relative; top: -5px; }
        .comfy-url-container { display: flex; gap: 10px; align-items: center; }
        .comfy-url-container 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; }

        /* 聊天内按钮组容器 */
        .comfy-button-group {
            display: inline-flex;
            align-items: center;
            gap: 5px;
            margin: 5px 4px;
        }

        /* 生成的图片容器样式 */
        .comfy-image-container {
            margin-top: 10px;
            max-width: 100%; /* 默认允许图片最大宽度为容器的100% */
        }
        .comfy-image-container img {
            /* 注意：这里的max-width将由JavaScript直接设置，CSS变量作为备用或默认值 */
            max-width: var(--comfy-image-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;
            }
        }

        /* 定义一个CSS变量，用于动态控制图片最大宽度 */
        :root {
            --comfy-image-max-width: 600px; /* 默认图片最大宽度 */
        }
    `);

    // A flag to prevent duplicate execution from touchstart and click
    let lastTapTimestamp = 0;
    const TAP_THRESHOLD = 300; // milliseconds to prevent double taps/clicks

    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-url-container">
                        <input id="comfyui-url" type="text" placeholder="例如: http://127.0.0.1:5001 或 http://127.0.0.1:8188">
                        <button id="comfyui-test-conn" class="comfy-button">测试连接</button>
                    </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>
                    <div>
                        <label for="comfyui-prompt-prefix">提示词固定前缀 (LoRA等):</label>
                        <input id="comfyui-prompt-prefix" type="text" placeholder="例如: <lora:cool_style:0.8> ">
                    </div>
                    <!-- 新增：最大图片宽度输入框 -->
                    <div>
                        <label for="comfyui-max-width">最大图片宽度 (px):</label>
                        <input id="comfyui-max-width" type="number" placeholder="例如: 600" min="100">
                    </div>
                    <p class="workflow-info">当使用调度器时，此处的工作流设置将被忽略。工作流由调度器决定。</p>
                    <label for="comfyui-workflow">工作流 (当不使用调度器时生效)</label>
                    <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 promptPrefixInput = document.getElementById('comfyui-prompt-prefix');
        const maxWidthInput = document.getElementById('comfyui-max-width');


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

        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('正在尝试连接服务...');

            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('连接成功！服务可用。');
                    } 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('连接超时！服务可能没有响应。');
                }
            });
        });

        clearCacheButton.addEventListener('click', () => {
            if (confirm('您确定要删除所有已生成的图片缓存吗？\n此操作不可撤销，但不会删除您本地ComfyUI输出文件夹中的文件。')) {
                GM_setValue(STORAGE_KEY_IMAGES, {});
                if (typeof toastr !== 'undefined') toastr.success('所有图片缓存已成功删除！请刷新页面以更新显示。');
            }
        });

        // 异步加载设置并应用
        loadSettings(urlInput, workflowInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput).then(() => {
            // 在设置加载完成后，将当前图片最大宽度应用到所有已存在的图片上
            applyCurrentMaxWidthToAllImages();
        });


        // 为所有输入框添加事件监听器
        [urlInput, workflowInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput].forEach(input => {
            input.addEventListener('input', async () => { // 修改为 async
                if(input === urlInput) testButton.classList.remove('success', 'error', 'testing');
                await saveSettings(urlInput, workflowInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput);
                // 每次保存设置时，也立即应用到所有已存在的图片
                applyCurrentMaxWidthToAllImages();
            });
        });
    }

    async function loadSettings(urlInput, workflowInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput) {
        cachedSettings.comfyuiUrl = await GM_getValue('comfyui_url', 'http://127.0.0.1:5001');
        cachedSettings.workflow = await GM_getValue('comfyui_workflow', '');
        cachedSettings.startTag = await GM_getValue('comfyui_start_tag', 'image###');
        cachedSettings.endTag = await GM_getValue('comfyui_end_tag', '###');
        cachedSettings.promptPrefix = await GM_getValue(STORAGE_KEY_PROMPT_PREFIX, '');
        cachedSettings.maxWidth = await GM_getValue(STORAGE_KEY_MAX_WIDTH, 600);

        urlInput.value = cachedSettings.comfyuiUrl;
        workflowInput.value = cachedSettings.workflow;
        startTagInput.value = cachedSettings.startTag;
        endTagInput.value = cachedSettings.endTag;
        promptPrefixInput.value = cachedSettings.promptPrefix;
        maxWidthInput.value = cachedSettings.maxWidth;

        document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
    }

    async function saveSettings(urlInput, workflowInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput) {
        cachedSettings.comfyuiUrl = urlInput.value.trim();
        cachedSettings.workflow = workflowInput.value;
        cachedSettings.startTag = startTagInput.value;
        cachedSettings.endTag = endTagInput.value;
        cachedSettings.promptPrefix = promptPrefixInput.value.trim();
        const newMaxWidth = parseInt(maxWidthInput.value);
        cachedSettings.maxWidth = isNaN(newMaxWidth) ? 600 : newMaxWidth;

        await GM_setValue('comfyui_url', cachedSettings.comfyuiUrl);
        await GM_setValue('comfyui_workflow', cachedSettings.workflow);
        await GM_setValue('comfyui_start_tag', cachedSettings.startTag);
        await GM_setValue('comfyui_end_tag', cachedSettings.endTag);
        await GM_setValue(STORAGE_KEY_PROMPT_PREFIX, cachedSettings.promptPrefix);
        await GM_setValue(STORAGE_KEY_MAX_WIDTH, cachedSettings.maxWidth);

        document.documentElement.style.setProperty('--comfy-image-max-width', cachedSettings.maxWidth + 'px');
    }

    /**
     * 将当前设置的最大宽度动态应用到所有已存在的图片元素上。
     * 这样做是为了确保图片大小能够即时更新，无需刷新页面。
     */
    async function applyCurrentMaxWidthToAllImages() {
        const images = document.querySelectorAll('.comfy-image-container img');
        const maxWidthPx = (cachedSettings.maxWidth || 600) + 'px';

        images.forEach(img => {
            img.style.maxWidth = maxWidthPx;
        });
    }

    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);
             }
        }
    }


    // --- 聊天消息处理与图片生成 ---

    function escapeRegex(string) {
        return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
    }

    function generateClientId() {
        return 'client-' + Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
    }

    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, imageUrl) {
        const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
        records[generationId] = imageUrl;
        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);
    }

    // 通用事件处理函数
    function handleComfyButtonClick(event, isTouch = false) {
        const button = event.target.closest('.comfy-chat-generate-button');
        if (!button) return; // 确保点击的是我们的按钮

        if (isTouch) {
            event.preventDefault(); // 阻止 300ms 延迟和潜在的双击缩放
            const now = Date.now();
            if (now - lastTapTimestamp < TAP_THRESHOLD) {
                // 这可能是快速双击或重复点击，忽略
                console.log('触摸事件被忽略：太快了。');
                return;
            }
            lastTapTimestamp = now;
            console.log('触摸事件触发：执行生成逻辑。');
            onGenerateButtonClickLogic(button);
        } else { // 这是点击事件路径
            // 如果最近的触摸事件已经触发了操作，则忽略此点击事件以防止重复执行
            if (Date.now() - lastTapTimestamp < TAP_THRESHOLD) {
                console.log('点击事件被忽略：最近由触摸事件触发。');
                return;
            }
            console.log('点击事件触发：执行生成逻辑。');
            onGenerateButtonClickLogic(button);
        }
    }

    async function processMessageForComfyButton(messageNode, savedImagesCache) { // 接收缓存的 savedImages
        const mesText = messageNode.querySelector('.mes_text');
        if (!mesText) return;

        // 使用缓存的设置
        const startTag = cachedSettings.startTag;
        const endTag = cachedSettings.endTag;
        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>`;
            });
        }

        // 使用传入的缓存 savedImagesCache
        const buttonGroups = mesText.querySelectorAll('.comfy-button-group');

        buttonGroups.forEach(group => {
            if (group.dataset.listenerAttached) return;

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

            // 检查全局冷却时间，如果有，则立即启动倒计时显示
            if (Date.now() < globalCooldownEndTime) {
                generateButton.dataset.cooldownEnd = globalCooldownEndTime.toString(); // 应用全局冷却时间
                startCooldownCountdown(generateButton, globalCooldownEndTime);
            } else if (savedImagesCache[generationId]) { // 使用 savedImagesCache
                displayImage(group, savedImagesCache[generationId]); // 使用 savedImagesCache
                setupGeneratedState(generateButton, generationId);
            }
            // 移除了直接的 addEventListener，改为依赖 chatObserver 中的事件委托
            group.dataset.listenerAttached = 'true'; // 标记已处理，防止重复修改 DOM
        });
    }

    function setupGeneratedState(generateButton, generationId) {
        generateButton.textContent = '重新生成';
        generateButton.disabled = false;
        generateButton.classList.remove('testing', 'success', 'error');
        delete generateButton.dataset.cooldownEnd; // 清除冷却标记

        // 确保重新生成事件已绑定
        // 移除了直接的 addEventListener，改为依赖 chatObserver 中的事件委托
        generateButton.dataset.regenerateListener = 'true'; // 标记已处理
        // 注意：这里不需要再添加事件监听器，因为我们将使用委托模式
        // 只要 handleComfyButtonClick 绑定在 .mes_text 上，它就会捕获所有内部按钮的事件
        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);
        }
    }

    // 核心生成逻辑函数
    async function onGenerateButtonClickLogic(button) {
        const group = button.closest('.comfy-button-group');
        let prompt = button.dataset.prompt;
        const generationId = group.dataset.generationId;

        // 在执行核心逻辑前再次检查是否被禁用，防止竞态条件
        if (button.disabled) {
            console.log('按钮已禁用，跳过生成逻辑。');
            return;
        }

        // --- GLOBAL: Cooldown Check at click time ---
        if (Date.now() < globalCooldownEndTime) {
            const remainingTime = Math.ceil((globalCooldownEndTime - Date.now()) / 1000);
            if (typeof toastr !== 'undefined') toastr.warning(`请稍候，图片生成功能正在冷却中 (${remainingTime}s)。`);
            return; // 阻止发送请求
        }
        // --- END GLOBAL Cooldown Check ---

        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;

        try {
            // 使用缓存的设置
            const url = cachedSettings.comfyuiUrl;
            let workflowString = cachedSettings.workflow;
            const promptPrefix = cachedSettings.promptPrefix;

            if (!url) throw new Error('调度器/ComfyUI URL 未配置。');

            if (promptPrefix) {
                prompt = promptPrefix + ' ' + prompt;
            }

            const clientId = generateClientId();
            let promptResponse;
            // 判断是否需要向调度器发送简化请求
            // 如果 URL 是调度器地址，并且用户没有在工作流字段填写内容，则认为是调度器模式
            const isScheduler = url.includes(':5001') || workflowString.trim() === '';

            if (isScheduler) {
                if (typeof toastr !== 'undefined') toastr.info('检测到调度器模式，正在发送简化请求...');
                // 修改此处：sendPromptRequestToScheduler 现在将处理 202 响应
                promptResponse = await sendPromptRequestToScheduler(url, {
                    client_id: clientId,
                    positive_prompt: prompt
                });

                // *** 新增：显示调度器分配的实例和队列长度 ***
                if (promptResponse.assigned_instance_name && typeof promptResponse.assigned_instance_queue_size !== 'undefined') {
                    if (typeof toastr !== 'undefined') {
                        toastr.success(`任务已分配到实例: ${promptResponse.assigned_instance_name}，当前队列长度: ${promptResponse.assigned_instance_queue_size}`);
                    }
                }

            } else {
                if (!workflowString.includes('%prompt%')) throw new Error('工作流中未找到必需的 %prompt% 占位符。');
                if (typeof toastr !== 'undefined') toastr.info('检测到直连ComfyUI模式，正在发送完整工作流...');
                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);
                promptResponse = await sendPromptRequestDirect(url, workflow, clientId);
            }

            // 对于调度器模式，promptResponse 应该包含 prompt_id
            const promptId = promptResponse.prompt_id; // 从调度器返回的 202 响应中获取 prompt_id

            if (!promptId) {
                throw new Error('调度器未返回有效的 Prompt ID。');
            }

            // 使用从调度器获得的 promptId 进行轮询
            const finalHistory = await pollForResult(url, promptId);
            const imageUrl = findImageUrlInHistory(finalHistory, promptId, url);
            if (!imageUrl) throw new Error('在服务返回结果中未找到图片。');

            // --- 成功生成新图片后，才移除旧图片并显示新图片 ---
            if (oldImageContainer) {
                oldImageContainer.remove();
            }
            displayImage(group, imageUrl); // 显示新图片
            await saveImageRecord(generationId, imageUrl); // 保存新图片记录

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

            setTimeout(() => {
                setupGeneratedState(button, generationId);
                if (deleteButton) deleteButton.style.display = 'inline-flex'; // 恢复删除按钮
            }, 2000);

        } catch (e) {
            // 记录详细错误到控制台
            console.error('ComfyUI 生图脚本错误（详细信息 - 仅供调试）:', e);

            let displayMessage = '图片生成失败，请检查调度器或ComfyUI服务。';
            let isRateLimitError = false;
            let actualCooldownSeconds = COOLDOWN_DURATION_MS / 1000; // 默认冷却时间为 60 秒

            // 尝试从错误消息中提取具体的速率限制信息
            // 调度器返回的错误信息格式： "请求频率过高，请稍后再试。限制: 1/60秒，请在 X 秒后重试。"
            const rateLimitMatch = e.message.match(/请在 (\d+) 秒后重试。/);
            if (rateLimitMatch && rateLimitMatch[1]) {
                actualCooldownSeconds = parseInt(rateLimitMatch[1], 10);
                displayMessage = `一分钟只能生成一张图片哦，请在 ${actualCooldownSeconds} 秒后重试。`;
                isRateLimitError = true;
            } else if (e.message.includes('请求频率过高')) {
                // 如果没有找到具体的秒数，但有“请求频率过高”的字样，也认为是速率限制
                displayMessage = '一分钟只能生成一张图片哦，请求频率过高，请稍后再试。'; // 给出通用提示
                isRateLimitError = true;
            } else if (e.message.includes('轮询结果超时')) {
                displayMessage = '生成任务超时，可能仍在处理或已失败。';
            } else if (e.message.includes('无法连接到调度器 API')) {
                displayMessage = '无法连接到调度器服务，请检查URL和网络。';
            } else if (e.message.includes('连接调度器 API 超时')) {
                displayMessage = '连接调度器服务超时，请检查网络。';
            } else if (e.message.includes('任务 ID 未找到或已过期')) {
                displayMessage = '生成任务ID无效或已过期，请尝试重新生成。';
            } else {
                // 对于其他未知错误，尝试提取更友好的信息
                const backendErrorMatch = e.message.match(/error:\s*"(.*?)"/);
                if (backendErrorMatch && backendErrorMatch[1]) {
                    displayMessage = `调度器错误: ${backendErrorMatch[1]}`;
                }
            }

            if (typeof toastr !== 'undefined') toastr.error(displayMessage);

            // --- 错误发生时，保持旧图片可见，并恢复删除按钮（如果存在） ---
            if (deleteButton) deleteButton.style.display = 'inline-flex'; // 恢复删除按钮

            // 检查是否为调度器速率限制错误，如果是，则启动全局冷却倒计时
            if (isRateLimitError) { // 使用 isRateLimitError 标志
                const newCooldownEndTime = Date.now() + (actualCooldownSeconds * 1000); // 使用实际的冷却秒数
                globalCooldownEndTime = newCooldownEndTime; // 设置全局冷却时间
                applyGlobalCooldown(newCooldownEndTime); // 将冷却状态应用到所有按钮
            } else {
                // 非速率限制错误，按钮显示“生成失败”并短时间后恢复
                button.textContent = '生成失败';
                button.classList.remove('testing');
                button.classList.add('error');

                setTimeout(() => {
                    const wasRegenerating = !!group.querySelector('.comfy-delete-button');
                    if (wasRegenerating) {
                        setupGeneratedState(button, generationId);
                    } else {
                         button.textContent = '开始生成';
                         button.disabled = false;
                         button.classList.remove('error');
                    }
                }, 3000);
            }
        }
    }

    /**
     * 将冷却状态应用到所有图片生成按钮。
     * @param {number} endTime - 冷却结束的时间戳 (ms)。
     */
    function applyGlobalCooldown(endTime) {
        const allGenerateButtons = document.querySelectorAll('.comfy-chat-generate-button');
        allGenerateButtons.forEach(button => {
            button.dataset.cooldownEnd = endTime.toString(); // 在每个按钮上设置冷却结束时间
            startCooldownCountdown(button, endTime); // 启动单个按钮的倒计时
        });
    }

    /**
     * 启动按钮的冷却倒计时。
     * @param {HTMLElement} button - 要冷却的按钮元素。
     * @param {number} endTime - 冷却结束的时间戳 (ms)。
     */
    function startCooldownCountdown(button, endTime) {
        button.disabled = true; // 确保按钮在冷却期间被禁用
        button.classList.remove('success', 'error', 'testing'); // 清除其他状态

        const updateCountdown = () => {
            const remainingTime = Math.max(0, endTime - Date.now());
            const seconds = Math.ceil(remainingTime / 1000);

            if (seconds > 0) {
                button.textContent = `冷却中 (${seconds}s)`;
                setTimeout(updateCountdown, 1000); // 每秒更新一次
            } else {
                // 冷却结束，恢复按钮状态
                button.disabled = false;
                delete button.dataset.cooldownEnd; // 移除冷却标记

                const group = button.closest('.comfy-button-group');
                const generationId = group.dataset.generationId;
                const deleteButtonPresent = group.querySelector('.comfy-delete-button');

                if (deleteButtonPresent) {
                    setupGeneratedState(button, generationId); // 恢复到“重新生成”状态
                } else {
                    button.textContent = '开始生成'; // 恢复到“开始生成”状态
                }
            }
        };
        updateCountdown(); // 立即执行一次以显示初始倒计时
    }


    function sendPromptRequestToScheduler(url, payload) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'POST',
                url: `${url}/generate`,
                headers: { 'Content-Type': 'application/json' },
                data: JSON.stringify(payload),
                timeout: 10000, // 从 30s 减少到 10s
                onload: (res) => {
                    // *** 关键修改: 处理 202 Accepted 状态码 ***
                    if (res.status === 202) {
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至调度器，任务已在后台排队。');
                        let responseData = {};
                        try {
                             responseData = JSON.parse(res.responseText);
                        } catch (e) {
                             console.warn('调度器 202 响应不是有效的 JSON 或为空。继续使用空响应数据。', e);
                        }
                        // 调度器 202 响应中现在应该包含 prompt_id, assigned_instance_name, assigned_instance_queue_size
                        resolve({
                            prompt_id: responseData.prompt_id,
                            message: responseData.message,
                            assigned_instance_name: responseData.assigned_instance_name, // 新增
                            assigned_instance_queue_size: responseData.assigned_instance_queue_size // 新增
                        });
                    } else if (res.status === 200) { // 兼容旧版调度器或同步返回 prompt_id 的情况
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至调度器，排队中...');
                        resolve(JSON.parse(res.responseText));
                    }
                    else {
                        let errorMessage = '';
                        try {
                            const errorJson = JSON.parse(res.responseText);
                            if (errorJson && errorJson.error) {
                                // 如果是JSON错误，直接使用其error字段
                                errorMessage = errorJson.error;
                            } else {
                                // 否则，使用状态码和原始响应文本
                                errorMessage = `调度器 API 错误: ${res.statusText || res.status} - ${res.responseText}`;
                            }
                        } catch (parseError) {
                            // 如果响应文本不是JSON，直接作为错误信息
                            errorMessage = `调度器 API 错误: ${res.statusText || res.status} - ${res.responseText}`;
                        }
                        reject(new Error(errorMessage));
                    }
                },
                onerror: (e) => reject(new Error('无法连接到调度器 API。请检查URL和网络连接。详细错误：' + (e.responseText || e.statusText || e.status))),
                ontimeout: () => reject(new Error('连接调度器 API 超时。请检查网络。')),
            });
        });
    }

    function sendPromptRequestDirect(url, workflow, clientId) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'POST',
                url: `${url}/prompt`,
                headers: { 'Content-Type': 'application/json' },
                data: JSON.stringify({ prompt: workflow, client_id: clientId }),
                timeout: 10000, // 从 30s 减少到 10s
                onload: (res) => {
                    if (res.status === 200) {
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至ComfyUI，排队中...');
                        resolve(JSON.parse(res.responseText));
                    } else {
                        // 对于直连ComfyUI的错误，保留更多细节，因为它可能不会返回统一的JSON错误格式
                        reject(new Error(`ComfyUI API 错误 (提示词): ${res.statusText || res.status} - ${res.responseText}`));
                    }
                },
                onerror: (e) => reject(new Error('无法连接到 ComfyUI API。请检查URL和网络连接。详细错误：' + (e.responseText || e.statusText || e.status))),
                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('轮询结果超时。任务可能仍在处理中或已失败。请查看调度器日志了解更多信息。'));
                    return;
                }
                GM_xmlhttpRequest({
                    method: 'GET',
                    url: `${url}/history/${promptId}`,
                    timeout: 65000, // 显式设置超时时间，略大于调度器的代理超时
                    onload: (res) => {
                        if (res.status === 200) {
                            const history = JSON.parse(res.responseText);
                            // 检查历史记录中是否存在该 promptId 并且其 outputs 不为空
                            if (history[promptId] && Object.keys(history[promptId].outputs).length > 0) {
                                clearInterval(poller);
                                resolve(history);
                            } else {
                                // 即使 200，如果 outputs 为空，也可能意味着任务仍在进行中
                                console.info(`轮询历史记录 ${promptId}: 任务仍在进行中。`); // 使用 console.info 避免频繁弹窗
                            }
                        } else if (res.status === 404) {
                             clearInterval(poller);
                             // 隐藏敏感信息
                             reject(new Error(`轮询结果失败: 任务 ID ${promptId} 未找到或已过期。请查看调度器日志了解更多信息。`));
                        }
                        else {
                            clearInterval(poller);
                            // 隐藏敏感信息，只显示状态码和通用信息
                            reject(new Error(`轮询结果失败: 后端服务返回状态码 ${res.status}。请查看调度器日志了解更多信息。`));
                        }
                    },
                    onerror: (e) => {
                        clearInterval(poller);
                        // 隐藏敏感信息，提供通用网络错误提示
                        reject(new Error('轮询结果网络错误或调度器/ComfyUI无响应。请检查网络连接或调度器日志。详细错误：' + (e.responseText || e.statusText || e.status)));
                    },
                    ontimeout: () => { // 处理单个轮询请求的超时
                        clearInterval(poller);
                        // 隐藏敏感信息
                        reject(new Error(`单个轮询请求超时。调度器在历史记录接口处无响应。请检查调度器日志了解更多信息。`));
                    }
                });
            }, 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,
                        prompt_id: promptId // 传递 prompt_id 给 /view 路由
                    });
                    return `${baseUrl}/view?${params.toString()}`;
                }
            }
        }
        return null;
    }

    /**
     * 显示图片，并确保其最大宽度设置被应用。
     * @param {HTMLElement} anchorElement - 图片按钮所在的父元素或相邻元素。
     * @param {string} imageUrl - 要显示的图片URL。
     */
    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 = 'ComfyUI 生成的图片'; // 更新 alt 文本
            container.appendChild(img);
            anchorElement.insertAdjacentElement('afterend', container);
        }
        const imgElement = container.querySelector('img');
        imgElement.src = imageUrl;

        // 直接从缓存中获取当前的最大宽度设置并应用到图片元素
        imgElement.style.maxWidth = (cachedSettings.maxWidth || 600) + 'px';
    }


    // --- 主执行逻辑 ---
    createComfyUIPanel();

    const chatObserver = new MutationObserver(async (mutations) => { // Make the callback async
        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));
                }
            });
            // 确保如果修改发生在 .mes 元素内部，也能被捕捉到
            if (mutation.target.nodeType === Node.ELEMENT_NODE && mutation.target.closest('.mes')) {
                 nodesToProcess.add(mutation.target.closest('.mes'));
            }
        }

        if (nodesToProcess.size > 0) {
            // Fetch savedImages once per batch of mutations
            const savedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});

            // Load settings into cache if not already loaded (or if forced refresh)
            // This ensures cachedSettings is always up-to-date for new messages
            await loadSettingsFromStorageAndApplyToCache(); // Ensure cachedSettings are current

            nodesToProcess.forEach(node => {
                // 在处理每个消息节点时，为其内部的 .mes_text 元素添加事件委托
                const mesTextElement = node.querySelector('.mes_text');
                if (mesTextElement && !mesTextElement.dataset.listenersAttached) {
                    mesTextElement.addEventListener('touchstart', (event) => handleComfyButtonClick(event, true), { passive: false });
                    mesTextElement.addEventListener('click', (event) => handleComfyButtonClick(event, false));
                    mesTextElement.dataset.listenersAttached = 'true'; // 标记已附加监听器
                }
                processMessageForComfyButton(node, savedImages); // Pass savedImages
            });
        }
    });

    // 独立函数，用于从存储中加载设置并更新缓存
    async function loadSettingsFromStorageAndApplyToCache() {
        const currentUrl = await GM_getValue('comfyui_url', 'http://127.0.0.1:5001');
        const currentWorkflow = await GM_getValue('comfyui_workflow', '');
        const currentStartTag = await GM_getValue('comfyui_start_tag', 'image###');
        const currentEndTag = await GM_getValue('comfyui_end_tag', '###');
        const currentPromptPrefix = await GM_getValue(STORAGE_KEY_PROMPT_PREFIX, '');
        const currentMaxWidth = await GM_getValue(STORAGE_KEY_MAX_WIDTH, 600);

        cachedSettings.comfyuiUrl = currentUrl;
        cachedSettings.workflow = currentWorkflow;
        cachedSettings.startTag = currentStartTag;
        cachedSettings.endTag = currentEndTag;
        cachedSettings.promptPrefix = currentPromptPrefix;
        cachedSettings.maxWidth = currentMaxWidth;

        // Apply max width to CSS variable immediately after loading/updating cache
        document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
    }


    function observeChat() {
        const chatElement = document.getElementById('chat');
        if (chatElement) {
            // On initial load, ensure settings are loaded and then process existing messages
            loadSettingsFromStorageAndApplyToCache().then(async () => {
                const initialSavedImages = await GM_getValue(STORAGE_KEY_IMAGES, {}); // Initial fetch of saved images
                chatElement.querySelectorAll('.mes').forEach(node => {
                    // 为已存在的聊天消息添加事件委托
                    const mesTextElement = node.querySelector('.mes_text');
                    if (mesTextElement && !mesTextElement.dataset.listenersAttached) {
                        mesTextElement.addEventListener('touchstart', (event) => handleComfyButtonClick(event, true), { passive: false });
                        mesTextElement.addEventListener('click', (event) => handleComfyButtonClick(event, false));
                        mesTextElement.dataset.listenersAttached = 'true';
                    }
                    processMessageForComfyButton(node, initialSavedImages)
                });
                // Adjust MutationObserver options, remove characterData: true
                chatObserver.observe(chatElement, { childList: true, subtree: 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'] });
        }
    });

})();