connectai/src/lib/contextManager.ts

/**
 * ============================================================
 * Context Manager (컨텍스트 한계 관리)
 *
 * "context length = 132k" 는 "답변을 132k 토큰까지 생성해도 된다" 가 아닙니다.
 *   시스템 프롬프트 + 대화 기록 + 입력 문서 + 생성될 답변 + 여유분 ≤ context length
 *
 * 이 모듈은 요청을 보내기 *전에* 입력 토큰을 추정하고,
 *   - 동적으로 출력 상한(maxTokens)을 계산하고,
 *   - 대화 기록이 예산을 넘으면 오래된 메시지를 잘라내고,
 *   - 그래도 넘으면 시스템 프롬프트의 [CONTEXT] 블록을 마지막 수단으로 줄이고,
 *   - 생성 종료 사유(stopReason / finish_reason)를 "정상 / 출력한계 / 컨텍스트초과 / 사용자중단"
 *     으로 분류해 호출자가 잘린 응답을 감지할 수 있게 합니다.
 * ============================================================
 */

export type ChatRole = 'user' | 'assistant' | 'system';

export interface BudgetMessage {
    role: ChatRole;
    content: string;
    /** internal/system bookkeeping messages that should be kept verbatim where possible */
    internal?: boolean;
}

export interface ContextLimits {
    /** 모델의 context window (프롬프트 + 생성 합산 한계). */
    contextLength: number;
    /** 한 응답에서 생성할 토큰 수의 상한 (이 값을 넘기지 않음). */
    maxOutputTokens: number;
    /** 추정 오차를 흡수하기 위한 여유분. */
    safetyMargin: number;
    /** 출력에 항상 확보해 둘 최소 토큰 수. */
    minOutputTokens: number;
}

export const DEFAULT_CONTEXT_LIMITS: ContextLimits = {
    contextLength: 32768,
    maxOutputTokens: 4096,
    safetyMargin: 2048,
    minOutputTokens: 512,
};

/** LM Studio `contextOverflowPolicy` 값 — 우리가 예산 계산에 실패했을 때의 안전망. */
export type ContextOverflowPolicy = 'stopAtLimit' | 'truncateMiddle' | 'rollingWindow';
export const DEFAULT_OVERFLOW_POLICY: ContextOverflowPolicy = 'stopAtLimit';

/**
 * 텍스트의 토큰 수를 대략 추정합니다.
 *
 * 정밀한 토크나이저가 없으므로 문자 기반 휴리스틱을 사용합니다:
 *  - CJK(한/중/일) 글자: ~1.6 토큰/글자 (byte-level BPE 기준 보수적)
 *  - 그 외(영문/코드/기호): ~0.30 토큰/글자
 * 약간 과대평가하는 쪽으로 잡아 컨텍스트 초과를 예방합니다.
 */
export function estimateTokens(text: string): number {
    if (!text) return 0;
    const cjkChars = (text.match(/[　-〿぀-ヿ㐀-䶿一-鿿가-힯＀-￯]/g) || []).length;
    const otherChars = text.length - cjkChars;
    return Math.ceil(cjkChars * 1.6 + otherChars * 0.3);
}

/**
 * 모델 식별자에서 파라미터 규모(B 단위)를 대략 추출합니다. 모르면 null.
 * 예: "qwen2.5-7b" → 7, "llama-3.1-8b-instruct" → 8, "gemma-3n-e2b" / "gemma4:e2b" → 2,
 *     "phi-3-mini" → null (숫자 없음), "qwen3-30b-a3b" → 30. "4bit" 같은 양자화 표기는 매칭 안 됨.
 */
export function estimateModelParamsB(modelId: string | null | undefined): number | null {
    if (!modelId) return null;
    const m = String(modelId).match(/(?:^|[-_/:.\s])e?(\d+(?:\.\d+)?)\s*b(?![a-z0-9])/i);
    if (!m) return null;
    const n = Number(m[1]);
    return Number.isFinite(n) && n > 0 && n < 2000 ? n : null;
}

/** role/구분자 등 메시지 1개당 발생하는 고정 오버헤드(대략). */
const PER_MESSAGE_TOKEN_OVERHEAD = 4;

export function estimateMessageTokens(msg: BudgetMessage): number {
    return estimateTokens(msg.content || '') + PER_MESSAGE_TOKEN_OVERHEAD;
}

export function estimateMessagesTokens(messages: BudgetMessage[]): number {
    return messages.reduce((sum, m) => sum + estimateMessageTokens(m), 0);
}

/**
 * 입력 토큰 수가 주어졌을 때 안전하게 생성할 수 있는 출력 토큰 상한을 계산합니다.
 *
 *   available = contextLength - inputTokens - safetyMargin
 *   maxOutput = clamp(available, minOutputTokens, maxOutputTokens)
 *
 * available 이 minOutputTokens 보다 작으면 입력이 이미 컨텍스트를 거의 다 먹은 상태이므로
 * `tight: true` 와 함께 minOutputTokens 를 그대로 돌려줍니다 (호출자가 추가로 줄여야 함).
 */
export function computeOutputBudget(
    inputTokens: number,
    limits: ContextLimits = DEFAULT_CONTEXT_LIMITS
): { maxOutputTokens: number; available: number; tight: boolean } {
    const { contextLength, maxOutputTokens, safetyMargin, minOutputTokens } = limits;
    const available = contextLength - inputTokens - safetyMargin;
    if (available <= minOutputTokens) {
        return { maxOutputTokens: minOutputTokens, available, tight: true };
    }
    return {
        maxOutputTokens: Math.max(minOutputTokens, Math.min(available, maxOutputTokens)),
        available,
        tight: false,
    };
}

export interface TrimResult<M extends BudgetMessage> {
    messages: M[];
    /** 잘려나간 메시지 개수 (0 이면 변화 없음). */
    droppedCount: number;
    /** 잘라낸 뒤의 입력 토큰 추정치. */
    tokensAfter: number;
}

/**
 * 대화 기록을 토큰 예산 안에 맞춥니다 (sliding window).
 *
 * 전략:
 *  1. 항상 마지막 메시지(보통 현재 사용자 질문)는 유지.
 *  2. 최근 메시지부터 역순으로 예산이 허용하는 만큼 채움.
 *  3. 하나라도 잘렸으면 맨 앞에 marker 를 끼워 모델이 맥락 누락을 인지하게 함.
 *     v2.2.69+ — marker 콜백은 droppedCount 뿐 아니라 *잘려나간 메시지 배열* 도 받아
 *     단순 count 가 아닌 진짜 요약/맥락을 작성할 수 있다.
 *
 * 주의: 여기서 잘라내는 것은 *요청에 보낼* 메시지 배열일 뿐, UI에 표시되는 전체 기록은 그대로 둡니다.
 */
export function trimHistoryToBudget<M extends BudgetMessage>(
    messages: M[],
    budgetTokens: number,
    makeMarker: (droppedCount: number, droppedMessages: M[]) => M
): TrimResult<M> {
    if (messages.length === 0) {
        return { messages, droppedCount: 0, tokensAfter: 0 };
    }
    const total = estimateMessagesTokens(messages);
    if (total <= budgetTokens) {
        return { messages, droppedCount: 0, tokensAfter: total };
    }

    // 최근 메시지부터 역순으로 채움. 최소 1개(마지막 메시지)는 무조건 유지.
    const kept: M[] = [];
    let used = 0;
    for (let i = messages.length - 1; i >= 0; i--) {
        const t = estimateMessageTokens(messages[i]);
        if (kept.length > 0 && used + t > budgetTokens) {
            break;
        }
        kept.unshift(messages[i]);
        used += t;
    }

    const droppedCount = messages.length - kept.length;
    if (droppedCount > 0) {
        const droppedMessages = messages.slice(0, droppedCount);
        const marker = makeMarker(droppedCount, droppedMessages);
        kept.unshift(marker);
        used += estimateMessageTokens(marker);
    }
    return { messages: kept, droppedCount, tokensAfter: used };
}

/** 시스템 프롬프트 안에서 "잘라내도 되는" 보조 컨텍스트 영역의 시작/끝 마커. */
export const CONTEXT_OPEN_MARKER = '\n\n[CONTEXT]\n';
export const CONTEXT_CLOSE_MARKER = '\n[/CONTEXT]\n';

/**
 * 시스템 프롬프트가 너무 클 때 마지막 수단으로 `[CONTEXT] … [/CONTEXT]` 사이의 보조 컨텍스트
 * (브레인/메모리/열린 파일/RAG 등 — 조립 단계에서 끼워 넣는 데이터)만 잘라냅니다.
 * 핵심 지시문(앞부분)과 마무리 지시문(예: negative constraints, agent system prompt — 뒷부분)은
 * 절대 건드리지 않습니다. `[/CONTEXT]` 마커가 없으면 `[CONTEXT]` 이후 전체를 trim 대상으로 봅니다.
 *
 * @param systemPrompt 조립이 끝난 전체 시스템 프롬프트
 * @param maxTokens    시스템 프롬프트에 허용할 토큰 상한
 */
export function truncateSystemPromptContext(
    systemPrompt: string,
    maxTokens: number
): { prompt: string; truncated: boolean } {
    if (estimateTokens(systemPrompt) <= maxTokens) {
        return { prompt: systemPrompt, truncated: false };
    }
    const openIdx = systemPrompt.indexOf(CONTEXT_OPEN_MARKER);
    if (openIdx < 0) {
        // 보조 컨텍스트 영역이 없으면 전체에서 뒤를 잘라낼 수밖에 없음.
        const approxChars = Math.max(1000, Math.floor(maxTokens / 0.3));
        return {
            prompt: systemPrompt.slice(0, approxChars) + '\n\n[…시스템 프롬프트가 컨텍스트 한계로 잘렸습니다…]',
            truncated: true,
        };
    }
    const bodyStart = openIdx + CONTEXT_OPEN_MARKER.length;
    const closeIdx = systemPrompt.indexOf(CONTEXT_CLOSE_MARKER, bodyStart);
    const head = systemPrompt.slice(0, bodyStart);              // 지시문 + "[CONTEXT]\n"
    const body = closeIdx >= 0 ? systemPrompt.slice(bodyStart, closeIdx) : systemPrompt.slice(bodyStart);
    const tail = closeIdx >= 0 ? systemPrompt.slice(closeIdx) : '';   // "[/CONTEXT]" + negative/agent 등

    const fixedTokens = estimateTokens(head) + estimateTokens(tail);
    const remainForBody = maxTokens - fixedTokens - 64;
    if (remainForBody <= 0) {
        return {
            prompt: head + '[…보조 컨텍스트는 컨텍스트 한계로 모두 생략되었습니다…]' + tail,
            truncated: true,
        };
    }
    // CJK 비중에 따라 글자수→토큰 비율이 달라지므로 보수적으로 0.4 토큰/글자로 환산.
    const approxChars = Math.floor(remainForBody / 0.4);
    const trimmedBody = body.length <= approxChars
        ? body
        : body.slice(0, approxChars) + '\n\n[…이하 보조 컨텍스트는 컨텍스트 한계로 생략됨…]';
    return { prompt: head + trimmedBody + tail, truncated: true };
}

export type GenerationStopKind =
    | 'complete'        // 정상 종료 (EOS / stop string)
    | 'output-limit'    // maxTokens 도달 — 답변이 중간에 잘림
    | 'context-overflow'// 입력+출력이 context window 초과
    | 'user-stopped'    // 사용자 취소
    | 'tool-calls'      // 툴 호출로 종료
    | 'error'           // 모델/런타임 오류
    | 'unknown';

/**
 * 엔진별 종료 사유 문자열을 공통 분류값으로 정규화합니다.
 *  - LM Studio SDK: `stats.stopReason` — eosFound / stopStringFound / maxPredictedTokensReached / contextLengthReached / userStopped / toolCalls / failed / modelUnloaded
 *  - OpenAI 호환 REST: `choices[].finish_reason` — stop / length / tool_calls / content_filter
 *  - Ollama: `done_reason` — stop / length / load
 */
export function classifyStopReason(raw: string | null | undefined): GenerationStopKind {
    if (!raw) return 'unknown';
    const r = String(raw).toLowerCase();
    if (/(maxpredictedtokensreached|^length$|max_tokens)/.test(r)) return 'output-limit';
    if (/(contextlengthreached|context_length|context_overflow|contextoverflow)/.test(r)) return 'context-overflow';
    if (/(eosfound|stopstringfound|^stop$|^end$|stop_sequence|content_filter)/.test(r)) return 'complete';
    if (/(userstopped|aborted|cancel)/.test(r)) return 'user-stopped';
    if (/(toolcalls|tool_calls)/.test(r)) return 'tool-calls';
    if (/(failed|error|modelunloaded)/.test(r)) return 'error';
    return 'unknown';
}

/**
 * 잘린 응답일 때 사용자에게 덧붙일 한 줄 안내. 정상 종료면 빈 문자열.
 * (output-limit 은 Astra 가 먼저 자동 이어쓰기를 시도하므로, 이 안내는 그래도 다 못 채웠을 때만 보입니다.
 *  그래서 "이어서 작성해줘" 같은 사용자 액션을 요구하지 않습니다.)
 */
export function truncationNotice(kind: GenerationStopKind): string {
    switch (kind) {
        case 'output-limit':
            return '\n\n> ⚠️ 답변이 길어 자동으로 이어 정리했지만 여전히 길이 한계에 닿았습니다. 더 좁은 주제로 나눠 질문하시면 완전한 답변을 받을 수 있어요.';
        case 'context-overflow':
            return '\n\n> ⚠️ 입력 컨텍스트가 모델의 context window 를 초과했습니다. 대화를 새로 시작하거나(`/newChat`) Settings 에서 `g1nation.contextLength` 를 모델 실제 값으로 맞추고, Brain/Skill 컨텍스트를 줄여보세요.';
        case 'error':
            return '\n\n> ⚠️ 모델이 비정상 종료했습니다 (컨텍스트 초과 또는 모델 용량 부족 가능). 더 큰 모델로 바꾸거나 컨텍스트를 줄여보세요.';
        default:
            return '';
    }
}

/**
 * Some local engines report `maxPredictedTokensReached` even when the visible
 * answer is short (for example after an internal retry or SDK stats mismatch).
 * Only show the "answer was cut off" notice when the generated answer actually
 * consumed most of the output budget.
 */
export function shouldShowTruncationNotice(
    kind: GenerationStopKind,
    outputTokens: number,
    maxOutputTokens: number
): boolean {
    if (kind === 'context-overflow' || kind === 'error') return true;
    if (kind !== 'output-limit') return false;
    const threshold = Math.max(128, Math.floor(maxOutputTokens * 0.85));
    return outputTokens >= threshold;
}