"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.TokenUsageTracker = void 0;
const logger_1 = __importDefault(require("../logger"));
/**
 * A utility class for tracking token usage across an evaluation
 */
class TokenUsageTracker {
    constructor() {
        this.providersMap = new Map();
    }
    /**
     * Get the singleton instance of TokenUsageTracker
     */
    static getInstance() {
        if (!TokenUsageTracker.instance) {
            TokenUsageTracker.instance = new TokenUsageTracker();
        }
        return TokenUsageTracker.instance;
    }
    /**
     * Track token usage for a provider
     * @param provider The provider to track usage for
     * @param usage The token usage to track
     */
    trackUsage(providerId, usage) {
        if (!usage) {
            return;
        }
        const current = this.providersMap.get(providerId) ??
            {
                total: 0,
                prompt: 0,
                completion: 0,
                cached: 0,
                numRequests: 0,
                completionDetails: { reasoning: 0, acceptedPrediction: 0, rejectedPrediction: 0 },
                assertions: { total: 0, prompt: 0, completion: 0, cached: 0 },
            };
        this.providersMap.set(providerId, this.mergeUsage(current, usage));
        logger_1.default.debug(`Tracked token usage for ${providerId}: total=${usage.total ?? 0}, cached=${usage.cached ?? 0}`);
    }
    /**
     * Get the cumulative token usage for a specific provider
     * @param providerId The ID of the provider to get usage for
     * @returns The token usage for the provider
     */
    getProviderUsage(providerId) {
        return this.providersMap.get(providerId);
    }
    /**
     * Get all provider IDs that have token usage tracked
     * @returns Array of provider IDs
     */
    getProviderIds() {
        return Array.from(this.providersMap.keys());
    }
    /**
     * Get aggregated token usage across all providers
     * @returns Aggregated token usage
     */
    getTotalUsage() {
        const result = {
            total: 0,
            prompt: 0,
            completion: 0,
            cached: 0,
            numRequests: 0,
            completionDetails: {
                reasoning: 0,
                acceptedPrediction: 0,
                rejectedPrediction: 0,
            },
            assertions: {
                total: 0,
                prompt: 0,
                completion: 0,
                cached: 0,
            },
        };
        // Accumulate totals from all providers
        for (const usage of this.providersMap.values()) {
            result.total += usage.total ?? 0;
            result.prompt += usage.prompt ?? 0;
            result.completion += usage.completion ?? 0;
            result.cached += usage.cached ?? 0;
            result.numRequests += usage.numRequests ?? 0;
            // Add completion details
            if (usage.completionDetails) {
                result.completionDetails.reasoning += usage.completionDetails.reasoning ?? 0;
                result.completionDetails.acceptedPrediction +=
                    usage.completionDetails.acceptedPrediction ?? 0;
                result.completionDetails.rejectedPrediction +=
                    usage.completionDetails.rejectedPrediction ?? 0;
            }
            // Add assertion statistics
            if (usage.assertions) {
                result.assertions.total += usage.assertions.total ?? 0;
                result.assertions.prompt += usage.assertions.prompt ?? 0;
                result.assertions.completion += usage.assertions.completion ?? 0;
                result.assertions.cached += usage.assertions.cached ?? 0;
            }
        }
        return result;
    }
    /**
     * Reset token usage for a specific provider
     * @param providerId The ID of the provider to reset
     */
    resetProviderUsage(providerId) {
        this.providersMap.delete(providerId);
    }
    /**
     * Reset token usage for all providers
     */
    resetAllUsage() {
        this.providersMap.clear();
    }
    /**
     * Cleanup method to prevent memory leaks
     */
    cleanup() {
        this.providersMap.clear();
    }
    /**
     * Merge token usage records
     * @param current The current token usage
     * @param update The new token usage to merge
     * @returns The merged token usage
     */
    mergeUsage(current, update) {
        // Create a result object with all properties initialized
        const result = {
            prompt: (current.prompt ?? 0) + (update.prompt ?? 0),
            completion: (current.completion ?? 0) + (update.completion ?? 0),
            cached: (current.cached ?? 0) + (update.cached ?? 0),
            total: (current.total ?? 0) + (update.total ?? 0),
            numRequests: (current.numRequests ?? 0) + (update.numRequests ?? 0),
            completionDetails: {
                reasoning: (current.completionDetails?.reasoning ?? 0) + (update.completionDetails?.reasoning ?? 0),
                acceptedPrediction: (current.completionDetails?.acceptedPrediction ?? 0) +
                    (update.completionDetails?.acceptedPrediction ?? 0),
                rejectedPrediction: (current.completionDetails?.rejectedPrediction ?? 0) +
                    (update.completionDetails?.rejectedPrediction ?? 0),
            },
            assertions: {
                total: (current.assertions?.total ?? 0) + (update.assertions?.total ?? 0),
                prompt: (current.assertions?.prompt ?? 0) + (update.assertions?.prompt ?? 0),
                completion: (current.assertions?.completion ?? 0) + (update.assertions?.completion ?? 0),
                cached: (current.assertions?.cached ?? 0) + (update.assertions?.cached ?? 0),
            },
        };
        return result;
    }
}
exports.TokenUsageTracker = TokenUsageTracker;
//# sourceMappingURL=tokenUsage.js.map