connectai/src/lib/contextBuilders/memoryContext.ts

import * as path from 'path';
import * as vscode from 'vscode';
import type { ChatMessage } from '../../agent';
import type { BrainProfile } from '../../config';
import { getConfig } from '../../config';
import type { MemoryManager } from '../../memory';
import type { RetrievalOrchestrator } from '../../retrieval';
import { buildLessonChecklistBlock } from '../../retrieval/lessonHelpers';
import { embedQuery, embedTexts } from '../../retrieval/embeddings';
import { backfillBrainEmbeddings } from '../../retrieval/brainIndex';
import { resolveScopeForAgent } from '../../skills/agentKnowledgeMap';
import {
    resolveKnowledgeMix,
    mapWeightToBrainFileLimit,
    mapWeightToRetrievalRatio,
    ResolvedKnowledgeMix,
} from '../../retrieval/knowledgeMix';
import { buildConflictWarningsBlock, ConflictThresholdSetting } from '../../retrieval/conflictBlock';
import { buildCoveChecklistBlock } from '../../retrieval/coveBlock';
import { captureWorkStateSignals } from '../../retrieval/actionabilityScoring';
import { getRecentSlashCommands } from '../../features/datacollect/slashRouter';
import { semanticRerank, DEFAULT_SEMANTIC_RERANK_OPTIONS } from '../../retrieval/semanticRerank';
import { detectAmbiguity, buildIntentClarificationBlock, IntentStrictness } from '../../retrieval/intentClarification';
import { buildCitationTraceBlock } from '../../retrieval/citationTrace';
import { buildTerminologyBlock } from '../../retrieval/terminologyBlock';

/**
 * 한 turn 의 RAG / 5-layer memory 컨텍스트 빌드.
 *
 * 옛 코드: agent.ts 의 130줄짜리 private `buildMemoryContext`. 인스턴스 state 6개
 * (memoryManager, chatHistory, retrievalOrchestrator, context, currentTaskId,
 * _turnCtx) 에 의존해 god-file 의 일부였음.
 *
 * 분리 방식: 호출자(provider) 가 모은 deps struct 를 받는 *순수 orchestration*
 * 함수로 격리. RetrievalOrchestrator / MemoryManager 자체는 그대로 둠 (이 함수는
 * 그 두 객체의 *사용 패턴* 만 표준화). 사이드 이펙트 두 가지는 명시:
 *   1) `deps.turnCtx` mutation — webview footer 가 읽는 retrieval/lessons/knowledgeMix.
 *   2) `backfillBrainEmbeddings` fire-and-forget — 다음 turn 의 score 향상용.
 *
 * 의도: agent.ts 로부터 130줄 빼면서 RAG 호출 패턴을 단위 테스트 대상 함수로 노출.
 * Provider 는 deps 만 채워 호출하면 되도록 줄임.
 */

/** TurnContext 의 retrieval 슬롯 모양. provider 의 `_turnCtx.retrieval` 와 일치해야 함. */
export interface TurnRetrievalSummary {
    agentName: string | null;
    scoped: boolean;
    source: string;
    configuredFolders: string[];
    usedBrainFiles: string[];
    usedMemoryLayers: string[];
    lessonFiles: string[];
    totalChunks: number;
    selectedChunks: number;
}

/**
 * Mutable turn-context sink — 호출자의 `_turnCtx` 와 같은 객체를 그대로 받아 함수
 * 안에서 채워준다. 매 호출 직전에 호출자가 `reset` 해서 비워야 함.
 */
export interface TurnContextSink {
    retrieval: TurnRetrievalSummary | null;
    lessons: string[];
    knowledgeMix: ResolvedKnowledgeMix | null;
    /**
     * 동적 시스템 프롬프트 블록 — id → 본문. memoryContext 가 이 turn 에 채우고,
     * buildAstraModeSystemPrompt 가 iterate 해서 [CONTEXT] 밖에 join 주입.
     * Casual conversation 모드면 자동 skip. 빈 본문은 자동 제외.
     *
     * 등록 순서대로 prompt 에 join — `intent-clarification` → `terminology` →
     * `conflict-warnings` → `cove-checklist` → `citation-trace`.
     */
    dynamicBlocks: Map<string, string>;
    /** Post-hoc Self-Check 용 — selected chunks 의 (title, excerpt) 요약. */
    selfCheckSources: Array<{ title: string; excerpt: string }>;
}

export interface MemoryContextDeps {
    currentPrompt: string;
    activeBrain: BrainProfile;
    agentSkillFile?: string;
    /** Visible + internal 합친 raw chat history. 함수 안에서 internal 필터링. */
    chatHistory: ChatMessage[];
    memoryManager: MemoryManager;
    retrievalOrchestrator: RetrievalOrchestrator;
    /** vscode ExtensionContext — chat_sessions globalState 읽기에 사용. */
    context: vscode.ExtensionContext;
    /** 현재 turn 의 session id — recentSessions 에서 자기 자신 제외. */
    currentTaskId: string;
    /** 함수가 채울 turn-context sink. 호출자는 호출 전에 비워둬야 한다. */
    turnCtx: TurnContextSink;
}

/**
 * 영구 저장된 chat_sessions 풀에서 medium-term 후보를 추리는 compact helper.
 * 활성 세션 자신은 제외, 빈 history 도 제외, 짧은 미리보기/요약만 보관해
 * orchestrator 입력에 들어가도 토큰 폭증 안 함.
 */
function compactRecentSessions(
    rawSessions: any[],
    activeSessionId: string | null,
    limit: number,
): Array<{ id: string; title: string; firstUserMsg: string; lastAssistantExcerpt: string; summary?: string; timestamp: number }> {
    if (!Array.isArray(rawSessions) || rawSessions.length === 0 || limit <= 0) return [];
    const pool = rawSessions.length > limit + 5 ? limit + 5 : rawSessions.length;
    const out: Array<{ id: string; title: string; firstUserMsg: string; lastAssistantExcerpt: string; summary?: string; timestamp: number }> = [];
    for (let i = 0; i < rawSessions.length && out.length < pool; i++) {
        const s = rawSessions[i];
        if (!s || typeof s !== 'object') continue;
        const id = String(s.id ?? '');
        if (!id || id === activeSessionId) continue;
        const history: any[] = Array.isArray(s.history) ? s.history : [];
        if (history.length === 0) continue;
        const firstUser = history.find((m) => m?.role === 'user');
        const lastAssistant = [...history].reverse().find((m) => m?.role === 'assistant');
        const firstUserMsg = String(firstUser?.content ?? '').replace(/\s+/g, ' ').trim().slice(0, 200);
        const lastTxt = String(lastAssistant?.content ?? '').replace(/\s+/g, ' ').trim();
        const lastAssistantExcerpt = lastTxt.length <= 200 ? lastTxt : lastTxt.slice(-200);
        const summary = typeof s.summary === 'string' ? s.summary.trim().slice(0, 600) : undefined;
        if (!firstUserMsg && !lastAssistantExcerpt && !summary) continue;
        out.push({
            id,
            title: String(s.title ?? '').trim() || firstUserMsg.slice(0, 50),
            firstUserMsg,
            lastAssistantExcerpt,
            summary,
            timestamp: typeof s.timestamp === 'number' ? s.timestamp : 0,
        });
    }
    return out;
}

export async function buildMemoryContext(deps: MemoryContextDeps): Promise<string> {
    const config = getConfig();
    if (!config.memoryEnabled) return '';

    // Settings 가 turn 사이에 바뀔 수 있으니 매번 동기화.
    deps.memoryManager.updateConfig({
        enabled: config.memoryEnabled,
        shortTermLimit: config.memoryShortTermMessages,
    });

    const visibleHistory = deps.chatHistory.filter((message) => !message.internal);
    const workspaceFolders = vscode.workspace.workspaceFolders;
    const workspacePath = workspaceFolders ? workspaceFolders[0].uri.fsPath : undefined;

    // Agent ↔ knowledge map. 매핑 없으면 folders=[] → orchestrator 가 whole-brain 사용 (legacy).
    const scope = resolveScopeForAgent(deps.agentSkillFile, deps.activeBrain.localBrainPath);

    // Context 윈도우 비례 retrieval 예산. 32K → 8K, 230K → 57K, 80K cap (scoring 속도).
    const scaledTotalBudget = Math.min(
        80000,
        Math.max(8000, Math.floor(config.contextLength * 0.25)),
    );

    // medium-term layer 용 옛 세션 후보. sidebar 가 직접 쓰는 key 를 read-through.
    const rawSessions = deps.context.globalState.get<any[]>('chat_sessions', []) || [];
    const recentSessions = compactRecentSessions(
        rawSessions,
        deps.currentTaskId,
        Math.max(0, config.memoryMediumTermSessions ?? 0),
    );

    // Hybrid retrieval (옵션): embedding model 있으면 query embedding 가져와 cosine
    // + TF-IDF blend. timeout 4초 — endpoint 가 느리면 그냥 pure TF-IDF 로 진행.
    let queryEmbedding: number[] | undefined;
    if (config.embeddingModel) {
        const EMBED_QUERY_TIMEOUT_MS = 4000;
        try {
            queryEmbedding = await Promise.race([
                embedQuery(deps.currentPrompt, { baseUrl: config.ollamaUrl, model: config.embeddingModel }),
                new Promise<undefined>((resolve) => setTimeout(() => resolve(undefined), EMBED_QUERY_TIMEOUT_MS)),
            ]);
        } catch {
            queryEmbedding = undefined;
        }
    }

    // Knowledge Mix 가중치 (per-agent → global → default). weight=50 이면 legacy 기본값과 동일.
    const knowledgeMix = resolveKnowledgeMix(deps.agentSkillFile);
    deps.turnCtx.knowledgeMix = knowledgeMix;
    const mixedBrainFileLimit = mapWeightToBrainFileLimit(knowledgeMix.weight, config.memoryLongTermFiles);
    const mixedRetrievalRatio = mapWeightToRetrievalRatio(knowledgeMix.weight);

    // Actionability — work-state 신호 캡처 (최근 슬래시 명령 + 열린 파일).
    // 설정으로 disable 가능. 신호 없으면 retrieve() 가 legacy 동작.
    const workStateSignals = config.actionabilityEnabled !== false
        ? captureWorkStateSignals(getRecentSlashCommands())
        : undefined;

    // Unified RAG Pipeline 호출.
    const result = deps.retrievalOrchestrator.retrieve(deps.currentPrompt, {
        brain: deps.activeBrain,
        memoryManager: deps.memoryManager,
        workspacePath,
        chatHistory: visibleHistory,
        contextBudget: {
            totalBudget: scaledTotalBudget,
            retrievalRatio: mixedRetrievalRatio,
        },
        brainFileLimit: mixedBrainFileLimit,
        scopeFolders: scope.folders,
        recentSessions,
        mediumTermLimit: config.memoryMediumTermSessions ?? 0,
        queryEmbedding,
        embeddingModel: config.embeddingModel || undefined,
        embeddingBlendAlpha: config.embeddingBlendAlpha,
        workStateSignals,
        hierarchicalReweightEnabled: config.hierarchicalReweightEnabled !== false,
    });

    // Semantic Re-rank (LLM, async) — selectedChunks 의 *순서* 만 재배치. 토큰 예산을
    // 통과한 chunks 안에서 의도-부합도 순으로 재정렬해 LLM attention bias 활용.
    // 기본 OFF — latency 우려. 사용자가 명시 enable 시만.
    if (config.semanticRerankEnabled && result.selectedChunks.length >= 3) {
        const rerankModel = (config.semanticRerankModel || '').trim() || config.defaultModel;
        if (rerankModel && config.ollamaUrl) {
            const rerankRes = await semanticRerank(deps.currentPrompt, result.selectedChunks, {
                ollamaUrl: config.ollamaUrl,
                model: rerankModel,
                candidateK: config.semanticRerankCandidateK ?? DEFAULT_SEMANTIC_RERANK_OPTIONS.candidateK,
                timeoutMs: (config.semanticRerankTimeoutSec ?? 8) * 1000,
                excerptLength: DEFAULT_SEMANTIC_RERANK_OPTIONS.excerptLength,
            });
            // In-place 교체 — buildContextString 가 이 배열을 그대로 읽음.
            result.selectedChunks = rerankRes.rerankedChunks;
            result.fusionLog.push(`Semantic re-rank: ${rerankRes.success ? '✓' : '✗'} ${rerankRes.note} (${rerankRes.durationMs}ms)`);
        }
    }

    // Fire-and-forget background embedding. Vector 없는 파일만 embed 하므로
    // steady-state turn 은 작업량 0 — 다음 turn 이 혜택.
    if (config.embeddingModel) {
        const scoredFilePaths = result.selectedChunks
            .filter((c) => c.source === 'brain-memory' && c.metadata.filePath)
            .map((c) => c.metadata.filePath!)
            .filter((p, i, arr) => arr.indexOf(p) === i);
        if (scoredFilePaths.length > 0) {
            void backfillBrainEmbeddings(
                deps.activeBrain.localBrainPath,
                scoredFilePaths,
                config.embeddingModel,
                (texts) => embedTexts(texts, { baseUrl: config.ollamaUrl, model: config.embeddingModel }),
            );
        }
    }

    // webview "scope used" footer 가 읽는 turn-context summary. brain-trace 는
    // 검색이 아니라 trace 표시용이라 usedMemoryLayers 에서 제외 (brain-memory 도 제외 —
    // 별도 usedBrainFiles 로 표시).
    const brainRoot = deps.activeBrain.localBrainPath;
    const rel = (p?: string) => (p ? (path.relative(brainRoot, p) || p) : '');
    const lessonChunks = result.lessonChunks || [];
    deps.turnCtx.retrieval = {
        agentName: scope.agent?.name ?? null,
        scoped: scope.folders.length > 0,
        source: String((scope as any).source ?? ''),
        configuredFolders: scope.folders.map((abs) => rel(abs)),
        usedBrainFiles: result.selectedChunks
            .filter((c) => c.source === 'brain-memory' && c.metadata.filePath)
            .map((c) => rel(c.metadata.filePath))
            .filter((p, i, arr) => p && arr.indexOf(p) === i),
        usedMemoryLayers: Array.from(new Set(
            result.selectedChunks
                .filter((c) => c.source !== 'brain-memory' && c.source !== 'brain-trace')
                .map((c) => c.source as string),
        )),
        lessonFiles: lessonChunks.map((c) => rel(c.metadata.filePath)).filter((p, i, arr) => p && arr.indexOf(p) === i),
        totalChunks: result.totalChunks,
        selectedChunks: result.selectedChunks.length,
    };

    deps.turnCtx.lessons = lessonChunks.map((c) => c.content);

    // 동적 시스템 프롬프트 블록 빌드 — 등록 순서대로 turnCtx.dynamicBlocks 에 set.
    // 옛 named field 5개 (conflictWarnings/coveChecklist/intentClarification/citationTrace/
    // terminology) 통합. 새 블록 추가 = 여기서 setBlock 한 줄.
    const blocks = deps.turnCtx.dynamicBlocks;

    // Intent Clarification — 답변보다 *역질문 우선*. 모호 아닐 때 빈 문자열 → join 시 자동 제외.
    if (config.intentClarificationEnabled !== false) {
        const strict = (config.intentClarificationStrictness || 'medium') as IntentStrictness;
        const ambig = detectAmbiguity(deps.currentPrompt, strict);
        blocks.set('intent-clarification', buildIntentClarificationBlock(ambig));
    }

    // Terminology Dictionary — 사용자 편집 글로서리. 파일 없으면 빈 문자열.
    if (config.glossaryEnabled !== false) {
        blocks.set('terminology', buildTerminologyBlock({
            relPath: config.glossaryPath || '.astra/glossary.md',
            maxBodyLength: config.glossaryMaxBodyLength ?? 4000,
        }));
    }

    // Conflict Surface — selectedChunks 의 per-doc conflictSeverity + 교차-문서 발산.
    if (config.conflictHighlightingEnabled !== false) {
        const threshold: ConflictThresholdSetting = (config.conflictSeverityThreshold || 'medium') as ConflictThresholdSetting;
        blocks.set('conflict-warnings', buildConflictWarningsBlock(result.selectedChunks, {
            selfFlagThreshold: threshold,
            crossDivergenceEnabled: config.conflictCrossDocEnabled !== false,
        }));
    }

    // CoVe — 답변 *작성 전* 그라운딩 체크리스트.
    if (config.coveEnabled !== false) {
        blocks.set('cove-checklist', buildCoveChecklistBlock(result.selectedChunks, deps.currentPrompt, {
            topSourcesCount: config.coveTopSourcesCount ?? 5,
            strictMode: config.coveStrictMode === true,
        }));
    }

    // Citation Trace — 답변 끝 출처 한 줄.
    if (config.citationTraceEnabled !== false && result.selectedChunks.length > 0) {
        blocks.set('citation-trace', buildCitationTraceBlock(result.selectedChunks));
    }

    // Self-Check 용 source 미리보기 — agent.ts 가 post-stream 에서 사용.
    deps.turnCtx.selfCheckSources = result.selectedChunks.slice(0, 5).map((c) => ({
        title: c.title || '(제목 없음)',
        excerpt: (c.content || '').slice(0, 200),
    }));

    // Lessons 블록은 일반 RAG context 보다 앞 — context-overflow truncation 에서 먼저
    // 살아남게.
    const lessonBlock = buildLessonChecklistBlock(lessonChunks.map((c) => ({ title: c.title, content: c.content })));
    const memoryBlock = deps.retrievalOrchestrator.buildContextString(result);
    return [lessonBlock, memoryBlock].filter(Boolean).join('\n\n');
}