// ==UserScript==
// @name         Google AI Studio 模型注入器
// @namespace    http://tampermonkey.net/
// @version      1.6.5
// @description  向 Google AI Studio 注入自定义模型，支持主题表情图标。拦截 XHR/Fetch 请求，处理数组结构的 JSON 数据
// @author       Generated by AI / HCPTangHY / Mozi / wisdgod / UserModified
// @match        https://aistudio.google.com/*
// @icon         https://www.google.com/s2/favicons?sz=64&domain=aistudio.google.com
// @grant        none
// @run-at       document-start
// @license      MIT
// @downloadURL none
// ==/UserScript==

(function() {
    'use strict';

    // ==================== 配置区域 ====================
    // 脚本版本号已更新为 1.6.5
    const SCRIPT_VERSION = "v1.6.5";
    const LOG_PREFIX = `[AI Studio 注入器 ${SCRIPT_VERSION}]`;
    const ANTI_HIJACK_PREFIX = ")]}'\n";

    // 模型配置列表
    // 已按要求将 jfdksal98a 放到 blacktooth 的下面
    const MODELS_TO_INJECT = [
        // Blacktooth 模型 (原 Toothless)
        {
            name: 'models/blacktooth-ab-test', // 已改为 blacktooth-ab-test
            displayName: `🏴‍☠️ Blacktooth (脚本 ${SCRIPT_VERSION})`, // emoji 改为 🏴‍☠️，名称改为 Blacktooth
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        // --- jfdksal98a 模型已移动到此处 ---
        {
            name: 'models/jfdksal98a',
            displayName: `🪐 jfdksal98a (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        // --- 其他模型 ---
        {
            name: 'models/gemini-2.5-pro-preview-03-25',
            displayName: `✨ Gemini 2.5 Pro 03-25 (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        {
            name: 'models/goldmane-ab-test',
            displayName: `🦁 Goldmane (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        {
            name: 'models/claybrook-ab-test',
            displayName: `💧 Claybrook (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        {
            name: 'models/frostwind-ab-test',
            displayName: `❄️ Frostwind (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        },
        {
            name: 'models/calmriver-ab-test',
            displayName: `🌊 Calmriver (脚本 ${SCRIPT_VERSION})`,
            description: `由脚本 ${SCRIPT_VERSION} 注入的模型`
        }
    ];

    // JSON 结构中的字段索引
    const MODEL_FIELDS = {
        NAME: 0,
        DISPLAY_NAME: 3,
        DESCRIPTION: 4,
        METHODS: 7
    };

    // ==================== 工具函数 ====================

    /**
     * 检查 URL 是否为目标 API 端点
     * @param {string} url - 要检查的 URL
     * @returns {boolean}
     */
    function isTargetURL(url) {
        return url && typeof url === 'string' &&
               url.includes('alkalimakersuite') &&
               url.includes('/ListModels');
    }

    /**
     * 递归查找模型列表数组
     * @param {any} obj - 要搜索的对象
     * @returns {Array|null} 找到的模型数组或 null
     */
    function findModelListArray(obj) {
        if (!obj) return null;

        // 检查是否为目标模型数组
        if (Array.isArray(obj) && obj.length > 0 && obj.every(
            item => Array.isArray(item) &&
                    typeof item[MODEL_FIELDS.NAME] === 'string' &&
                    String(item[MODEL_FIELDS.NAME]).startsWith('models/')
        )) {
            return obj;
        }

        // 递归搜索子对象
        if (typeof obj === 'object') {
            for (const key in obj) {
                if (Object.prototype.hasOwnProperty.call(obj, key) &&
                    typeof obj[key] === 'object' &&
                    obj[key] !== null) {
                    const result = findModelListArray(obj[key]);
                    if (result) return result;
                }
            }
        }
        return null;
    }

    /**
     * 查找合适的模板模型
     * @param {Array} modelsArray - 模型数组
     * @returns {Array|null} 模板模型或 null
     */
    function findTemplateModel(modelsArray) {
        // 优先查找包含特定关键词的模型
        const templateModel =
            modelsArray.find(m => Array.isArray(m) &&
                                 m[MODEL_FIELDS.NAME] &&
                                 String(m[MODEL_FIELDS.NAME]).includes('pro') &&
                                 Array.isArray(m[MODEL_FIELDS.METHODS])) ||
            modelsArray.find(m => Array.isArray(m) &&
                                 m[MODEL_FIELDS.NAME] &&
                                 String(m[MODEL_FIELDS.NAME]).includes('flash') &&
                                 Array.isArray(m[MODEL_FIELDS.METHODS])) ||
            modelsArray.find(m => Array.isArray(m) &&
                                 m[MODEL_FIELDS.NAME] &&
                                 Array.isArray(m[MODEL_FIELDS.METHODS]));

        return templateModel;
    }

    /**
     * 更新已存在模型的显示名称
     * @param {Array} existingModel - 现有模型
     * @param {Object} modelToInject - 要注入的模型配置
     * @returns {boolean} 是否进行了更新
     */
    function updateExistingModel(existingModel, modelToInject) {
        if (!existingModel || existingModel[MODEL_FIELDS.DISPLAY_NAME] === modelToInject.displayName) {
            return false;
        }

        // 提取基础名称（去除版本号和表情）
        // 更新正则表达式以匹配 vX.Y.Z 格式
        const cleanName = (name) => String(name)
            .replace(/ \(脚本 v\d+\.\d+(\.\d+)?(-beta\d*)?\)/, '')
            // 包含所有当前使用的表情，包括新增的 🏴‍☠️, 🤖, 🪐
            .replace(/^[✨🦁💧❄️🌊🐉🏴‍☠️🤖🪐]\s*/, '')
            .trim();

        const baseExistingName = cleanName(existingModel[MODEL_FIELDS.DISPLAY_NAME]);
        const baseInjectName = cleanName(modelToInject.displayName);

        if (baseExistingName === baseInjectName) {
            // 仅更新版本号和表情
            existingModel[MODEL_FIELDS.DISPLAY_NAME] = modelToInject.displayName;
            console.log(LOG_PREFIX, `已更新表情/版本号: ${modelToInject.displayName}`);
        } else {
            // 标记为原始模型
            existingModel[MODEL_FIELDS.DISPLAY_NAME] = modelToInject.displayName + " (原始)";
            console.log(LOG_PREFIX, `已更新官方模型 ${modelToInject.name} 的显示名称`);
        }
        return true;
    }

    /**
     * 创建新模型
     * @param {Array} templateModel - 模板模型
     * @param {Object} modelToInject - 要注入的模型配置
     * @param {string} templateName - 模板名称
     * @returns {Array} 新模型数组
     */
    function createNewModel(templateModel, modelToInject, templateName) {
        const newModel = structuredClone(templateModel);

        newModel[MODEL_FIELDS.NAME] = modelToInject.name;
        newModel[MODEL_FIELDS.DISPLAY_NAME] = modelToInject.displayName;
        newModel[MODEL_FIELDS.DESCRIPTION] = `${modelToInject.description} (基于 ${templateName} 结构)`;

        if (!Array.isArray(newModel[MODEL_FIELDS.METHODS])) {
            newModel[MODEL_FIELDS.METHODS] = [
                "generateContent",
                "countTokens",
                "createCachedContent",
                "batchGenerateContent"
            ];
        }

        return newModel;
    }

    // ==================== 核心处理函数 ====================

    /**
     * 处理并修改 JSON 数据
     * @param {Object} jsonData - 原始 JSON 数据
     * @param {string} url - 请求 URL
     * @returns {Object} 包含处理后数据和修改标志的对象
     */
    function processJsonData(jsonData, url) {
        let modificationMade = false;
        const modelsArray = findModelListArray(jsonData);

        if (!modelsArray || !Array.isArray(modelsArray)) {
            console.warn(LOG_PREFIX, '在 JSON 中未找到有效的模型列表结构:', url);
            return { data: jsonData, modified: false };
        }

        // 查找模板模型
        const templateModel = findTemplateModel(modelsArray);
        const templateName = templateModel?.[MODEL_FIELDS.NAME] || 'unknown';

        if (!templateModel) {
            console.warn(LOG_PREFIX, '未找到合适的模板模型，无法注入新模型');
        }

        // 反向遍历以保持显示顺序 (配置中靠前的模型显示在最上面)
        [...MODELS_TO_INJECT].reverse().forEach(modelToInject => {
            const existingModel = modelsArray.find(
                model => Array.isArray(model) && model[MODEL_FIELDS.NAME] === modelToInject.name
            );

            if (!existingModel) {
                // 注入新模型
                if (!templateModel) {
                    console.warn(LOG_PREFIX, `无法注入 ${modelToInject.name}：缺少模板`);
                    return;
                }

                const newModel = createNewModel(templateModel, modelToInject, templateName);
                modelsArray.unshift(newModel); // unshift 将模型添加到数组开头
                modificationMade = true;
                console.log(LOG_PREFIX, `成功注入: ${modelToInject.displayName}`);
            } else {
                // 更新现有模型
                if (updateExistingModel(existingModel, modelToInject)) {
                    modificationMade = true;
                }
            }
        });

        return { data: jsonData, modified: modificationMade };
    }

    /**
     * 修改响应体
     * @param {string} originalText - 原始响应文本
     * @param {string} url - 请求 URL
     * @returns {string} 修改后的响应文本
     */
    function modifyResponseBody(originalText, url) {
        if (!originalText || typeof originalText !== 'string') {
            return originalText;
        }

        try {
            let textBody = originalText;
            let hasPrefix = false;

            // 处理反劫持前缀
            if (textBody.startsWith(ANTI_HIJACK_PREFIX)) {
                textBody = textBody.substring(ANTI_HIJACK_PREFIX.length);
                hasPrefix = true;
            }

            if (!textBody.trim()) return originalText;

            const jsonData = JSON.parse(textBody);
            const result = processJsonData(jsonData, url);

            if (result.modified) {
                let newBody = JSON.stringify(result.data);
                if (hasPrefix) {
                    newBody = ANTI_HIJACK_PREFIX + newBody;
                }
                return newBody;
            }
        } catch (error) {
            console.error(LOG_PREFIX, '处理响应体时出错:', url, error);
        }

        return originalText;
    }

    // ==================== 请求拦截 ====================

    // 拦截 Fetch API
    const originalFetch = window.fetch;
    window.fetch = async function(...args) {
        const resource = args[0];
        const url = (resource instanceof Request) ? resource.url : String(resource);
        const response = await originalFetch.apply(this, args);

        if (isTargetURL(url) && response.ok) {
            console.log(LOG_PREFIX, '[Fetch] 拦截到目标请求:', url);
            try {
                const cloneResponse = response.clone();
                const originalText = await cloneResponse.text();
                const newBody = modifyResponseBody(originalText, url);

                if (newBody !== originalText) {
                    return new Response(newBody, {
                        status: response.status,
                        statusText: response.statusText,
                        headers: response.headers
                    });
                }
            } catch (e) {
                console.error(LOG_PREFIX, '[Fetch] 处理错误:', e);
            }
        }
        return response;
    };

    // 拦截 XMLHttpRequest
    const xhrProto = XMLHttpRequest.prototype;
    const originalOpen = xhrProto.open;
    const originalResponseTextDescriptor = Object.getOwnPropertyDescriptor(xhrProto, 'responseText');
    const originalResponseDescriptor = Object.getOwnPropertyDescriptor(xhrProto, 'response');
    let interceptionCount = 0;

    // 重写 open 方法
    xhrProto.open = function(method, url) {
        this._interceptorUrl = url;
        this._isTargetXHR = isTargetURL(url);

        if (this._isTargetXHR) {
            interceptionCount++;
            console.log(LOG_PREFIX, `[XHR] 检测到目标请求 (${interceptionCount})：`, url);
        }

        return originalOpen.apply(this, arguments);
    };

    /**
     * 处理 XHR 响应
     * @param {XMLHttpRequest} xhr - XHR 对象
     * @param {any} originalValue - 原始响应值
     * @param {string} type - 响应类型
     * @returns {any} 处理后的响应值
     */
    const handleXHRResponse = (xhr, originalValue, type = 'text') => {
        if (!xhr._isTargetXHR || xhr.readyState !== 4 || xhr.status !== 200) {
            return originalValue;
        }

        const cacheKey = '_modifiedResponseCache_' + type;

        if (xhr[cacheKey] === undefined) {
            const originalText = (type === 'text' || typeof originalValue !== 'object' || originalValue === null)
                ? String(originalValue || '')
                : JSON.stringify(originalValue);

            xhr[cacheKey] = modifyResponseBody(originalText, xhr._interceptorUrl);
        }

        const cachedResponse = xhr[cacheKey];

        try {
            if (type === 'json' && typeof cachedResponse === 'string') {
                const textToParse = cachedResponse.replace(ANTI_HIJACK_PREFIX, '');
                return textToParse ? JSON.parse(textToParse) : null;
            }
        } catch (e) {
            console.error(LOG_PREFIX, '[XHR] 解析缓存的 JSON 时出错:', e);
            return originalValue;
        }

        return cachedResponse;
    };

    // 重写 responseText 属性
    if (originalResponseTextDescriptor?.get) {
        Object.defineProperty(xhrProto, 'responseText', {
            get: function() {
                const originalText = originalResponseTextDescriptor.get.call(this);

                if (this.responseType && this.responseType !== 'text' && this.responseType !== "") {
                    return originalText;
                }

                return handleXHRResponse(this, originalText, 'text');
            },
            configurable: true
        });
    }

    // 重写 response 属性
    if (originalResponseDescriptor?.get) {
        Object.defineProperty(xhrProto, 'response', {
            get: function() {
                const originalResponse = originalResponseDescriptor.get.call(this);

                if (this.responseType === 'json') {
                    return handleXHRResponse(this, originalResponse, 'json');
                }

                if (!this.responseType || this.responseType === 'text' || this.responseType === "") {
                    return handleXHRResponse(this, originalResponse, 'text');
                }

                return originalResponse;
            },
            configurable: true
        });
    }

    console.log(LOG_PREFIX, '脚本已激活，Fetch 和 XHR 拦截已启用');
})();