// Story Summary - Prompt Injection
// 注入格式：世界状态 → 亲身经历(含闪回) → 相关背景(含闪回) → 记忆碎片 → 人物弧光

import { getContext } from "../../../../../../extensions.js";
import { extension_prompts, extension_prompt_types, extension_prompt_roles } from "../../../../../../../script.js";
import { xbLog } from "../../../core/debug-core.js";
import { getSummaryStore } from "../data/store.js";
import { getVectorConfig, getSummaryPanelConfig, getSettings } from "../data/config.js";
import { recallMemory, buildQueryText } from "../vector/recall.js";

const MODULE_ID = "summaryPrompt";
const SUMMARY_PROMPT_KEY = "LittleWhiteBox_StorySummary";

const BUDGET = { total: 10000, l3Max: 2000, l2Max: 5000 };
const MAX_CHUNKS_PER_EVENT = 2;
const MAX_ORPHAN_CHUNKS = 6;

// ═══════════════════════════════════════════════════════════════════════════
// 工具函数
// ═══════════════════════════════════════════════════════════════════════════

function estimateTokens(text) {
    if (!text) return 0;
    const s = String(text);
    const zh = (s.match(/[\u4e00-\u9fff]/g) || []).length;
    return Math.ceil(zh + (s.length - zh) / 4);
}

function pushWithBudget(lines, text, state) {
    const t = estimateTokens(text);
    if (state.used + t > state.max) return false;
    lines.push(text);
    state.used += t;
    return true;
}

// 从 summary 解析楼层范围：(#321-322) 或 (#321)
function parseFloorRange(summary) {
    if (!summary) return null;
    
    // 匹配 (#123-456) 或 (#123)
    const match = String(summary).match(/\(#(\d+)(?:-(\d+))?\)/);
    if (!match) return null;
    
    const start = parseInt(match[1], 10);
    const end = match[2] ? parseInt(match[2], 10) : start;
    
    return { start, end };
}

// 去掉 summary 末尾的楼层标记
function cleanSummary(summary) {
    return String(summary || '').replace(/\s*\(#\d+(?:-\d+)?\)\s*$/, '').trim();
}

// ═══════════════════════════════════════════════════════════════════════════
// L1 → L2 归属
// ═══════════════════════════════════════════════════════════════════════════

function attachChunksToEvents(events, chunks) {
    const usedChunkIds = new Set();
    
    // 给每个 event 挂载 chunks
    for (const e of events) {
        e._chunks = [];
        const range = parseFloorRange(e.event?.summary);
        if (!range) continue;
        
        for (const c of chunks) {
            if (c.floor >= range.start && c.floor <= range.end) {
                if (!usedChunkIds.has(c.chunkId)) {
                    e._chunks.push(c);
                    usedChunkIds.add(c.chunkId);
                }
            }
        }
        
        // 每个事件最多保留 N 条，按相似度排序
        e._chunks.sort((a, b) => (b.similarity || 0) - (a.similarity || 0));
        e._chunks = e._chunks.slice(0, MAX_CHUNKS_PER_EVENT);
    }
    
    // 找出无归属的 chunks（记忆碎片）
    const orphans = chunks
        .filter(c => !usedChunkIds.has(c.chunkId))
        .sort((a, b) => (b.similarity || 0) - (a.similarity || 0))
        .slice(0, MAX_ORPHAN_CHUNKS);
    
    return { events, orphans };
}

// ═══════════════════════════════════════════════════════════════════════════
// 格式化函数
// ═══════════════════════════════════════════════════════════════════════════

function formatWorldLines(world) {
    return [...(world || [])]
        .sort((a, b) => (b.floor || 0) - (a.floor || 0))
        .map(w => `- ${w.topic}：${w.content}`);
}

function formatChunkLine(c) {
    const text = String(c.text || '');
    const preview = text.length > 80 ? text.slice(0, 80) + '...' : text;
    return `› #${c.floor} ${preview}`;
}

function formatEventBlock(e, idx) {
    const ev = e.event || {};
    const time = ev.timeLabel || '';
    const people = (ev.participants || []).join(' / ');
    const summary = cleanSummary(ev.summary);
    
    const lines = [];
    
    // 标题行
    const header = time ? `${idx}.【${time}】${people}` : `${idx}. ${people}`;
    lines.push(header);
    
    // 摘要
    lines.push(`  ${summary}`);
    
    // 挂载的闪回
    for (const c of (e._chunks || [])) {
        lines.push(`  ${formatChunkLine(c)}`);
    }
    
    return lines.join('\n');
}

function formatArcLine(a) {
    const moments = (a.moments || [])
        .map(m => typeof m === 'string' ? m : m.text)
        .filter(Boolean);
    
    if (moments.length) {
        return `- ${a.name}：${moments.join(' → ')}（当前：${a.trajectory}）`;
    }
    return `- ${a.name}：${a.trajectory}`;
}

// ═══════════════════════════════════════════════════════════════════════════
// 主构建函数
// ═══════════════════════════════════════════════════════════════════════════

function buildMemoryPromptVectorEnabled(store, recallResult) {
    const data = store.json || {};
    const total = { used: 0, max: BUDGET.total };
    const sections = [];
    
    // ─────────────────────────────────────────────────────────────────────
    // [世界状态]
    // ─────────────────────────────────────────────────────────────────────
    const worldLines = formatWorldLines(data.world);
    if (worldLines.length) {
        const l3 = { used: 0, max: Math.min(BUDGET.l3Max, total.max) };
        const l3Lines = [];
        
        for (const line of worldLines) {
            if (!pushWithBudget(l3Lines, line, l3)) break;
        }
        
        if (l3Lines.length) {
            sections.push(`[世界状态] 请严格遵守\n${l3Lines.join('\n')}`);
            total.used += l3.used;
        }
    }
    
    // ─────────────────────────────────────────────────────────────────────
    // L1 → L2 归属处理
    // ─────────────────────────────────────────────────────────────────────
    const events = recallResult?.events || [];
    const chunks = recallResult?.chunks || [];
    const { events: eventsWithChunks, orphans } = attachChunksToEvents(events, chunks);
    
    // 分离 DIRECT 和 SIMILAR
    const directEvents = eventsWithChunks.filter(e => e._recallType === 'DIRECT');
    const similarEvents = eventsWithChunks.filter(e => e._recallType !== 'DIRECT');
    
    // ─────────────────────────────────────────────────────────────────────
    // [亲身经历] - DIRECT
    // ─────────────────────────────────────────────────────────────────────
    if (directEvents.length) {
        const l2 = { used: 0, max: Math.min(BUDGET.l2Max * 0.7, total.max - total.used) };
        const lines = [];
        
        let idx = 1;
        for (const e of directEvents) {
            const block = formatEventBlock(e, idx);
            if (!pushWithBudget(lines, block, l2)) break;
            idx++;
        }
        
        if (lines.length) {
            sections.push(`[亲身经历]\n\n${lines.join('\n\n')}`);
            total.used += l2.used;
        }
    }
    
    // ─────────────────────────────────────────────────────────────────────
    // [相关背景] - SIMILAR
    // ─────────────────────────────────────────────────────────────────────
    if (similarEvents.length) {
        const l2s = { used: 0, max: Math.min(BUDGET.l2Max * 0.3, total.max - total.used) };
        const lines = [];
        
        let idx = directEvents.length + 1;
        for (const e of similarEvents) {
            const block = formatEventBlock(e, idx);
            if (!pushWithBudget(lines, block, l2s)) break;
            idx++;
        }
        
        if (lines.length) {
            sections.push(`[相关背景]\n\n${lines.join('\n\n')}`);
            total.used += l2s.used;
        }
    }
    
    // ─────────────────────────────────────────────────────────────────────
    // [记忆碎片] - 无归属的 chunks
    // ─────────────────────────────────────────────────────────────────────
    if (orphans.length && total.used < total.max) {
        const l1 = { used: 0, max: total.max - total.used };
        const lines = [];
        
        for (const c of orphans) {
            const line = formatChunkLine(c);
            if (!pushWithBudget(lines, line, l1)) break;
        }
        
        if (lines.length) {
            sections.push(`[记忆碎片]\n${lines.join('\n')}`);
            total.used += l1.used;
        }
    }
    
    // ─────────────────────────────────────────────────────────────────────
    // [人物弧光]
    // ─────────────────────────────────────────────────────────────────────
    if (data.arcs?.length && total.used < total.max) {
        const arcLines = data.arcs.map(formatArcLine);
        const arcText = `[人物弧光]\n${arcLines.join('\n')}`;
        
        if (total.used + estimateTokens(arcText) <= total.max) {
            sections.push(arcText);
        }
    }
    
    // ─────────────────────────────────────────────────────────────────────
    // 组装
    // ─────────────────────────────────────────────────────────────────────
    if (!sections.length) return '';
    
    return `<剧情记忆>\n\n${sections.join('\n\n')}\n\n</剧情记忆>`;
}

function buildMemoryPromptVectorDisabled(store) {
    const data = store.json || {};
    const sections = [];
    
    // 世界状态
    if (data.world?.length) {
        const lines = formatWorldLines(data.world);
        sections.push(`[世界状态] 请严格遵守\n${lines.join('\n')}`);
    }
    
    // 全部事件（无召回，按时间）
    if (data.events?.length) {
        const lines = data.events.map((ev, i) => {
            const time = ev.timeLabel || '';
            const people = (ev.participants || []).join(' / ');
            const summary = cleanSummary(ev.summary);
            const header = time ? `${i + 1}.【${time}】${people}` : `${i + 1}. ${people}`;
            return `${header}\n  ${summary}`;
        });
        sections.push(`[剧情记忆]\n\n${lines.join('\n\n')}`);
    }
    
    // 弧光
    if (data.arcs?.length) {
        const lines = data.arcs.map(formatArcLine);
        sections.push(`[人物弧光]\n${lines.join('\n')}`);
    }
    
    if (!sections.length) return '';
    
    return `<剧情记忆>\n\n${sections.join('\n\n')}\n\n</剧情记忆>`;
}

// ═══════════════════════════════════════════════════════════════════════════
// 导出
// ═══════════════════════════════════════════════════════════════════════════

export function formatPromptWithMemory(store, recallResult) {
    const vectorCfg = getVectorConfig();
    return vectorCfg?.enabled
        ? buildMemoryPromptVectorEnabled(store, recallResult)
        : buildMemoryPromptVectorDisabled(store);
}

export async function recallAndInjectPrompt(excludeLastAi = false, postToFrame = null) {
    if (!getSettings().storySummary?.enabled) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    const { chat } = getContext();
    const store = getSummaryStore();

    if (!store?.json) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    const allEvents = store.json.events || [];
    const lastIdx = store.lastSummarizedMesId ?? 0;
    const length = chat?.length || 0;

    if (lastIdx >= length) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    const vectorCfg = getVectorConfig();
    let recallResult = { events: [], chunks: [] };

    if (vectorCfg?.enabled) {
        try {
            const queryText = buildQueryText(chat, 2, excludeLastAi);
            recallResult = await recallMemory(queryText, allEvents, vectorCfg, { excludeLastAi });
            postToFrame?.({ type: "RECALL_LOG", text: recallResult.logText || "" });
        } catch (e) {
            xbLog.error(MODULE_ID, "召回失败", e);
        }
    }

    injectPrompt(store, recallResult, chat);
}

export function updateSummaryExtensionPrompt() {
    if (!getSettings().storySummary?.enabled) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    const { chat } = getContext();
    const store = getSummaryStore();

    if (!store?.json || (store.lastSummarizedMesId ?? 0) >= (chat?.length || 0)) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    injectPrompt(store, { events: [], chunks: [] }, chat);
}

function injectPrompt(store, recallResult, chat) {
    const length = chat?.length || 0;

    let text = formatPromptWithMemory(store, recallResult);

    const cfg = getSummaryPanelConfig();
    if (cfg.trigger?.wrapperHead) {
        text = cfg.trigger.wrapperHead + "\n" + text;
    }
    if (cfg.trigger?.wrapperTail) {
        text = text + "\n" + cfg.trigger.wrapperTail;
    }

    if (!text.trim()) {
        delete extension_prompts[SUMMARY_PROMPT_KEY];
        return;
    }

    const lastIdx = store.lastSummarizedMesId ?? 0;
    let depth = length - lastIdx - 1;
    if (depth < 0) depth = 0;

    if (cfg.trigger?.forceInsertAtEnd) {
        depth = 10000;
    }

    extension_prompts[SUMMARY_PROMPT_KEY] = {
        value: text,
        position: extension_prompt_types.IN_CHAT,
        depth,
        role: extension_prompt_roles.ASSISTANT,
    };
}

export function clearSummaryExtensionPrompt() {
    delete extension_prompts[SUMMARY_PROMPT_KEY];
}