// ST-BME: 主入口
// 事件钩子、设置管理、流程调度

import { extension_settings, getContext, saveMetadataDebounced } from '../../extensions.js';
import { eventSource, event_types, saveSettingsDebounced } from '../../../script.js';
import { renderExtensionTemplateAsync } from '../../extensions.js';

import { createEmptyGraph, deserializeGraph, serializeGraph, exportGraph, importGraph, getGraphStats } from './graph.js';
import { DEFAULT_NODE_SCHEMA, validateSchema } from './schema.js';
import { retrieve } from './retriever.js';
import { extractMemories, generateSynopsis } from './extractor.js';
import { compressAll, sleepCycle } from './compressor.js';
import { formatInjection, estimateTokens } from './injector.js';
import { testConnection as testEmbeddingConnection } from './embedding.js';
import { evolveMemories } from './evolution.js';

const MODULE_NAME = 'st_bme';
const GRAPH_METADATA_KEY = 'st_bme_graph';

// ==================== 默认设置 ====================

const defaultSettings = {
    enabled: false,

    // 提取设置
    extractEvery: 1,              // 每 N 条 assistant 回复提取一次
    extractContextTurns: 2,       // 提取时包含的上下文楼层数

    // 召回设置
    recallEnabled: true,
    recallTopK: 15,               // 混合评分 Top-K
    recallMaxNodes: 8,            // LLM 召回最大节点数
    recallEnableLLM: true,        // 是否启用 LLM 精确召回

    // 注入设置
    injectPosition: 'atDepth',    // 注入位置
    injectDepth: 4,               // 注入深度（atDepth 模式）
    injectRole: 0,                // 0=system, 1=user, 2=assistant

    // 混合评分权重
    graphWeight: 0.6,
    vectorWeight: 0.3,
    importanceWeight: 0.1,

    // Embedding API 配置
    embeddingApiUrl: '',
    embeddingApiKey: '',
    embeddingModel: 'text-embedding-3-small',

    // Schema
    nodeTypeSchema: null,         // null 表示使用默认

    // 自定义提示词
    extractPrompt: '',

    // ====== v2 增强设置 ======

    // ③ A-MEM 记忆进化
    enableEvolution: true,        // 启用记忆进化
    evoNeighborCount: 5,          // 近邻搜索数量
    evoConsolidateEvery: 50,      // 每 N 次进化后整理

    // ② Mem0 精确对照
    enablePreciseConflict: true,  // 启用精确对照
    conflictThreshold: 0.85,      // 相似度阈值

    // ⑨ 全局故事概要
    enableSynopsis: true,         // 启用全局概要
    synopsisEveryN: 5,            // 每 N 次提取后更新概要

    // ⑥ 认知边界过滤（P1）
    enableVisibility: false,      // 启用认知边界
    // ⑦ 双记忆交叉检索（P1）
    enableCrossRecall: false,     // 启用交叉检索

    // ① 惊奇度分割（P2）
    enableSmartTrigger: false,    // 启用惊奇度分割
    triggerPatterns: '',          // 自定义触发正则

    // ⑤ 主动遗忘（P2）
    enableSleepCycle: false,      // 启用主动遗忘
    forgetThreshold: 0.5,         // 保留价值阈值
    sleepEveryN: 10,              // 每 N 次提取后执行

    // ⑧ 概率触发回忆（P2）
    enableProbRecall: false,      // 启用概率触发
    probRecallChance: 0.15,       // 触发概率

    // ⑩ 反思条目（P2）
    enableReflection: false,      // 启用反思
    reflectEveryN: 10,            // 每 N 次提取后反思
};

// ==================== 状态 ====================

let currentGraph = null;
let isExtracting = false;
let isRecalling = false;
let lastInjectionContent = '';
let extractionCount = 0;    // v2: 提取次数计数器（定期触发概要/遗忘/反思）

// ==================== 设置管理 ====================

function getSettings() {
    if (!extension_settings[MODULE_NAME]) {
        extension_settings[MODULE_NAME] = { ...defaultSettings };
    }
    return extension_settings[MODULE_NAME];
}

function getSchema() {
    const settings = getSettings();
    return settings.nodeTypeSchema || DEFAULT_NODE_SCHEMA;
}

function getEmbeddingConfig() {
    const settings = getSettings();
    return {
        apiUrl: settings.embeddingApiUrl,
        apiKey: settings.embeddingApiKey,
        model: settings.embeddingModel,
    };
}

// ==================== 图状态持久化 ====================

function loadGraphFromChat() {
    const context = getContext();
    if (!context.chatMetadata) {
        currentGraph = createEmptyGraph();
        return;
    }

    const savedData = context.chatMetadata[GRAPH_METADATA_KEY];
    if (savedData) {
        currentGraph = deserializeGraph(savedData);
        console.log('[ST-BME] 从聊天数据加载图谱:', getGraphStats(currentGraph));
    } else {
        currentGraph = createEmptyGraph();
    }
}

function saveGraphToChat() {
    const context = getContext();
    if (!context.chatMetadata || !currentGraph) return;

    context.chatMetadata[GRAPH_METADATA_KEY] = currentGraph;
    saveMetadataDebounced();
}

// ==================== 核心流程 ====================

/**
 * 提取管线：处理未提取的对话楼层
 */
async function runExtraction() {
    if (isExtracting || !currentGraph) return;

    const settings = getSettings();
    if (!settings.enabled) return;

    const context = getContext();
    const chat = context.chat;
    if (!chat || chat.length === 0) return;

    // 找出 assistant 楼层序号
    const assistantTurns = [];
    for (let i = 0; i < chat.length; i++) {
        if (chat[i].is_user === false && !chat[i].is_system) {
            assistantTurns.push(i);
        }
    }

    const lastProcessed = currentGraph.lastProcessedSeq;
    const unprocessedStarts = assistantTurns.filter(i => i > lastProcessed);

    if (unprocessedStarts.length === 0) return;

    // 按 extractEvery 批次处理
    if (unprocessedStarts.length < settings.extractEvery) return;

    isExtracting = true;

    try {
        // 收集要处理的消息
        const startIdx = unprocessedStarts[0];
        const endIdx = unprocessedStarts[unprocessedStarts.length - 1];

        // 包含上下文
        const contextStart = Math.max(0, startIdx - settings.extractContextTurns * 2);
        const messages = [];
        for (let i = contextStart; i <= endIdx && i < chat.length; i++) {
            const msg = chat[i];
            if (msg.is_system) continue;
            messages.push({
                role: msg.is_user ? 'user' : 'assistant',
                content: msg.mes || '',
            });
        }

        console.log(`[ST-BME] 开始提取: 楼层 ${startIdx}-${endIdx}`);

        const result = await extractMemories({
            graph: currentGraph,
            messages,
            startSeq: endIdx,
            schema: getSchema(),
            embeddingConfig: getEmbeddingConfig(),
            extractPrompt: settings.extractPrompt || undefined,
            v2Options: {
                enablePreciseConflict: settings.enablePreciseConflict,
                conflictThreshold: settings.conflictThreshold,
            },
        });

        if (result.success) {
            extractionCount++;

            // v2: A-MEM 记忆进化
            if (settings.enableEvolution && result.newNodeIds?.length > 0) {
                try {
                    await evolveMemories({
                        graph: currentGraph,
                        newNodeIds: result.newNodeIds,
                        embeddingConfig: getEmbeddingConfig(),
                        options: { neighborCount: settings.evoNeighborCount },
                    });
                } catch (e) {
                    console.error('[ST-BME] 记忆进化失败:', e);
                }
            }

            // v2: 全局故事概要（每 N 次提取更新一次）
            if (settings.enableSynopsis && extractionCount % settings.synopsisEveryN === 0) {
                try {
                    await generateSynopsis({
                        graph: currentGraph,
                        schema: getSchema(),
                        currentSeq: endIdx,
                    });
                } catch (e) {
                    console.error('[ST-BME] 概要生成失败:', e);
                }
            }

            // v2: 主动遗忘（每 N 次提取执行）
            if (settings.enableSleepCycle && extractionCount % settings.sleepEveryN === 0) {
                try {
                    sleepCycle(currentGraph, settings);
                } catch (e) {
                    console.error('[ST-BME] 主动遗忘失败:', e);
                }
            }

            // 压缩检查
            await compressAll(currentGraph, getSchema(), getEmbeddingConfig());
            saveGraphToChat();
        }
    } catch (e) {
        console.error('[ST-BME] 提取失败:', e);
    } finally {
        isExtracting = false;
    }
}

/**
 * 召回管线：检索并注入记忆
 */
async function runRecall() {
    if (isRecalling || !currentGraph) return;

    const settings = getSettings();
    if (!settings.enabled || !settings.recallEnabled) return;

    const context = getContext();
    const chat = context.chat;
    if (!chat || chat.length === 0) return;

    isRecalling = true;

    try {
        // 获取最新用户消息
        let userMessage = '';
        const recentMessages = [];

        for (let i = chat.length - 1; i >= 0 && recentMessages.length < 4; i--) {
            const msg = chat[i];
            if (msg.is_system) continue;

            if (msg.is_user && !userMessage) {
                userMessage = msg.mes || '';
            }
            recentMessages.unshift(`[${msg.is_user ? 'user' : 'assistant'}]: ${msg.mes || ''}`);
        }

        if (!userMessage) return;

        console.log('[ST-BME] 开始召回');

        const result = await retrieve({
            graph: currentGraph,
            userMessage,
            recentMessages,
            embeddingConfig: getEmbeddingConfig(),
            schema: getSchema(),
            options: {
                topK: settings.recallTopK,
                maxRecallNodes: settings.recallMaxNodes,
                enableLLMRecall: settings.recallEnableLLM,
                weights: {
                    graphWeight: settings.graphWeight,
                    vectorWeight: settings.vectorWeight,
                    importanceWeight: settings.importanceWeight,
                },
                // v2 options
                enableVisibility: settings.enableVisibility ?? false,
                visibilityFilter: context.name2 || null,
                enableCrossRecall: settings.enableCrossRecall ?? false,
                enableProbRecall: settings.enableProbRecall ?? false,
                probRecallChance: settings.probRecallChance ?? 0.15,
            },
        });

        // 格式化注入文本
        const injectionText = formatInjection(result, getSchema());
        lastInjectionContent = injectionText;

        if (injectionText) {
            const tokens = estimateTokens(injectionText);
            console.log(`[ST-BME] 注入 ${tokens} 估算 tokens, Core=${result.stats.coreCount}, Recall=${result.stats.recallCount}`);

            // 使用 ST 的 extension prompt API 注入
            context.setExtensionPrompt(
                MODULE_NAME,
                injectionText,
                1,  // extension_prompt_types.IN_PROMPT
                settings.injectDepth,
            );
        }

        // 保存召回结果和访问强化
        currentGraph.lastRecallResult = result.selectedNodeIds;
        saveGraphToChat();
    } catch (e) {
        console.error('[ST-BME] 召回失败:', e);
    } finally {
        isRecalling = false;
    }
}

// ==================== 事件钩子 ====================

function onChatChanged() {
    loadGraphFromChat();
    lastInjectionContent = '';
}

async function onGenerationAfterCommands() {
    await runExtraction();
}

async function onBeforeCombinePrompts() {
    await runRecall();
}

function onMessageReceived() {
    // 新消息到达，图状态可能需要更新
    if (currentGraph) {
        saveGraphToChat();
    }
}

// ==================== UI 操作 ====================

async function onViewGraph() {
    if (!currentGraph) {
        toastr.warning('当前没有加载的图谱');
        return;
    }

    const stats = getGraphStats(currentGraph);
    const statsText = [
        `节点: ${stats.activeNodes} 活跃 / ${stats.archivedNodes} 归档`,
        `边: ${stats.totalEdges}`,
        `最后处理楼层: ${stats.lastProcessedSeq}`,
        `类型分布: ${Object.entries(stats.typeCounts).map(([k, v]) => `${k}=${v}`).join(', ') || '(空)'}`,
    ].join('\n');

    toastr.info(statsText, 'ST-BME 图谱状态', { timeOut: 10000 });
}

async function onRebuild() {
    if (!confirm('确定要从当前聊天重建图谱？这将清除现有图谱数据。')) return;

    currentGraph = createEmptyGraph();
    saveGraphToChat();

    toastr.info('图谱已重置，将在下次生成时重新提取');
}

async function onManualCompress() {
    if (!currentGraph) return;

    const result = await compressAll(currentGraph, getSchema(), getEmbeddingConfig(), false);
    saveGraphToChat();

    toastr.info(`压缩完成: 新建 ${result.created}, 归档 ${result.archived}`);
}

async function onExportGraph() {
    if (!currentGraph) return;

    const json = exportGraph(currentGraph);
    const blob = new Blob([json], { type: 'application/json' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `st-bme-graph-${Date.now()}.json`;
    a.click();
    URL.revokeObjectURL(url);

    toastr.success('图谱已导出');
}

async function onImportGraph() {
    const input = document.createElement('input');
    input.type = 'file';
    input.accept = '.json';
    input.onchange = async (e) => {
        const file = e.target.files[0];
        if (!file) return;

        try {
            const text = await file.text();
            currentGraph = importGraph(text);
            saveGraphToChat();
            toastr.success('图谱已导入');
        } catch (err) {
            toastr.error(`导入失败: ${err.message}`);
        }
    };
    input.click();
}

async function onViewLastInjection() {
    if (!lastInjectionContent) {
        toastr.info('暂无注入内容');
        return;
    }

    // 简单弹窗显示
    const popup = document.createElement('div');
    popup.style.cssText = 'position:fixed;top:50%;left:50%;transform:translate(-50%,-50%);background:#1a1a2e;color:#eee;padding:24px;border-radius:12px;max-width:80vw;max-height:80vh;overflow:auto;z-index:99999;white-space:pre-wrap;font-size:13px;box-shadow:0 8px 32px rgba(0,0,0,0.5);';
    popup.textContent = lastInjectionContent;

    const close = document.createElement('button');
    close.textContent = '关闭';
    close.style.cssText = 'position:absolute;top:8px;right:12px;background:#e94560;color:white;border:none;padding:4px 12px;border-radius:4px;cursor:pointer;';
    close.onclick = () => popup.remove();
    popup.appendChild(close);

    document.body.appendChild(popup);
}

async function onTestEmbedding() {
    const config = getEmbeddingConfig();
    if (!config.apiUrl || !config.apiKey) {
        toastr.warning('请先配置 Embedding API 地址和 Key');
        return;
    }

    toastr.info('正在测试 Embedding API 连通性...');
    const result = await testEmbeddingConnection(config);

    if (result.success) {
        toastr.success(`连接成功！向量维度: ${result.dimensions}`);
    } else {
        toastr.error(`连接失败: ${result.error}`);
    }
}

// ==================== 设置 UI ====================

function bindSettingsUI() {
    const settings = getSettings();

    // 开关
    $('#st_bme_enabled').prop('checked', settings.enabled).on('change', function () {
        settings.enabled = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // 提取频率
    $('#st_bme_extract_every').val(settings.extractEvery).on('input', function () {
        settings.extractEvery = Math.max(1, parseInt($(this).val()) || 1);
        saveSettingsDebounced();
    });

    // 召回开关
    $('#st_bme_recall_enabled').prop('checked', settings.recallEnabled).on('change', function () {
        settings.recallEnabled = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // LLM 精确召回
    $('#st_bme_recall_llm').prop('checked', settings.recallEnableLLM).on('change', function () {
        settings.recallEnableLLM = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // 注入深度
    $('#st_bme_inject_depth').val(settings.injectDepth).on('input', function () {
        settings.injectDepth = Math.max(0, parseInt($(this).val()) || 4);
        saveSettingsDebounced();
    });

    // 评分权重
    $('#st_bme_graph_weight').val(settings.graphWeight).on('input', function () {
        settings.graphWeight = parseFloat($(this).val()) || 0.6;
        saveSettingsDebounced();
    });
    $('#st_bme_vector_weight').val(settings.vectorWeight).on('input', function () {
        settings.vectorWeight = parseFloat($(this).val()) || 0.3;
        saveSettingsDebounced();
    });
    $('#st_bme_importance_weight').val(settings.importanceWeight).on('input', function () {
        settings.importanceWeight = parseFloat($(this).val()) || 0.1;
        saveSettingsDebounced();
    });

    // Embedding API
    $('#st_bme_embed_url').val(settings.embeddingApiUrl).on('input', function () {
        settings.embeddingApiUrl = $(this).val().trim();
        saveSettingsDebounced();
    });
    $('#st_bme_embed_key').val(settings.embeddingApiKey).on('input', function () {
        settings.embeddingApiKey = $(this).val().trim();
        saveSettingsDebounced();
    });
    $('#st_bme_embed_model').val(settings.embeddingModel).on('input', function () {
        settings.embeddingModel = $(this).val().trim();
        saveSettingsDebounced();
    });

    // 操作按钮
    $('#st_bme_btn_view_graph').on('click', onViewGraph);
    $('#st_bme_btn_rebuild').on('click', onRebuild);
    $('#st_bme_btn_compress').on('click', onManualCompress);
    $('#st_bme_btn_export').on('click', onExportGraph);
    $('#st_bme_btn_import').on('click', onImportGraph);
    $('#st_bme_btn_view_injection').on('click', onViewLastInjection);
    $('#st_bme_btn_test_embed').on('click', onTestEmbedding);

    // ====== v2 增强设置 UI 绑定 ======

    // P0: 记忆进化
    $('#st_bme_evolution').prop('checked', settings.enableEvolution).on('change', function () {
        settings.enableEvolution = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_evo_neighbors').val(settings.evoNeighborCount).on('input', function () {
        settings.evoNeighborCount = Math.max(1, parseInt($(this).val()) || 5);
        saveSettingsDebounced();
    });

    // P0: 精确对照
    $('#st_bme_precise_conflict').prop('checked', settings.enablePreciseConflict).on('change', function () {
        settings.enablePreciseConflict = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_conflict_threshold').val(settings.conflictThreshold).on('input', function () {
        settings.conflictThreshold = parseFloat($(this).val()) || 0.85;
        saveSettingsDebounced();
    });

    // P0: 全局概要
    $('#st_bme_synopsis').prop('checked', settings.enableSynopsis).on('change', function () {
        settings.enableSynopsis = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_synopsis_every').val(settings.synopsisEveryN).on('input', function () {
        settings.synopsisEveryN = Math.max(1, parseInt($(this).val()) || 5);
        saveSettingsDebounced();
    });

    // P1: 认知边界
    $('#st_bme_visibility').prop('checked', settings.enableVisibility ?? false).on('change', function () {
        settings.enableVisibility = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // P1: 交叉检索
    $('#st_bme_cross_recall').prop('checked', settings.enableCrossRecall ?? false).on('change', function () {
        settings.enableCrossRecall = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // P2: 惊奇度分割
    $('#st_bme_smart_trigger').prop('checked', settings.enableSmartTrigger).on('change', function () {
        settings.enableSmartTrigger = $(this).prop('checked');
        saveSettingsDebounced();
    });

    // P2: 主动遗忘
    $('#st_bme_sleep_cycle').prop('checked', settings.enableSleepCycle).on('change', function () {
        settings.enableSleepCycle = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_forget_threshold').val(settings.forgetThreshold).on('input', function () {
        settings.forgetThreshold = parseFloat($(this).val()) || 0.5;
        saveSettingsDebounced();
    });

    // P2: 概率触发
    $('#st_bme_prob_recall').prop('checked', settings.enableProbRecall).on('change', function () {
        settings.enableProbRecall = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_prob_chance').val(settings.probRecallChance).on('input', function () {
        settings.probRecallChance = parseFloat($(this).val()) || 0.15;
        saveSettingsDebounced();
    });

    // P2: 反思条目
    $('#st_bme_reflection').prop('checked', settings.enableReflection).on('change', function () {
        settings.enableReflection = $(this).prop('checked');
        saveSettingsDebounced();
    });
    $('#st_bme_reflect_every').val(settings.reflectEveryN).on('input', function () {
        settings.reflectEveryN = Math.max(3, parseInt($(this).val()) || 10);
        saveSettingsDebounced();
    });
}

// ==================== 初始化 ====================

(async function init() {
    // 加载设置面板 HTML
    const settingsHtml = await renderExtensionTemplateAsync('third-party/st-bme', 'settings');
    $('#extensions_settings2').append(settingsHtml);

    // 绑定 UI
    bindSettingsUI();

    // 注册事件钩子
    eventSource.on(event_types.CHAT_CHANGED, onChatChanged);
    eventSource.on(event_types.GENERATION_AFTER_COMMANDS, onGenerationAfterCommands);
    eventSource.on(event_types.GENERATE_BEFORE_COMBINE_PROMPTS, onBeforeCombinePrompts);
    eventSource.on(event_types.MESSAGE_RECEIVED, onMessageReceived);

    // 加载当前聊天的图谱
    loadGraphFromChat();

    console.log('[ST-BME] 初始化完成');
})();