/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
import { ApiRequestEvent, ApiResponseEvent, ApiErrorEvent, } from '../telemetry/types.js';
import { logApiError, logApiRequest, logApiResponse, } from '../telemetry/loggers.js';
import { CodeAssistServer } from '../code_assist/server.js';
import { toContents } from '../code_assist/converter.js';
import { isStructuredError } from '../utils/quotaErrorDetection.js';
import { runInDevTraceSpan } from '../telemetry/trace.js';
import { TokenUsageTracker } from './tokenUsageTracker.js';
/**
 * A decorator that wraps a ContentGenerator to add logging to API calls.
 */
export class LoggingContentGenerator {
    wrapped;
    config;
    constructor(wrapped, config) {
        this.wrapped = wrapped;
        this.config = config;
    }
    getWrapped() {
        return this.wrapped;
    }
    logApiRequest(contents, model, promptId) {
        const requestText = JSON.stringify(contents);
        logApiRequest(this.config, new ApiRequestEvent(model, promptId, requestText));
    }
    _getEndpointUrl(req, method) {
        // Case 1: Authenticated with a Google account (`gcloud auth login`).
        // Requests are routed through the internal CodeAssistServer.
        if (this.wrapped instanceof CodeAssistServer) {
            const url = new URL(this.wrapped.getMethodUrl(method));
            const port = url.port
                ? parseInt(url.port, 10)
                : url.protocol === 'https:'
                    ? 443
                    : 80;
            return { address: url.hostname, port };
        }
        const genConfig = this.config.getContentGeneratorConfig();
        // Case 2: Using an API key for Vertex AI.
        if (genConfig?.vertexai) {
            const location = process.env['GOOGLE_CLOUD_LOCATION'];
            if (location) {
                return { address: `${location}-aiplatform.googleapis.com`, port: 443 };
            }
            else {
                return { address: 'unknown', port: 0 };
            }
        }
        // Case 3: Default to the public Gemini API endpoint.
        // This is used when an API key is provided but not for Vertex AI.
        return { address: `generativelanguage.googleapis.com`, port: 443 };
    }
    _logApiResponse(requestContents, durationMs, model, prompt_id, responseId, responseCandidates, usageMetadata, responseText, generationConfig, serverDetails) {
        logApiResponse(this.config, new ApiResponseEvent(model, durationMs, {
            prompt_id,
            contents: requestContents,
            generate_content_config: generationConfig,
            server: serverDetails,
        }, {
            candidates: responseCandidates,
            response_id: responseId,
        }, this.config.getContentGeneratorConfig()?.authType, usageMetadata, responseText));
    }
    _logApiError(durationMs, error, model, prompt_id, requestContents, generationConfig, serverDetails) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        const errorType = error instanceof Error ? error.name : 'unknown';
        logApiError(this.config, new ApiErrorEvent(model, errorMessage, durationMs, {
            prompt_id,
            contents: requestContents,
            generate_content_config: generationConfig,
            server: serverDetails,
        }, this.config.getContentGeneratorConfig()?.authType, errorType, isStructuredError(error)
            ? error.status
            : undefined));
    }
    async generateContent(req, userPromptId) {
        return runInDevTraceSpan({
            name: 'generateContent',
        }, async ({ metadata: spanMetadata }) => {
            spanMetadata.input = { request: req, userPromptId, model: req.model };
            const startTime = Date.now();
            const contents = toContents(req.contents);
            this.logApiRequest(toContents(req.contents), req.model, userPromptId);
            const serverDetails = this._getEndpointUrl(req, 'generateContent');
            try {
                const response = await this.wrapped.generateContent(req, userPromptId);
                spanMetadata.output = {
                    response,
                    usageMetadata: response.usageMetadata,
                };
                const durationMs = Date.now() - startTime;
                this._logApiResponse(contents, durationMs, response.modelVersion || req.model, userPromptId, response.responseId, response.candidates, response.usageMetadata, JSON.stringify(response), req.config, serverDetails);
                // Track token usage for subscription management
                if (response.usageMetadata?.totalTokenCount) {
                    await TokenUsageTracker.addTokenUsage(response.usageMetadata.totalTokenCount);
                    // Check if user should be blocked after token usage (async, non-blocking)
                    import('./subscriptionManager.js').then(async (m) => {
                        const status = await m.SubscriptionManager.getSubscriptionStatus();
                        if (status.shouldBlock) {
                            // Emit subscription blocked event
                            const { coreEvents } = await import('../utils/events.js');
                            coreEvents.emitSubscriptionBlocked(status.shouldBlock, status.totalTokensUsed, status.isSubscriber);
                        }
                    }).catch(() => {
                        // Ignore errors to not break API calls
                    });
                }
                return response;
            }
            catch (error) {
                const durationMs = Date.now() - startTime;
                this._logApiError(durationMs, error, req.model, userPromptId, contents, req.config, serverDetails);
                throw error;
            }
        });
    }
    async generateContentStream(req, userPromptId) {
        return runInDevTraceSpan({
            name: 'generateContentStream',
            noAutoEnd: true,
        }, async ({ metadata: spanMetadata, endSpan }) => {
            spanMetadata.input = { request: req, userPromptId, model: req.model };
            const startTime = Date.now();
            this.logApiRequest(toContents(req.contents), req.model, userPromptId);
            const serverDetails = this._getEndpointUrl(req, 'generateContentStream');
            let stream;
            try {
                stream = await this.wrapped.generateContentStream(req, userPromptId);
            }
            catch (error) {
                const durationMs = Date.now() - startTime;
                this._logApiError(durationMs, error, req.model, userPromptId, toContents(req.contents), req.config, serverDetails);
                throw error;
            }
            return this.loggingStreamWrapper(req, stream, startTime, userPromptId, spanMetadata, endSpan);
        });
    }
    async *loggingStreamWrapper(req, stream, startTime, userPromptId, spanMetadata, endSpan) {
        const responses = [];
        let lastUsageMetadata;
        const serverDetails = this._getEndpointUrl(req, 'generateContentStream');
        const requestContents = toContents(req.contents);
        try {
            for await (const response of stream) {
                responses.push(response);
                if (response.usageMetadata) {
                    lastUsageMetadata = response.usageMetadata;
                }
                yield response;
            }
            // Only log successful API response if no error occurred
            const durationMs = Date.now() - startTime;
            this._logApiResponse(requestContents, durationMs, responses[0]?.modelVersion || req.model, userPromptId, responses[0]?.responseId, responses.flatMap((response) => response.candidates || []), lastUsageMetadata, JSON.stringify(responses), req.config, serverDetails);
            // Track token usage for subscription management
            if (lastUsageMetadata?.totalTokenCount) {
                await TokenUsageTracker.addTokenUsage(lastUsageMetadata.totalTokenCount);
                // Check if user should be blocked after token usage (async, non-blocking)
                import('./subscriptionManager.js').then(async (m) => {
                    const status = await m.SubscriptionManager.getSubscriptionStatus();
                    if (status.shouldBlock) {
                        // Emit subscription blocked event
                        const { coreEvents } = await import('../utils/events.js');
                        coreEvents.emitSubscriptionBlocked(status.shouldBlock, status.totalTokensUsed, status.isSubscriber);
                    }
                }).catch(() => {
                    // Ignore errors to not break API calls
                });
            }
            spanMetadata.output = {
                streamChunks: responses.map((r) => ({
                    content: r.candidates?.[0]?.content ?? null,
                })),
                usageMetadata: lastUsageMetadata,
                durationMs,
            };
        }
        catch (error) {
            spanMetadata.error = error;
            const durationMs = Date.now() - startTime;
            this._logApiError(durationMs, error, responses[0]?.modelVersion || req.model, userPromptId, requestContents, req.config, serverDetails);
            throw error;
        }
        finally {
            endSpan();
        }
    }
    async countTokens(req) {
        return this.wrapped.countTokens(req);
    }
    async embedContent(req) {
        return runInDevTraceSpan({
            name: 'embedContent',
        }, async ({ metadata: spanMetadata }) => {
            spanMetadata.input = { request: req };
            const output = await this.wrapped.embedContent(req);
            spanMetadata.output = output;
            return output;
        });
    }
}
//# sourceMappingURL=loggingContentGenerator.js.map