/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
import { getDirectoryContextString, getEnvironmentContext, } from '../utils/environmentContext.js';
import { CompressionStatus } from './turn.js';
import { Turn, GeminiEventType } from './turn.js';
import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js';
import { getResponseText } from '../utils/partUtils.js';
import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
import { reportError } from '../utils/errorReporting.js';
import { GeminiChat } from './geminiChat.js';
import { retryWithBackoff } from '../utils/retry.js';
import { getErrorMessage } from '../utils/errors.js';
import { tokenLimit } from './tokenLimits.js';
import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL_AUTO, DEFAULT_THINKING_MODE, getEffectiveModel, } from '../config/models.js';
import { LoopDetectionService } from '../services/loopDetectionService.js';
import { ideContextStore } from '../ide/ideContext.js';
import { logChatCompression, logNextSpeakerCheck, } from '../telemetry/loggers.js';
import { makeChatCompressionEvent, NextSpeakerCheckEvent, } from '../telemetry/types.js';
import { handleFallback } from '../fallback/handler.js';
import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
export function isThinkingSupported(model) {
    return model.startsWith('gemini-2.5') || model === DEFAULT_GEMINI_MODEL_AUTO;
}
export function isThinkingDefault(model) {
    if (model.startsWith('gemini-2.5-flash-lite')) {
        return false;
    }
    return model.startsWith('gemini-2.5') || model === DEFAULT_GEMINI_MODEL_AUTO;
}
/**
 * Returns the index of the oldest item to keep when compressing. May return
 * contents.length which indicates that everything should be compressed.
 *
 * Exported for testing purposes.
 */
export function findCompressSplitPoint(contents, fraction) {
    if (fraction <= 0 || fraction >= 1) {
        throw new Error('Fraction must be between 0 and 1');
    }
    const charCounts = contents.map((content) => JSON.stringify(content).length);
    const totalCharCount = charCounts.reduce((a, b) => a + b, 0);
    const targetCharCount = totalCharCount * fraction;
    let lastSplitPoint = 0; // 0 is always valid (compress nothing)
    let cumulativeCharCount = 0;
    for (let i = 0; i < contents.length; i++) {
        cumulativeCharCount += charCounts[i];
        const content = contents[i];
        if (content.role === 'user' &&
            !content.parts?.some((part) => !!part.functionResponse)) {
            if (cumulativeCharCount >= targetCharCount) {
                return i;
            }
            lastSplitPoint = i;
        }
    }
    // We found no split points after targetCharCount.
    // Check if it's safe to compress everything.
    const lastContent = contents[contents.length - 1];
    if (lastContent?.role === 'model' &&
        !lastContent?.parts?.some((part) => part.functionCall)) {
        return contents.length;
    }
    // Can't compress everything so just compress at last splitpoint.
    return lastSplitPoint;
}
const MAX_TURNS = 100;
/**
 * Threshold for compression token count as a fraction of the model's token limit.
 * If the chat history exceeds this threshold, it will be compressed.
 */
const COMPRESSION_TOKEN_THRESHOLD = 0.7;
/**
 * The fraction of the latest chat history to keep. A value of 0.3
 * means that only the last 30% of the chat history will be kept after compression.
 */
const COMPRESSION_PRESERVE_THRESHOLD = 0.3;
export class GeminiClient {
    config;
    chat;
    generateContentConfig = {
        temperature: 0,
        topP: 1,
    };
    sessionTurnCount = 0;
    loopDetector;
    lastPromptId;
    currentSequenceModel = null;
    lastSentIdeContext;
    forceFullIdeContext = true;
    /**
     * At any point in this conversation, was compression triggered without
     * being forced and did it fail?
     */
    hasFailedCompressionAttempt = false;
    constructor(config) {
        this.config = config;
        this.loopDetector = new LoopDetectionService(config);
        this.lastPromptId = this.config.getSessionId();
    }
    async initialize() {
        this.chat = await this.startChat();
    }
    getContentGeneratorOrFail() {
        if (!this.config.getContentGenerator()) {
            throw new Error('Content generator not initialized');
        }
        return this.config.getContentGenerator();
    }
    async addHistory(content) {
        this.getChat().addHistory(content);
    }
    getChat() {
        if (!this.chat) {
            throw new Error('Chat not initialized');
        }
        return this.chat;
    }
    isInitialized() {
        return this.chat !== undefined;
    }
    getHistory() {
        return this.getChat().getHistory();
    }
    stripThoughtsFromHistory() {
        this.getChat().stripThoughtsFromHistory();
    }
    setHistory(history) {
        this.getChat().setHistory(history);
        this.forceFullIdeContext = true;
    }
    async setTools() {
        const toolRegistry = this.config.getToolRegistry();
        const toolDeclarations = toolRegistry.getFunctionDeclarations();
        const tools = [{ functionDeclarations: toolDeclarations }];
        this.getChat().setTools(tools);
    }
    async resetChat() {
        this.chat = await this.startChat();
    }
    getChatRecordingService() {
        return this.chat?.getChatRecordingService();
    }
    getLoopDetectionService() {
        return this.loopDetector;
    }
    async addDirectoryContext() {
        if (!this.chat) {
            return;
        }
        this.getChat().addHistory({
            role: 'user',
            parts: [{ text: await getDirectoryContextString(this.config) }],
        });
    }
    async startChat(extraHistory) {
        this.forceFullIdeContext = true;
        this.hasFailedCompressionAttempt = false;
        const envParts = await getEnvironmentContext(this.config);
        const toolRegistry = this.config.getToolRegistry();
        const toolDeclarations = toolRegistry.getFunctionDeclarations();
        const tools = [{ functionDeclarations: toolDeclarations }];
        const history = [
            {
                role: 'user',
                parts: envParts,
            },
            {
                role: 'model',
                parts: [{ text: 'Got it. Thanks for the context!' }],
            },
            ...(extraHistory ?? []),
        ];
        try {
            const userMemory = this.config.getUserMemory();
            const systemInstruction = getCoreSystemPrompt(userMemory);
            const model = this.config.getModel();
            const config = { ...this.generateContentConfig };
            if (isThinkingSupported(model)) {
                config.thinkingConfig = {
                    includeThoughts: true,
                    thinkingBudget: DEFAULT_THINKING_MODE,
                };
            }
            return new GeminiChat(this.config, {
                systemInstruction,
                ...config,
                tools,
            }, history);
        }
        catch (error) {
            await reportError(error, 'Error initializing Gemini chat session.', history, 'startChat');
            throw new Error(`Failed to initialize chat: ${getErrorMessage(error)}`);
        }
    }
    getIdeContextParts(forceFullContext) {
        const currentIdeContext = ideContextStore.get();
        if (!currentIdeContext) {
            return { contextParts: [], newIdeContext: undefined };
        }
        if (forceFullContext || !this.lastSentIdeContext) {
            // Send full context as JSON
            const openFiles = currentIdeContext.workspaceState?.openFiles || [];
            const activeFile = openFiles.find((f) => f.isActive);
            const otherOpenFiles = openFiles
                .filter((f) => !f.isActive)
                .map((f) => f.path);
            const contextData = {};
            if (activeFile) {
                contextData['activeFile'] = {
                    path: activeFile.path,
                    cursor: activeFile.cursor
                        ? {
                            line: activeFile.cursor.line,
                            character: activeFile.cursor.character,
                        }
                        : undefined,
                    selectedText: activeFile.selectedText || undefined,
                };
            }
            if (otherOpenFiles.length > 0) {
                contextData['otherOpenFiles'] = otherOpenFiles;
            }
            if (Object.keys(contextData).length === 0) {
                return { contextParts: [], newIdeContext: currentIdeContext };
            }
            const jsonString = JSON.stringify(contextData, null, 2);
            const contextParts = [
                "Here is the user's editor context as a JSON object. This is for your information only.",
                '```json',
                jsonString,
                '```',
            ];
            if (this.config.getDebugMode()) {
                console.log(contextParts.join('\n'));
            }
            return {
                contextParts,
                newIdeContext: currentIdeContext,
            };
        }
        else {
            // Calculate and send delta as JSON
            const delta = {};
            const changes = {};
            const lastFiles = new Map((this.lastSentIdeContext.workspaceState?.openFiles || []).map((f) => [f.path, f]));
            const currentFiles = new Map((currentIdeContext.workspaceState?.openFiles || []).map((f) => [
                f.path,
                f,
            ]));
            const openedFiles = [];
            for (const [path] of currentFiles.entries()) {
                if (!lastFiles.has(path)) {
                    openedFiles.push(path);
                }
            }
            if (openedFiles.length > 0) {
                changes['filesOpened'] = openedFiles;
            }
            const closedFiles = [];
            for (const [path] of lastFiles.entries()) {
                if (!currentFiles.has(path)) {
                    closedFiles.push(path);
                }
            }
            if (closedFiles.length > 0) {
                changes['filesClosed'] = closedFiles;
            }
            const lastActiveFile = (this.lastSentIdeContext.workspaceState?.openFiles || []).find((f) => f.isActive);
            const currentActiveFile = (currentIdeContext.workspaceState?.openFiles || []).find((f) => f.isActive);
            if (currentActiveFile) {
                if (!lastActiveFile || lastActiveFile.path !== currentActiveFile.path) {
                    changes['activeFileChanged'] = {
                        path: currentActiveFile.path,
                        cursor: currentActiveFile.cursor
                            ? {
                                line: currentActiveFile.cursor.line,
                                character: currentActiveFile.cursor.character,
                            }
                            : undefined,
                        selectedText: currentActiveFile.selectedText || undefined,
                    };
                }
                else {
                    const lastCursor = lastActiveFile.cursor;
                    const currentCursor = currentActiveFile.cursor;
                    if (currentCursor &&
                        (!lastCursor ||
                            lastCursor.line !== currentCursor.line ||
                            lastCursor.character !== currentCursor.character)) {
                        changes['cursorMoved'] = {
                            path: currentActiveFile.path,
                            cursor: {
                                line: currentCursor.line,
                                character: currentCursor.character,
                            },
                        };
                    }
                    const lastSelectedText = lastActiveFile.selectedText || '';
                    const currentSelectedText = currentActiveFile.selectedText || '';
                    if (lastSelectedText !== currentSelectedText) {
                        changes['selectionChanged'] = {
                            path: currentActiveFile.path,
                            selectedText: currentSelectedText,
                        };
                    }
                }
            }
            else if (lastActiveFile) {
                changes['activeFileChanged'] = {
                    path: null,
                    previousPath: lastActiveFile.path,
                };
            }
            if (Object.keys(changes).length === 0) {
                return { contextParts: [], newIdeContext: currentIdeContext };
            }
            delta['changes'] = changes;
            const jsonString = JSON.stringify(delta, null, 2);
            const contextParts = [
                "Here is a summary of changes in the user's editor context, in JSON format. This is for your information only.",
                '```json',
                jsonString,
                '```',
            ];
            if (this.config.getDebugMode()) {
                console.log(contextParts.join('\n'));
            }
            return {
                contextParts,
                newIdeContext: currentIdeContext,
            };
        }
    }
    async *sendMessageStream(request, signal, prompt_id, turns = MAX_TURNS) {
        if (this.lastPromptId !== prompt_id) {
            this.loopDetector.reset(prompt_id);
            this.lastPromptId = prompt_id;
            this.currentSequenceModel = null;
        }
        this.sessionTurnCount++;
        if (this.config.getMaxSessionTurns() > 0 &&
            this.sessionTurnCount > this.config.getMaxSessionTurns()) {
            yield { type: GeminiEventType.MaxSessionTurns };
            return new Turn(this.getChat(), prompt_id);
        }
        // Ensure turns never exceeds MAX_TURNS to prevent infinite loops
        const boundedTurns = Math.min(turns, MAX_TURNS);
        if (!boundedTurns) {
            return new Turn(this.getChat(), prompt_id);
        }
        const compressed = await this.tryCompressChat(prompt_id, false);
        if (compressed.compressionStatus === CompressionStatus.COMPRESSED) {
            yield { type: GeminiEventType.ChatCompressed, value: compressed };
        }
        // Prevent context updates from being sent while a tool call is
        // waiting for a response. The Gemini API requires that a functionResponse
        // part from the user immediately follows a functionCall part from the model
        // in the conversation history . The IDE context is not discarded; it will
        // be included in the next regular message sent to the model.
        const history = this.getHistory();
        const lastMessage = history.length > 0 ? history[history.length - 1] : undefined;
        const hasPendingToolCall = !!lastMessage &&
            lastMessage.role === 'model' &&
            (lastMessage.parts?.some((p) => 'functionCall' in p) || false);
        if (this.config.getIdeMode() && !hasPendingToolCall) {
            const { contextParts, newIdeContext } = this.getIdeContextParts(this.forceFullIdeContext || history.length === 0);
            if (contextParts.length > 0) {
                this.getChat().addHistory({
                    role: 'user',
                    parts: [{ text: contextParts.join('\n') }],
                });
            }
            this.lastSentIdeContext = newIdeContext;
            this.forceFullIdeContext = false;
        }
        const turn = new Turn(this.getChat(), prompt_id);
        const controller = new AbortController();
        const linkedSignal = AbortSignal.any([signal, controller.signal]);
        const loopDetected = await this.loopDetector.turnStarted(signal);
        if (loopDetected) {
            yield { type: GeminiEventType.LoopDetected };
            return turn;
        }
        const routingContext = {
            history: this.getChat().getHistory(/*curated=*/ true),
            request,
            signal,
        };
        let modelToUse;
        // Determine Model (Stickiness vs. Routing)
        if (this.currentSequenceModel) {
            modelToUse = this.currentSequenceModel;
        }
        else {
            const router = await this.config.getModelRouterService();
            const decision = await router.route(routingContext);
            modelToUse = decision.model;
            // Lock the model for the rest of the sequence
            this.currentSequenceModel = modelToUse;
        }
        const resultStream = turn.run(modelToUse, request, linkedSignal);
        for await (const event of resultStream) {
            if (this.loopDetector.addAndCheck(event)) {
                yield { type: GeminiEventType.LoopDetected };
                controller.abort();
                return turn;
            }
            yield event;
            if (event.type === GeminiEventType.Error) {
                return turn;
            }
        }
        if (!turn.pendingToolCalls.length && signal && !signal.aborted) {
            // Check if next speaker check is needed
            if (this.config.getQuotaErrorOccurred()) {
                return turn;
            }
            if (this.config.getSkipNextSpeakerCheck()) {
                return turn;
            }
            const nextSpeakerCheck = await checkNextSpeaker(this.getChat(), this.config.getBaseLlmClient(), signal, prompt_id);
            logNextSpeakerCheck(this.config, new NextSpeakerCheckEvent(prompt_id, turn.finishReason?.toString() || '', nextSpeakerCheck?.next_speaker || ''));
            if (nextSpeakerCheck?.next_speaker === 'model') {
                const nextRequest = [{ text: 'Please continue.' }];
                // This recursive call's events will be yielded out, but the final
                // turn object will be from the top-level call.
                yield* this.sendMessageStream(nextRequest, signal, prompt_id, boundedTurns - 1);
            }
        }
        return turn;
    }
    async generateContent(contents, generationConfig, abortSignal, model) {
        let currentAttemptModel = model;
        const configToUse = {
            ...this.generateContentConfig,
            ...generationConfig,
        };
        try {
            const userMemory = this.config.getUserMemory();
            const systemInstruction = getCoreSystemPrompt(userMemory);
            const requestConfig = {
                abortSignal,
                ...configToUse,
                systemInstruction,
            };
            const apiCall = () => {
                const modelToUse = this.config.isInFallbackMode()
                    ? DEFAULT_GEMINI_FLASH_MODEL
                    : model;
                currentAttemptModel = modelToUse;
                return this.getContentGeneratorOrFail().generateContent({
                    model: modelToUse,
                    config: requestConfig,
                    contents,
                }, this.lastPromptId);
            };
            const onPersistent429Callback = async (authType, error) => 
            // Pass the captured model to the centralized handler.
            await handleFallback(this.config, currentAttemptModel, authType, error);
            const result = await retryWithBackoff(apiCall, {
                onPersistent429: onPersistent429Callback,
                authType: this.config.getContentGeneratorConfig()?.authType,
            });
            return result;
        }
        catch (error) {
            if (abortSignal.aborted) {
                throw error;
            }
            await reportError(error, `Error generating content via API with model ${currentAttemptModel}.`, {
                requestContents: contents,
                requestConfig: configToUse,
            }, 'generateContent-api');
            throw new Error(`Failed to generate content with model ${currentAttemptModel}: ${getErrorMessage(error)}`);
        }
    }
    async tryCompressChat(prompt_id, force = false) {
        // If the model is 'auto', we will use a placeholder model to check.
        // Compression occurs before we choose a model, so calling `count_tokens`
        // before the model is chosen would result in an error.
        const configModel = this.config.getModel();
        let model = configModel === DEFAULT_GEMINI_MODEL_AUTO
            ? DEFAULT_GEMINI_MODEL
            : configModel;
        // Check if the model needs to be a fallback
        model = getEffectiveModel(this.config.isInFallbackMode(), model);
        const curatedHistory = this.getChat().getHistory(true);
        // Regardless of `force`, don't do anything if the history is empty.
        if (curatedHistory.length === 0 ||
            (this.hasFailedCompressionAttempt && !force)) {
            return {
                originalTokenCount: 0,
                newTokenCount: 0,
                compressionStatus: CompressionStatus.NOOP,
            };
        }
        const { totalTokens: originalTokenCount } = await this.getContentGeneratorOrFail().countTokens({
            model,
            contents: curatedHistory,
        });
        if (originalTokenCount === undefined) {
            console.warn(`Could not determine token count for model ${model}.`);
            this.hasFailedCompressionAttempt = !force && true;
            return {
                originalTokenCount: 0,
                newTokenCount: 0,
                compressionStatus: CompressionStatus.COMPRESSION_FAILED_TOKEN_COUNT_ERROR,
            };
        }
        const contextPercentageThreshold = this.config.getChatCompression()?.contextPercentageThreshold;
        // Don't compress if not forced and we are under the limit.
        if (!force) {
            const threshold = contextPercentageThreshold ?? COMPRESSION_TOKEN_THRESHOLD;
            if (originalTokenCount < threshold * tokenLimit(model)) {
                return {
                    originalTokenCount,
                    newTokenCount: originalTokenCount,
                    compressionStatus: CompressionStatus.NOOP,
                };
            }
        }
        const splitPoint = findCompressSplitPoint(curatedHistory, 1 - COMPRESSION_PRESERVE_THRESHOLD);
        const historyToCompress = curatedHistory.slice(0, splitPoint);
        const historyToKeep = curatedHistory.slice(splitPoint);
        const summaryResponse = await this.config
            .getContentGenerator()
            .generateContent({
            model,
            contents: [
                ...historyToCompress,
                {
                    role: 'user',
                    parts: [
                        {
                            text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
                        },
                    ],
                },
            ],
            config: {
                systemInstruction: { text: getCompressionPrompt() },
            },
        }, prompt_id);
        const summary = getResponseText(summaryResponse) ?? '';
        const chat = await this.startChat([
            {
                role: 'user',
                parts: [{ text: summary }],
            },
            {
                role: 'model',
                parts: [{ text: 'Got it. Thanks for the additional context!' }],
            },
            ...historyToKeep,
        ]);
        this.forceFullIdeContext = true;
        const { totalTokens: newTokenCount } = await this.getContentGeneratorOrFail().countTokens({
            model,
            contents: chat.getHistory(),
        });
        if (newTokenCount === undefined) {
            console.warn('Could not determine compressed history token count.');
            this.hasFailedCompressionAttempt = !force && true;
            return {
                originalTokenCount,
                newTokenCount: originalTokenCount,
                compressionStatus: CompressionStatus.COMPRESSION_FAILED_TOKEN_COUNT_ERROR,
            };
        }
        uiTelemetryService.setLastPromptTokenCount(newTokenCount);
        logChatCompression(this.config, makeChatCompressionEvent({
            tokens_before: originalTokenCount,
            tokens_after: newTokenCount,
        }));
        if (newTokenCount > originalTokenCount) {
            this.getChat().setHistory(curatedHistory);
            this.hasFailedCompressionAttempt = !force && true;
            return {
                originalTokenCount,
                newTokenCount,
                compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
            };
        }
        else {
            this.chat = chat; // Chat compression successful, set new state.
        }
        return {
            originalTokenCount,
            newTokenCount,
            compressionStatus: CompressionStatus.COMPRESSED,
        };
    }
}
export const TEST_ONLY = {
    COMPRESSION_PRESERVE_THRESHOLD,
    COMPRESSION_TOKEN_THRESHOLD,
};
//# sourceMappingURL=client.js.map