// The main script for the extension
// The following are examples of some basic extension functionality

//You'll likely need to import extension_settings, getContext, and loadExtensionSettings from extensions.js
import { extension_settings, getContext } from '../../../extensions.js';
//You'll likely need to import some other functions from the main script
import {
    saveSettingsDebounced,
    eventSource,
    event_types,
    updateMessageBlock,
    ensureMessageMediaIsArray,
} from '../../../../script.js';
import { appendMediaToMessage } from '../../../../script.js';
import { regexFromString } from '../../../utils.js';
import { SlashCommandParser } from '../../../slash-commands/SlashCommandParser.js';
import { findSecret, SECRET_KEYS } from '../../../secrets.js';

// 扩展名称和路径
const extensionName = 'st-image-auto-generation';
// /scripts/extensions/third-party
const extensionFolderPath = `/scripts/extensions/third-party/${extensionName}`;
const apiKeyStorageKey = `${extensionName}.apiKey`;

// 插入类型常量
const INSERT_TYPE = {
    DISABLED: 'disabled',
    INLINE: 'inline',
    NEW_MESSAGE: 'new',
    REPLACE: 'replace',
};

/**
 * Escapes characters for safe inclusion inside HTML attribute values.
 * @param {string} value
 * @returns {string}
 */
function escapeHtmlAttribute(value) {
    if (typeof value !== 'string') {
        return '';
    }

    return value
        .replace(/&/g, '&')
        .replace(/"/g, '"')
        .replace(/'/g, ''')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;');
}

// 默认设置
const defaultSettings = {
    insertType: INSERT_TYPE.DISABLED,
    customApi: {
        enabled: true,
        baseUrl: 'https://freeapi.dgbmc.top/v1',
        model: 'grok-imagine-image',
        size: '1024x1024',
        responseFormat: 'auto',
    },
    promptInjection: {
        enabled: true,
        prompt: `<image_generation>
You must insert a <pic prompt="example prompt"> at end of the reply. Prompts are used for stable diffusion image generation, based on the plot and character to output appropriate prompts to generate captivating images.
</image_generation>`,
        regex: '/<pic[^>]*\\sprompt="([^"]*)"[^>]*?>/g',
        position: 'deep_system', // deep_system, deep_user, deep_assistant
        depth: 0, // 0表示添加到末尾，>0表示从末尾往前数第几个位置
    },
};

/**
 * Reads the API key stored for this extension in the current browser.
 * It is deliberately kept out of extension_settings so shared settings
 * exports and plugin packages do not contain the key.
 * @returns {string}
 */
function getStoredApiKey() {
    try {
        return window.localStorage.getItem(apiKeyStorageKey) || '';
    } catch {
        return '';
    }
}

/**
 * Stores or clears the API key for this extension in the current browser.
 * @param {string} value
 */
function setStoredApiKey(value) {
    try {
        const normalizedValue = String(value || '').trim();
        if (normalizedValue) {
            window.localStorage.setItem(apiKeyStorageKey, normalizedValue);
        } else {
            window.localStorage.removeItem(apiKeyStorageKey);
        }
    } catch (error) {
        console.warn(`[${extensionName}] 无法保存图片 API Key:`, error);
    }
}

// 从设置更新UI
function updateUI() {
    const settings = extension_settings[extensionName];

    if (!settings) {
        return;
    }

    // 根据insertType设置开关状态
    $('#auto_generation').toggleClass(
        'selected',
        settings.insertType !== INSERT_TYPE.DISABLED,
    );

    // 只在表单元素存在时更新它们
    if ($('#image_generation_insert_type').length) {
        $('#image_generation_insert_type').val(
            settings.insertType,
        );
        $('#custom_image_api_enabled').prop('checked', settings.customApi.enabled);
        $('#custom_image_api_base_url').val(settings.customApi.baseUrl);
        $('#custom_image_api_model').val(settings.customApi.model);
        $('#custom_image_api_size').val(settings.customApi.size);
        $('#custom_image_api_response_format').val(settings.customApi.responseFormat);
        $('#custom_image_api_key').val(getStoredApiKey());
        $('#custom_image_api_settings').toggle(settings.customApi.enabled);
        $('#prompt_injection_enabled').prop(
            'checked',
            settings.promptInjection.enabled,
        );
        $('#prompt_injection_text').val(
            settings.promptInjection.prompt,
        );
        $('#prompt_injection_regex').val(
            settings.promptInjection.regex,
        );
        $('#prompt_injection_position').val(
            settings.promptInjection.position,
        );
        $('#prompt_injection_depth').val(
            settings.promptInjection.depth,
        );
    }
}

// 加载设置
async function loadSettings() {
    extension_settings[extensionName] = extension_settings[extensionName] || {};
    const settings = extension_settings[extensionName];

    if (settings.insertType === undefined) {
        settings.insertType = defaultSettings.insertType;
    }

    settings.customApi = {
        ...defaultSettings.customApi,
        ...(settings.customApi || {}),
    };
    settings.customApi.enabled = settings.customApi.enabled !== false;
    settings.customApi.baseUrl = String(
        settings.customApi.baseUrl || defaultSettings.customApi.baseUrl,
    ).trim();
    settings.customApi.model = String(
        settings.customApi.model || defaultSettings.customApi.model,
    ).trim();
    settings.customApi.size = String(
        settings.customApi.size || defaultSettings.customApi.size,
    ).trim();
    if (!['auto', 'url', 'b64_json'].includes(settings.customApi.responseFormat)) {
        settings.customApi.responseFormat = defaultSettings.customApi.responseFormat;
    }

    settings.promptInjection = {
        ...defaultSettings.promptInjection,
        ...(settings.promptInjection || {}),
    };

    updateUI();
}

// 创建设置页面
async function createSettings(settingsHtml) {
    // 创建一个容器来存放设置，确保其正确显示在扩展设置面板中
    if (!$('#image_auto_generation_container').length) {
        $('#extensions_settings2').append(
            '<div id="image_auto_generation_container" class="extension_container"></div>',
        );
    }

    // 使用传入的settingsHtml而不是重新获取
    $('#image_auto_generation_container').empty().append(settingsHtml);

    // 添加设置变更事件处理
    $('#image_generation_insert_type').on('change', function () {
        const newValue = $(this).val();
        extension_settings[extensionName].insertType = newValue;
        updateUI();
        saveSettingsDebounced();
    });

    $('#custom_image_api_enabled').on('change', function () {
        extension_settings[extensionName].customApi.enabled = $(this).prop('checked');
        updateUI();
        saveSettingsDebounced();
    });

    $('#custom_image_api_base_url').on('input', function () {
        extension_settings[extensionName].customApi.baseUrl = String($(this).val()).trim();
        saveSettingsDebounced();
    });

    $('#custom_image_api_model').on('input', function () {
        extension_settings[extensionName].customApi.model = String($(this).val()).trim();
        saveSettingsDebounced();
    });

    $('#custom_image_api_size').on('change', function () {
        extension_settings[extensionName].customApi.size = String($(this).val()).trim();
        saveSettingsDebounced();
    });

    $('#custom_image_api_response_format').on('change', function () {
        extension_settings[extensionName].customApi.responseFormat = $(this).val();
        saveSettingsDebounced();
    });

    $('#custom_image_api_key').on('input', function () {
        setStoredApiKey($(this).val());
    });

    $('#custom_image_api_key_clear').on('click', function () {
        setStoredApiKey('');
        $('#custom_image_api_key').val('');
        toastr.success('图片 API Key 已清除');
    });

    // 添加提示词注入设置的事件处理
    $('#prompt_injection_enabled').on('change', function () {
        extension_settings[extensionName].promptInjection.enabled =
            $(this).prop('checked');
        saveSettingsDebounced();
    });

    $('#prompt_injection_text').on('input', function () {
        extension_settings[extensionName].promptInjection.prompt =
            $(this).val();
        saveSettingsDebounced();
    });

    $('#prompt_injection_regex').on('input', function () {
        extension_settings[extensionName].promptInjection.regex = $(this).val();
        saveSettingsDebounced();
    });

    $('#prompt_injection_position').on('change', function () {
        extension_settings[extensionName].promptInjection.position =
            $(this).val();
        saveSettingsDebounced();
    });

    // 深度设置事件处理
    $('#prompt_injection_depth').on('input', function () {
        const value = parseInt(String($(this).val()));
        extension_settings[extensionName].promptInjection.depth = isNaN(value)
            ? 0
            : value;
        saveSettingsDebounced();
    });

    // 初始化设置值
    updateUI();
}

// 设置变更处理函数
function onExtensionButtonClick() {
    // 直接访问扩展设置面板
    const extensionsDrawer = $('#extensions-settings-button .drawer-toggle');

    // 如果抽屉是关闭的，点击打开它
    if ($('#rm_extensions_block').hasClass('closedDrawer')) {
        extensionsDrawer.trigger('click');
    }

    // 等待抽屉打开后滚动到我们的设置容器
    setTimeout(() => {
        // 找到我们的设置容器
        const container = $('#image_auto_generation_container');
        if (container.length) {
            // 滚动到设置面板位置
            $('#rm_extensions_block').animate(
                {
                    scrollTop:
                        container.offset().top -
                        $('#rm_extensions_block').offset().top +
                        $('#rm_extensions_block').scrollTop(),
                },
                500,
            );

            // 使用SillyTavern原生的抽屉展开方式
            // 检查抽屉内容是否可见
            const drawerContent = container.find('.inline-drawer-content');
            const drawerHeader = container.find('.inline-drawer-header');

            // 只有当内容被隐藏时才触发展开
            if (drawerContent.is(':hidden') && drawerHeader.length) {
                // 直接使用原生点击事件触发，而不做任何内部处理
                drawerHeader.trigger('click');
            }
        }
    }, 500);
}

// 初始化扩展
$(function () {
    (async function () {
        // 获取设置HTML (只获取一次)
        const settingsHtml = await $.get(
            `${extensionFolderPath}/settings.html`,
        );

        // 添加扩展到菜单
        $('#extensionsMenu')
            .append(`<div id="auto_generation" class="list-group-item flex-container flexGap5">
            <div class="fa-solid fa-robot"></div>
            <span data-i18n="Image Auto Generation">Image Auto Generation</span>
        </div>`);

        // 修改点击事件，打开设置面板而不是切换状态
        $('#auto_generation').off('click').on('click', onExtensionButtonClick);

        await loadSettings();

        // 创建设置 - 将获取的HTML传递给createSettings
        await createSettings(settingsHtml);

        // 确保设置面板可见时，设置值是正确的
        $('#extensions-settings-button').on('click', function () {
            setTimeout(() => {
                updateUI();
            }, 200);
        });
    })();
});
// 获取消息角色
function getMesRole() {
    // 确保对象路径存在
    if (
        !extension_settings[extensionName] ||
        !extension_settings[extensionName].promptInjection ||
        !extension_settings[extensionName].promptInjection.position
    ) {
        return 'system'; // 默认返回system角色
    }

    switch (extension_settings[extensionName].promptInjection.position) {
        case 'deep_system':
            return 'system';
        case 'deep_user':
            return 'user';
        case 'deep_assistant':
            return 'assistant';
        default:
            return 'system';
    }
}

/**
 * Builds an OpenAI-compatible image generation endpoint from a base URL.
 * A full /images/generations URL is also accepted for convenience.
 * @param {string} baseUrl
 * @returns {string}
 */
function getImageGenerationUrl(baseUrl) {
    const normalizedBaseUrl = String(baseUrl || '').trim().replace(/\/+$/, '');
    if (!normalizedBaseUrl) {
        throw new Error('自定义图片 API 地址不能为空');
    }

    return /\/images\/generations$/i.test(normalizedBaseUrl)
        ? normalizedBaseUrl
        : `${normalizedBaseUrl}/images/generations`;
}

/**
 * Reads an error body without leaking the API key into the error message.
 * @param {Response} response
 * @returns {Promise<string>}
 */
async function getImageApiError(response) {
    const text = await response.text();
    if (!text) {
        return `${response.status} ${response.statusText}`.trim();
    }

    try {
        const data = JSON.parse(text);
        return (
            data?.error?.message ||
            (typeof data?.error === 'string' ? data.error : '') ||
            data?.message ||
            text
        );
    } catch {
        return text;
    }
}

/**
 * Extracts a browser-loadable image URL from an OpenAI-compatible response.
 * @param {any} data
 * @returns {string}
 */
function extractImageUrl(data) {
    const item = Array.isArray(data?.data)
        ? data.data[0]
        : Array.isArray(data?.images)
            ? data.images[0]
            : data?.data;

    if (typeof item === 'string' && item.trim()) {
        return item.trim();
    }

    const imageUrl = item?.url || item?.image_url || data?.url;
    if (typeof imageUrl === 'string' && imageUrl.trim()) {
        return imageUrl.trim();
    }

    const base64 = item?.b64_json || item?.base64 || data?.b64_json;
    if (typeof base64 === 'string' && base64.trim()) {
        if (base64.startsWith('data:')) {
            return base64;
        }

        const format = String(item?.format || data?.format || 'png')
            .replace(/[^a-z0-9.+-]/gi, '')
            .toLowerCase() || 'png';
        return `data:image/${format};base64,${base64}`;
    }

    return '';
}

/**
 * Generates an image through a user-configured OpenAI-compatible endpoint.
 * @param {string} prompt
 * @returns {Promise<string>}
 */
async function generateCustomImage(prompt) {
    const customApi = extension_settings[extensionName].customApi;
    const apiKey = getStoredApiKey() || await findSecret(SECRET_KEYS.GENERIC);
    const headers = { 'Content-Type': 'application/json' };
    const model = String(customApi.model || '').trim();
    const size = String(customApi.size || '').trim();

    if (!model) {
        throw new Error('自定义图片模型 ID 不能为空');
    }

    if (!size) {
        throw new Error('图片尺寸不能为空');
    }

    if (apiKey) {
        const normalizedKey = String(apiKey).trim();
        headers.Authorization = /^bearer\s/i.test(normalizedKey)
            ? normalizedKey
            : `Bearer ${normalizedKey}`;
    }

    const body = {
        prompt: prompt.trim(),
        model,
        size,
        n: 1,
    };

    if (customApi.responseFormat !== 'auto') {
        body.response_format = customApi.responseFormat;
    }

    const response = await fetch(getImageGenerationUrl(customApi.baseUrl), {
        method: 'POST',
        headers,
        body: JSON.stringify(body),
    });

    if (!response.ok) {
        throw new Error(await getImageApiError(response));
    }

    const imageUrl = extractImageUrl(await response.json());
    if (!imageUrl) {
        throw new Error('图片 API 返回中没有找到 url 或 b64_json');
    }

    return imageUrl;
}

/**
 * Generates an image using either the custom endpoint or SillyTavern's native
 * stable-diffusion command.
 * @param {string} prompt
 * @param {string} insertType
 * @returns {Promise<string>}
 */
async function generateImage(prompt, insertType) {
    const settings = extension_settings[extensionName];
    if (settings.customApi.enabled) {
        return await generateCustomImage(prompt);
    }

    const command = SlashCommandParser.commands['sd'];
    if (!command?.callback) {
        throw new Error('酒馆图片生成命令尚未加载，请刷新页面后重试');
    }

    return await command.callback(
        {
            quiet: insertType === INSERT_TYPE.NEW_MESSAGE ? 'false' : 'true',
        },
        prompt,
    );
}

/**
 * Adds a custom-API result as a new assistant message.
 * @param {any} context
 * @param {string} prompt
 * @param {string} imageUrl
 * @returns {Promise<void>}
 */
async function addImageMessage(context, prompt, imageUrl) {
    const message = {
        name: context.name2 || 'Image Generation',
        is_user: false,
        is_system: false,
        send_date: new Date().toISOString(),
        mes: '',
        extra: {
            media: [
                {
                    type: 'image',
                    url: imageUrl,
                    title: prompt,
                    source: 'generated',
                },
            ],
            media_display: 'gallery',
            media_index: 0,
            inline_image: false,
        },
    };

    context.chat.push(message);
    const messageId = context.chat.length - 1;
    await eventSource.emit(event_types.MESSAGE_RECEIVED, messageId, 'extension');
    context.addOneMessage(message);
    await eventSource.emit(event_types.CHARACTER_MESSAGE_RENDERED, messageId, 'extension');
    await context.saveChat();
}

// 监听CHAT_COMPLETION_PROMPT_READY事件以注入提示词
eventSource.on(
    event_types.CHAT_COMPLETION_PROMPT_READY,
    async function (eventData) {
        try {
            // 确保设置对象和promptInjection对象都存在
            if (
                !extension_settings[extensionName] ||
                !extension_settings[extensionName].promptInjection ||
                !extension_settings[extensionName].promptInjection.enabled ||
                extension_settings[extensionName].insertType ===
                    INSERT_TYPE.DISABLED
            ) {
                return;
            }

            const prompt =
                extension_settings[extensionName].promptInjection.prompt;
            const depth =
                extension_settings[extensionName].promptInjection.depth || 0;
            const role = getMesRole();

            console.log(
                `[${extensionName}] 准备注入提示词: 角色=${role}, 深度=${depth}`,
            );
            console.log(
                `[${extensionName}] 提示词内容: ${prompt.substring(0, 50)}...`,
            );

            // 根据depth参数决定插入位置
            if (depth === 0) {
                // 添加到末尾
                eventData.chat.push({ role: role, content: prompt });
                console.log(`[${extensionName}] 提示词已添加到聊天末尾`);
            } else {
                // 从末尾向前插入
                eventData.chat.splice(-depth, 0, {
                    role: role,
                    content: prompt,
                });
                console.log(
                    `[${extensionName}] 提示词已插入到聊天中，从末尾往前第 ${depth} 个位置`,
                );
            }
        } catch (error) {
            console.error(`[${extensionName}] 提示词注入错误:`, error);
            toastr.error(`提示词注入错误: ${error}`);
        }
    },
);

// 监听消息接收事件
eventSource.on(event_types.MESSAGE_RECEIVED, handleIncomingMessage);
async function handleIncomingMessage() {
    // 确保设置对象存在
    if (
        !extension_settings[extensionName] ||
        extension_settings[extensionName].insertType === INSERT_TYPE.DISABLED
    ) {
        return;
    }

    const context = getContext();
    const message = context.chat[context.chat.length - 1];

    // 检查是否是AI消息
    if (!message || message.is_user) {
        return;
    }

    // 确保promptInjection对象和regex属性存在
    if (
        !extension_settings[extensionName].promptInjection ||
        !extension_settings[extensionName].promptInjection.regex ||
        typeof message.mes !== 'string'
    ) {
        console.error('Prompt injection settings not properly initialized');
        return;
    }

    // 使用正则表达式search
    const imgTagRegex = regexFromString(
        extension_settings[extensionName].promptInjection.regex,
    );
    // const testRegex = regexFromString(extension_settings[extensionName].promptInjection.regex);
    let matches;
    if (imgTagRegex.global) {
        matches = [...message.mes.matchAll(imgTagRegex)];
    } else {
        const singleMatch = message.mes.match(imgTagRegex);
        matches = singleMatch ? [singleMatch] : [];
    }
    console.log(imgTagRegex, matches);
    if (matches.length > 0) {
        // 延迟执行图片生成，确保消息首先显示出来
        setTimeout(async () => {
            try {
                toastr.info(`Generating ${matches.length} images...`);
                const insertType = extension_settings[extensionName].insertType;

                const messageId = context.chat.indexOf(message);
                if (messageId < 0 || context.chat[messageId] !== message) {
                    throw new Error('当前聊天已发生变化，已取消图片插入');
                }

                // 在当前消息中插入图片
                if (!message.extra) {
                    message.extra = {};
                }
                ensureMessageMediaIsArray(message);
                if (!Array.isArray(message.extra.media)) {
                    message.extra.media = [];
                }

                // 获取消息元素用于稍后更新
                const messageElement = $(`.mes[mesid="${messageId}"]`);

                // 处理每个匹配的图片标签
                for (const match of matches) {
                    const prompt =
                        typeof match?.[1] === 'string' ? match[1] : '';
                    if (!prompt.trim()) {
                        continue;
                    }

                    // @ts-ignore
                    const imageUrl = await generateImage(prompt, insertType);
                    if (typeof imageUrl !== 'string' || !imageUrl.trim()) {
                        throw new Error('图片生成接口返回了空结果');
                    }

                    const customApiEnabled = extension_settings[extensionName].customApi.enabled;

                    if (insertType === INSERT_TYPE.NEW_MESSAGE) {
                        // Native /sd already creates the new message when quiet=false.
                        if (customApiEnabled) {
                            await addImageMessage(context, prompt, imageUrl);
                        }
                        continue;
                    }

                    // 统一插入到extra里
                    if (insertType === INSERT_TYPE.INLINE) {
                        message.extra.media.push({
                            type: 'image',
                            url: imageUrl,
                            title: prompt,
                            source: 'generated',
                        });
                        message.extra.media_display = 'gallery';
                        message.extra.media_index = message.extra.media.length - 1;
                        message.extra.inline_image = true;

                        // 更新UI
                        appendMediaToMessage(message, messageElement);

                        // 保存聊天记录
                        await context.saveChat();
                    } else if (insertType === INSERT_TYPE.REPLACE) {
                        // Find the original image tag in the message
                        const originalTag =
                            typeof match?.[0] === 'string' ? match[0] : '';
                        if (!originalTag) {
                            continue;
                        }

                        // Replace it with an actual image tag
                        const escapedUrl = escapeHtmlAttribute(imageUrl);
                        const escapedPrompt = escapeHtmlAttribute(prompt);
                        const newImageTag = `<img src="${escapedUrl}" title="${escapedPrompt}" alt="${escapedPrompt}">`;
                        message.mes = message.mes.replace(originalTag, newImageTag);

                        // Update the message display using updateMessageBlock
                        updateMessageBlock(messageId, message);
                        await eventSource.emit(event_types.MESSAGE_UPDATED, messageId);

                        // Save the chat
                        await context.saveChat();
                    }
                }
                toastr.success(
                    `${matches.length} images generated successfully`,
                );
            } catch (error) {
                toastr.error(`Image generation error: ${error}`);
                console.error('Image generation error:', error);
            }
        }, 0); //防阻塞UI渲染
    }
}