import { AgentVariables, PriorityContainer } from '../../index.js';
import { Rag } from '../../rag/Rag.js';
import { MessageRecombiner } from '../../chunking/MessageRecombiner.js';
export class ContextualizedChat {
    _rawChat;
    _chunker;
    _variables;
    _chunks = {
        persistent: [],
        temporary: [],
    };
    _rags;
    _functionCallResultFirstChunks = {};
    _lastTwoMsgs = {
        persistent: new PriorityContainer(2, (a, b) => b - a),
        temporary: new PriorityContainer(2, (a, b) => b - a),
    };
    constructor(_context, _rawChat, _chunker, _variables) {
        this._rawChat = _rawChat;
        this._chunker = _chunker;
        this._variables = _variables;
        this._rags = {
            persistent: Rag.standard(_context).selector(x => x.json),
            temporary: Rag.standard(_context).selector(x => x.json),
        };
    }
    get rawChat() {
        return this._rawChat;
    }
    async contextualize(contextVector, tokenLimits) {
        await this._processNewMessages();
        const persistentLargeChunks = await this._rags.persistent
            .query(contextVector)
            .recombine(MessageRecombiner.standard(tokenLimits.persistent, this._rawChat.chatLogs, 'persistent', this._lastTwoMsgs.persistent.getItems()));
        const temporaryChunks = await this._rags.temporary
            .query(contextVector)
            .recombine(MessageRecombiner.standard(tokenLimits.temporary, this._rawChat.chatLogs, 'temporary', this._lastTwoMsgs.temporary.getItems()));
        const sorted = {
            persistent: persistentLargeChunks.map(x => JSON.parse(x.json)),
            temporary: temporaryChunks.map(x => JSON.parse(x.json)),
        };
        sorted.persistent = postProcessMessages(sorted.persistent);
        sorted.temporary = postProcessMessages(sorted.temporary);
        const chat = this._rawChat.cloneEmpty();
        await chat.persistent(sorted.persistent);
        await chat.temporary(sorted.temporary);
        return chat;
    }
    async _processNewMessages() {
        this._processNewMessagesByType('persistent');
        this._processNewMessagesByType('temporary');
    }
    _processNewMessagesByType(type) {
        const messages = this._rawChat.chatLogs.get(type).msgs;
        if (messages.length === 0) {
            return;
        }
        const lastProcessedIdx = getLastProcessedMessageIndex(this._chunks[type]);
        if (lastProcessedIdx === messages.length - 1) {
            return;
        }
        for (let i = lastProcessedIdx + 1; i < messages.length; ++i) {
            const message = messages[i];
            this._processNewMessage(message, i, type);
        }
    }
    _processNewMessage(message, msgIdx, type) {
        const newChunks = [];
        const varName = message.content || '';
        let isVariable = false;
        if (AgentVariables.hasSyntax(varName)) {
            const data = this._variables.get(varName);
            if (data) {
                message.content = data;
                isVariable = true;
            }
        }
        if (this._chunker.shouldChunk(message)) {
            if (isVariable) {
                newChunks.push(...this._chunker.chunk(message).map((chunk, index) => variableChunkText(chunk, index, varName)));
            }
            else {
                newChunks.push(...this._chunker.chunk(message));
            }
        }
        else if (isVariable) {
            newChunks.push(variableChunkText(message, 0, varName));
        }
        else {
            newChunks.push(JSON.stringify(message));
        }
        const chunks = this._chunks[type];
        const startChunkIdx = chunks.length;
        for (const _ of newChunks) {
            chunks.push({ msgIdx });
        }
        this._rags[type].addItems(newChunks.map((chunk, index) => ({
            chunkIdx: startChunkIdx + index,
            msgIdx,
            tokens: this._rawChat.tokenizer.encode(chunk).length,
            json: chunk,
        })));
        if (message.role === 'function' || 'function_call' in message) {
            this._functionCallResultFirstChunks[startChunkIdx] = {
                text: newChunks[0],
                metadata: { index: startChunkIdx },
            };
        }
        this._lastTwoMsgs[type].addItem(startChunkIdx);
    }
}
function getLastProcessedMessageIndex(chunks) {
    const lastIdx = chunks.length - 1;
    if (lastIdx < 0) {
        return -1;
    }
    return chunks[lastIdx].msgIdx;
}
function postProcessMessages(messages) {
    const result = [];
    const varChunkPrefix = 'Variable "${';
    const varNameRegex = /\$\{([^}]+)\}/;
    let prevVarName;
    for (const message of messages) {
        let varName;
        if (message.content?.startsWith(varChunkPrefix)) {
            const match = message.content.match(varNameRegex);
            varName = (match && match[1]) || undefined;
        }
        if (varName && prevVarName === varName) {
            const lastMessage = result.at(-1);
            lastMessage.content += `\n\n${message.content}`;
            result[result.length - 1] = lastMessage;
        }
        else {
            result.push(message);
        }
        if (prevVarName && !varName) {
            const lastMessage = result.at(-1);
            lastMessage.content += `\nThe above function result was too large, so it was stored in the variable "\${${prevVarName}}".`;
        }
        prevVarName = varName;
    }
    return result;
}
function variableChunkText(chunk, index, varName) {
    const message = typeof chunk === 'string' ? JSON.parse(chunk) : chunk;
    return JSON.stringify({
        ...message,
        content: `Variable "${varName}" chunk #${index}\n\`\`\`\n${message.content}\n\`\`\``,
    });
}
