const calculateTotalTokens = (chunks) => {
    return chunks.reduce((acc, chunk) => acc + chunk.tokens, 0);
};
const tryGetFunctionMessagePair = (chunk, chunkIdx, originalItems, chatLogs, chatLogType) => {
    const msg = JSON.parse(chunk.json);
    if ((msg.role === 'assistant' && 'function_call' in msg) || msg.role === 'function') {
        const direction = msg.role === 'assistant' ? 1 : -1;
        const search = msg.role === 'assistant' ? 'function' : 'assistant';
        let funcChunkIdx;
        let nextChunkIdx = chunkIdx + direction;
        while (!funcChunkIdx) {
            if (originalItems.length >= nextChunkIdx || nextChunkIdx < 0) {
                break;
            }
            const msg = chatLogs.getMsg(chatLogType, originalItems[nextChunkIdx].msgIdx);
            if (msg?.role === search) {
                funcChunkIdx = nextChunkIdx;
            }
            else {
                nextChunkIdx += direction;
            }
        }
        if (funcChunkIdx) {
            return funcChunkIdx;
        }
    }
    return undefined;
};
export class MessageRecombiner {
    static standard(tokenLimit, chatLogs, chatLogType, initChunks) {
        return async (results, originalItems) => {
            if (calculateTotalTokens(originalItems) < tokenLimit) {
                return originalItems;
            }
            const chunks = [];
            const chunksAdded = new Set();
            const msgsAdded = {};
            let tokenCount = 0;
            const canAddTokens = (tokens) => tokenCount + tokens <= tokenLimit;
            const addChunk = (chunkIdx, recursive = false) => {
                const chunk = originalItems[chunkIdx];
                if (chunksAdded.has(chunkIdx)) {
                    return true;
                }
                chunksAdded.add(chunkIdx);
                if (!canAddTokens(chunk.tokens)) {
                    return false;
                }
                if (!msgsAdded[chunk.msgIdx]) {
                    msgsAdded[chunk.msgIdx] = 0;
                }
                if (msgsAdded[chunk.msgIdx] >= 4) {
                    return true;
                }
                msgsAdded[chunk.msgIdx] += 1;
                chunks.push(chunk);
                tokenCount += chunk.tokens;
                const firstChunkIdx = originalItems.slice(0, chunkIdx).findIndex(item => item.msgIdx === chunk.msgIdx);
                if (firstChunkIdx !== -1) {
                    addChunk(firstChunkIdx);
                }
                if (!recursive) {
                    const funcPairChunk = tryGetFunctionMessagePair(chunk, chunkIdx, originalItems, chatLogs, chatLogType);
                    if (funcPairChunk) {
                        return addChunk(funcPairChunk, true);
                    }
                }
                return true;
            };
            const iterator = await results();
            if (chatLogType === 'persistent') {
                const smallPersistentChunks = [];
                let lastMsgIndex = -1;
                for (const chunk of originalItems) {
                    if (chunk.msgIdx !== lastMsgIndex) {
                        smallPersistentChunks.push(chunk);
                        lastMsgIndex = chunk.msgIdx;
                    }
                    else {
                        smallPersistentChunks.pop();
                        break;
                    }
                }
                for (const chunk of smallPersistentChunks) {
                    addChunk(chunk.chunkIdx);
                }
            }
            for (const chunk of initChunks) {
                addChunk(chunk);
            }
            for await (const result of iterator) {
                const chunkIdx = result.metadata().index;
                addChunk(chunkIdx);
            }
            return chunks.sort((a, b) => a.chunkIdx - b.chunkIdx);
        };
    }
}
