"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.TokenCounter = void 0;
const sdk_1 = __importDefault(require("@anthropic-ai/sdk"));
const tiktoken_1 = require("tiktoken");
const cliconfig_1 = require("../database/cliconfig");
const path_1 = __importDefault(require("path"));
const fs_1 = __importDefault(require("fs"));
const chalk_1 = __importDefault(require("chalk"));
// Function to count OpenAI tokens
function countOpenAITokens(messages, model) {
    const encoder = (0, tiktoken_1.encoding_for_model)(model);
    let tokenCount = 0;
    for (const msg of messages) {
        tokenCount += encoder.encode(msg.role || "").length;
        if (msg.role === "assistant") {
            if (typeof msg.content === "string") {
                tokenCount += encoder.encode(msg.content).length;
            }
            else if (Array.isArray(msg.content)) {
                for (let block of msg.content) {
                    if (block.type === "text") {
                        tokenCount += encoder.encode(block.text).length;
                    }
                    else {
                        throw new Error("Cannot encode non text tokens yet");
                    }
                }
            }
            if (msg.tool_calls) {
                for (let toolCall of msg.tool_calls) {
                    tokenCount += encoder.encode(toolCall.id).length;
                    tokenCount += encoder.encode(toolCall.type).length;
                    tokenCount += encoder.encode(toolCall.function.name).length;
                    tokenCount += encoder.encode(toolCall.function.arguments).length;
                }
            }
        }
        else if (msg.role === "tool") {
            tokenCount += encoder.encode(msg.tool_call_id).length;
            if (typeof msg.content === "string") {
                tokenCount += encoder.encode(msg.content).length;
            }
            else {
                for (let block of msg.content) {
                    if (block.type === "text") {
                        tokenCount += encoder.encode(block.text).length;
                    }
                    else {
                        throw new Error("Cannot encode non text tokens yet");
                    }
                }
            }
        }
        else if (msg.role === "user") {
            if (typeof msg.content === "string") {
                tokenCount += encoder.encode(msg.content).length;
            }
            else {
                for (let block of msg.content) {
                    if (block.type === "text") {
                        tokenCount += encoder.encode(block.text).length;
                    }
                    else {
                        throw new Error("Cannot encode non text tokens yet");
                    }
                }
            }
        }
    }
    return tokenCount;
}
// TokenCounter class
class TokenCounter {
    /**
     * Creates a new TokenCounter instance
     * @param tools The tools to use for token counting
     */
    constructor(tools = [], agentType = "", model = "", provider = "openai") {
        this.tools = tools;
        this.inMemoryInputMessages = [];
        this.agentType = agentType;
        this.model = model;
        this.provider = provider;
        this.inMemoryOutputMessages = [];
        this.amazonUsage = [];
    }
    /**
     * Sets the tools for this token counter
     * @param tools The tools to use
     */
    setTools(tools) {
        this.tools = tools;
    }
    appendInMemoryInputMessages(messages) {
        if (this.provider !== "amazon") {
            this.inMemoryInputMessages.push(messages);
        }
    }
    appendInMemoryOutputMessages(messages) {
        if (this.provider !== "amazon") {
            this.inMemoryOutputMessages.push(...messages);
        }
    }
    getInputTokenCount() {
        return __awaiter(this, void 0, void 0, function* () {
            if (this.provider === "amazon") {
                if (this.amazonUsage.length > 0) {
                    let total = this.amazonUsage.reduce((totalInput, usage) => totalInput + (usage.cacheReadInputTokens || 0) + (usage.inputTokens || 0), 0);
                    return {
                        tokenCount: total,
                        model: this.model,
                        provider: this.provider,
                        agentType: this.agentType,
                    };
                }
                else {
                    throw new Error("Uninitialized usage in amazon");
                }
            }
            else {
                // Get the system prompt based on the agent type
                let systemPrompt = "You are a helpful assistant.";
                if (!this.agentType) {
                    throw new Error(`Agent type not specified`);
                }
                if (this.agentType === "widget") {
                    systemPrompt = fs_1.default.readFileSync(path_1.default.resolve(__dirname, "../systemPrompts/latest.html"), { encoding: "utf8" });
                }
                else if (this.agentType === "app") {
                    systemPrompt = fs_1.default.readFileSync(path_1.default.resolve(__dirname, "../systemPrompts/appPlanner.html"), { encoding: "utf8" });
                }
                else {
                    throw new Error(`Unknown agent type: ${this.agentType}`);
                }
                let tokenCount;
                if (this.provider === "claude") {
                    try {
                        const rec = (0, cliconfig_1.getAPIKey)("claude");
                        if (!rec) {
                            throw new Error("Claude API key not found");
                        }
                        const claude = new sdk_1.default({ apiKey: rec.apikey });
                        const tokenCounts = [0];
                        yield Promise.all(this.inMemoryInputMessages.map((messages) => __awaiter(this, void 0, void 0, function* () {
                            try {
                                const tokenCountResult = yield claude.messages.countTokens({
                                    model: this.model,
                                    system: systemPrompt,
                                    messages: (0, cliconfig_1.toAnthropicMessages)((0, cliconfig_1.sequenceToolCallsAndAddEmptyResults)(messages)),
                                    tools: this.tools,
                                });
                                tokenCounts.push(tokenCountResult.input_tokens);
                            }
                            catch (error) {
                                console.error("Failed to count tokens for messages: ", error, JSON.stringify(messages, null, 2));
                            }
                        })));
                        tokenCount = tokenCounts.reduce((acc, count) => acc + count, 0);
                    }
                    catch (err) {
                        console.error("Failed to count input tokens", err);
                        throw err;
                    }
                }
                else {
                    // For OpenAI, use our own implementation of token counting
                    const tokenCounts = [0];
                    yield Promise.all(this.inMemoryInputMessages.map((messages) => __awaiter(this, void 0, void 0, function* () {
                        const openaiMsgs = (0, cliconfig_1.toOpenaiMessages)(messages, systemPrompt);
                        const count = countOpenAITokens(openaiMsgs, this.model);
                        tokenCounts.push(count);
                    })));
                    tokenCount = tokenCounts.reduce((acc, count) => acc + count, 0);
                    // Add an estimate for tool definitions (very rough estimate)
                    const toolsTokenEstimate = this.tools.reduce((acc, tool) => {
                        // Rough estimate: name + description + JSON.stringify(inputSchema)
                        return (acc +
                            tool.name.length +
                            (tool.description || "").length +
                            JSON.stringify(tool.input_schema).length);
                    }, 0);
                    // Convert characters to tokens (rough estimate: 4 chars = 1 token)
                    tokenCount += Math.ceil(toolsTokenEstimate / 4);
                }
                // Return token count, cost, and tools information
                return {
                    tokenCount,
                    model: this.model,
                    provider: this.provider,
                    agentType: this.agentType,
                };
            }
        });
    }
    /**
     * Gets the token count for output messages
     * @returns The token count result for output messages
     */
    getOutputTokenCount() {
        return __awaiter(this, void 0, void 0, function* () {
            if (this.provider === "amazon") {
                if (this.amazonUsage) {
                    let total = this.amazonUsage.reduce((totalOutput, usage) => totalOutput + (usage.outputTokens || 0), 0);
                    return {
                        tokenCount: total,
                        model: this.model,
                        provider: this.provider,
                        agentType: this.agentType,
                    };
                }
                else {
                    throw new Error("Uninitialized usage in amazon");
                }
            }
            else {
                // For output messages, we don't need system prompt or tools
                let tokenCount;
                if (this.provider === "claude") {
                    try {
                        const rec = (0, cliconfig_1.getAPIKey)("claude");
                        if (!rec) {
                            throw new Error("Claude API key not found");
                        }
                        const claude = new sdk_1.default({ apiKey: rec.apikey });
                        const tokenCountResult = yield claude.messages.countTokens({
                            model: this.model,
                            messages: (0, cliconfig_1.toAnthropicMessages)((0, cliconfig_1.sequenceToolCallAndResponses)(this.inMemoryOutputMessages)),
                        });
                        tokenCount = tokenCountResult.input_tokens;
                    }
                    catch (err) {
                        console.error("Failed to count output tokens", err, chalk_1.default.redBright(JSON.stringify(this.inMemoryOutputMessages.map((it) => {
                            return {
                                role: it.role,
                                tool_call_id: it.tool_call_id,
                                content_type: it.content_type,
                            };
                        }), null, 2)));
                        throw err;
                    }
                }
                else {
                    // For OpenAI, use our own implementation of token counting
                    const openaiMsgs = (0, cliconfig_1.toOpenaiMessages)(this.inMemoryOutputMessages, "");
                    tokenCount = countOpenAITokens(openaiMsgs, this.model);
                }
                // Return token count, cost, and tools information
                return {
                    tokenCount,
                    model: this.model,
                    provider: this.provider,
                    agentType: this.agentType,
                };
            }
        });
    }
    setAmazonUsage(usage) {
        this.amazonUsage.push(usage);
    }
    /**
     * Gets the token count for both input and output messages
     * @returns The combined token count result
     */
    getTokenCount() {
        return __awaiter(this, void 0, void 0, function* () {
            console.info(chalk_1.default.red("Token usage: " + JSON.stringify(this.amazonUsage, null, 2)));
            if (this.provider === "amazon" && this.amazonUsage) {
                const totalInput = this.amazonUsage.reduce((totalInput, usage) => totalInput + (usage.cacheReadInputTokens || 0) + (usage.inputTokens || 0), 0);
                const totalOutput = this.amazonUsage.reduce((totalOutput, usage) => totalOutput + (usage.outputTokens || 0), 0);
                return {
                    inputTokenCount: totalInput,
                    outputTokenCount: totalOutput,
                    model: this.model,
                    provider: this.provider,
                    agentType: this.agentType
                };
            }
            else {
                try {
                    const inputResult = yield this.getInputTokenCount();
                    const outputResult = yield this.getOutputTokenCount();
                    // Combine the results
                    return {
                        inputTokenCount: inputResult.tokenCount,
                        outputTokenCount: outputResult.tokenCount,
                        model: this.model,
                        provider: this.provider,
                        agentType: this.agentType,
                    };
                }
                catch (err) {
                    console.error("Failed to get token counts: ", err);
                    return {
                        inputTokenCount: 0,
                        outputTokenCount: 0,
                        model: "",
                        provider: "claude",
                        agentType: "",
                    };
                }
            }
        });
    }
}
exports.TokenCounter = TokenCounter;
