"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.expandTokensDefinition = exports.TokenCompressor = void 0;
const JsonCompressor_1 = require("./JsonCompressor");
class TokenCompressor {
    storage = {};
    keyPrefix = "compressed_";
    toolName = exports.expandTokensDefinition.function.name;
    compressionThreshold = 4000;
    characterLimit = this.compressionThreshold * 4;
    maxTokens = this.compressionThreshold * 2;
    jsonCompressor;
    constructor(toolsService) {
        this.jsonCompressor = new JsonCompressor_1.JsonCompressor(this, this.compressionThreshold, this.maxTokens, this.toolName);
        this.registerTool(toolsService);
    }
    estimateTokens(text) {
        return Math.ceil(text.length / 4);
    }
    setCompressionThreshold(threshold) {
        this.compressionThreshold = threshold;
        this.characterLimit = threshold * 4;
        this.jsonCompressor.updateSettings(threshold, this.maxTokens);
    }
    setMaxTokens(maxTokens) {
        if (maxTokens > this.maxTokens) {
            this.maxTokens = maxTokens;
            this.jsonCompressor.updateSettings(this.compressionThreshold, maxTokens);
        }
    }
    tryParseJson(content) {
        return this.jsonCompressor.tryParseJson(content);
    }
    compressStringInChunks(content, path = "") {
        if (content.length <= this.characterLimit) {
            return content;
        }
        const chunks = [];
        const chunkKeys = [];
        let remaining = content;
        while (remaining.length > this.characterLimit) {
            const chunkStart = remaining.length - this.characterLimit;
            const chunk = remaining.substring(chunkStart);
            chunks.unshift(chunk);
            remaining = remaining.substring(0, chunkStart);
        }
        if (remaining.length > 0) {
            chunks.unshift(remaining);
        }
        for (let i = 0; i < chunks.length; i++) {
            const key = this.generateKey();
            chunkKeys.push(key);
        }
        for (let i = 0; i < chunks.length; i++) {
            let chunkContent = chunks[i];
            if (i < chunks.length - 1) {
                const nextKey = chunkKeys[i + 1];
                chunkContent += `\n\n[NEXT_CHUNK_KEY: ${nextKey}]`;
            }
            this.storeString(chunkKeys[i], chunkContent);
        }
        const firstKey = chunkKeys[0];
        const totalTokens = this.estimateTokens(content);
        const chunkCount = chunks.length;
        return `[COMPRESSED_STRING - ${totalTokens} tokens in ${chunkCount} chunks]\nKey: ${firstKey}\nPath: ${path}\nPreview: ${content.substring(0, 200)}...\n[Use ${this.toolName} tool with key "${firstKey}" to retrieve content. Follow NEXT_CHUNK_KEY references for complete content]\n[TIP: try jqToolResponse,grepToolResponse,tailToolResponse to filter/search/map without repeated expandTokens calls]`;
    }
    isAlreadyCompressed(content) {
        if (content.includes("[COMPRESSED_STRING")) {
            return true;
        }
        const parsed = this.tryParseJson(content);
        if (parsed && parsed._schema_key && typeof parsed._schema_key === "string") {
            return true;
        }
        return false;
    }
    compressContent(content, path = "") {
        if (this.isAlreadyCompressed(content)) {
            return content;
        }
        const tokens = this.estimateTokens(content);
        const thresholdToUse = path === "" ? this.compressionThreshold : this.maxTokens;
        if (tokens <= thresholdToUse) {
            return content;
        }
        const jsonObj = this.tryParseJson(content);
        if (jsonObj) {
            const dataToCompress = jsonObj._mcp_format ? jsonObj._data : jsonObj;
            const schema = this.jsonCompressor.generateSchema(jsonObj);
            const schemaKey = this.generateKey();
            this.storeString(schemaKey, JSON.stringify(schema));
            const compressedObj = this.compressJsonProperties(dataToCompress, path || "data");
            const finalCompressedObj = jsonObj._mcp_format
                ? {
                    _mcp_format: true,
                    _raw_structure: jsonObj._raw_structure,
                    _data: compressedObj,
                }
                : compressedObj;
            const resultWithSchema = typeof finalCompressedObj === "object" &&
                !Array.isArray(finalCompressedObj)
                ? { ...finalCompressedObj, _schema_key: schemaKey }
                : { _schema_key: schemaKey, data: finalCompressedObj };
            const compressedContent = JSON.stringify(resultWithSchema, null, 2);
            const compressedTokens = this.estimateTokens(compressedContent);
            const compressionThreshold = 0.6;
            if (compressedTokens < tokens * compressionThreshold) {
                return compressedContent;
            }
        }
        return this.compressStringInChunks(content, path);
    }
    compressJsonProperties(obj, path = "") {
        return this.jsonCompressor.compressJsonProperties(obj, path);
    }
    generateKey() {
        return `${this.keyPrefix}${Date.now()}_${Math.random()
            .toString(36)
            .substr(2, 9)}`;
    }
    async compressMessage(message) {
        if (typeof message.content === "string") {
            const compressed = this.compressContent(message.content);
            if (message.role === "tool" && message.tool_call_id && compressed !== message.content) {
                message.content = compressed + `\n[toolCallId: ${message.tool_call_id}]`;
            }
            else {
                message.content = compressed;
            }
        }
        else if (Array.isArray(message.content)) {
            for (const item of message.content) {
                if (item.type === "text" && item.text) {
                    item.text = this.compressContent(item.text);
                }
            }
        }
    }
    createProcessor(filterFn) {
        return async (originalMessages, modifiedMessages) => {
            for (const message of modifiedMessages) {
                if (filterFn && !filterFn(message)) {
                    continue;
                }
                await this.compressMessage(message);
            }
        };
    }
    retrieveString(key) {
        return this.storage[key] || null;
    }
    storeString(key, value) {
        if (this.estimateTokens(value) > this.maxTokens) {
            this.setMaxTokens(this.estimateTokens(value) + 1);
        }
        this.storage[key] = value;
    }
    clearStorage() {
        this.storage = {};
        this.jsonCompressor.clearDeduplication();
    }
    getStorageKeys() {
        return Object.keys(this.storage);
    }
    getStorageSize() {
        return Object.keys(this.storage).length;
    }
    registerTool(toolsService) {
        if (toolsService) {
            toolsService.addTools([exports.expandTokensDefinition]);
            toolsService.addFunctions({
                [this.toolName]: (key) => {
                    const data = this.retrieveString(key);
                    if (!data) {
                        return `Error: No data found for key "${key}". Available keys: ${this.getStorageKeys().join(", ")}`;
                    }
                    return data;
                },
            });
        }
    }
    getSchema(key) {
        const schemaKey = `${key}_schema`;
        const schemaStr = this.storage[schemaKey];
        if (!schemaStr) {
            return null;
        }
        try {
            return JSON.parse(schemaStr);
        }
        catch (e) {
            return null;
        }
    }
    getCompressedProperties(key) {
        const content = this.storage[key];
        if (!content) {
            return null;
        }
        try {
            const metadata = JSON.parse(content);
            return metadata.compressed_properties || null;
        }
        catch (e) {
            return null;
        }
    }
    getFullObject(mainObj, compressedKey) {
        if (!mainObj || typeof mainObj !== "object") {
            return mainObj;
        }
        const compressed = this.getCompressedProperties(compressedKey);
        if (!compressed) {
            return mainObj;
        }
        const { _compressed_properties_key, _compressed_property_names, _compression_info, ...highSignal } = mainObj;
        return { ...highSignal, ...compressed };
    }
    extractKeys(content) {
        const keys = [];
        const keyPattern = /\$expandTokens\[([^\]]+)\]|Key:\s*([^\s\n]+)/g;
        let match;
        while ((match = keyPattern.exec(content)) !== null) {
            const key = match[1] || match[2];
            if (key && !keys.includes(key)) {
                keys.push(key);
            }
        }
        return keys;
    }
    getKeyChain(key) {
        const chain = [];
        let currentKey = key;
        while (currentKey) {
            const content = this.storage[currentKey];
            if (!content)
                break;
            chain.push({
                key: currentKey,
                size: content.length,
                preview: content.substring(0, 100),
            });
            const nextMatch = content.match(/NEXT_CHUNK_KEY:\s*([^\s\n]+)/);
            currentKey = nextMatch ? nextMatch[1] : null;
        }
        return chain;
    }
}
exports.TokenCompressor = TokenCompressor;
exports.expandTokensDefinition = {
    type: "function",
    function: {
        name: "expandTokens",
        description: "Retrieve a chunk of compressed data that was stored during message processing. The returned content may contain a `NEXT_CHUNK_KEY` to retrieve subsequent chunks. NOTE: Can also use jqToolResponse, grepToolResponse, tailToolResponse to access/filter/search compressed content without needing to expand all tokens.",
        parameters: {
            type: "object",
            positional: true,
            properties: {
                key: {
                    type: "string",
                    description: "The key of the compressed data to retrieve",
                },
            },
            required: ["key"],
        },
    },
};
//# sourceMappingURL=TokenCompressor.js.map