
import { raceCancellation } from '../../../base/common/async.js';
import { CancellationToken } from '../../../base/common/cancellation.js';
import { CancellationError } from '../../../base/common/errors.js';
import { toDisposable } from '../../../base/common/lifecycle.js';
import { revive } from '../../../base/common/marshalling.js';
import { generateUuid } from '../../../base/common/uuid.js';
import { isToolInvocationContext } from '../../contrib/chat/common/tools/languageModelToolsService.js';
import { computeCombinationKey } from '../../contrib/chat/common/tools/languageModelToolsConfirmationService.js';
import { InternalEditToolId, ExtensionEditToolId } from '../../contrib/chat/common/tools/builtinTools/editFileTool.js';
import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/builtinTools/tools.js';
import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
import { isProposedApiEnabled, checkProposedApiEnabled } from '../../services/extensions/common/extensions.js';
import { SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
import { MainContext } from './extHost.protocol.js';
import { LanguageModelToolSource, LanguageModelToolResult, MarkdownString, IconPath } from './extHostTypeConverters.js';
import { URI } from '../../../base/common/uri.js';

class Tool {
    constructor(data) {
        this._data = data;
    }
    update(newData) {
        this._data = newData;
        this._apiObject = undefined;
        this._apiObjectWithChatParticipantAdditions = undefined;
    }
    get data() {
        return this._data;
    }
    get apiObject() {
        if (!this._apiObject) {
            this._apiObject = ( Object.freeze({
                name: this._data.id,
                description: this._data.modelDescription,
                inputSchema: this._data.inputSchema,
                tags: this._data.tags ?? [],
                source: undefined
            }));
        }
        return this._apiObject;
    }
    get apiObjectWithChatParticipantAdditions() {
        if (!this._apiObjectWithChatParticipantAdditions) {
            this._apiObjectWithChatParticipantAdditions = ( Object.freeze({
                name: this._data.id,
                description: this._data.modelDescription,
                inputSchema: this._data.inputSchema,
                tags: this._data.tags ?? [],
                source: LanguageModelToolSource.to(this._data.source),
                fullReferenceName: this._data.fullReferenceName
            }));
        }
        return this._apiObjectWithChatParticipantAdditions;
    }
}
class ExtHostLanguageModelTools {
    constructor(mainContext, _languageModels) {
        this._languageModels = _languageModels;
        this._registeredTools = ( new Map());
        this._tokenCountFuncs = ( new Map());
        this._allTools = ( new Map());
        this._proxy = ( mainContext.getProxy(MainContext.MainThreadLanguageModelTools));
        this._proxy.$getTools().then(tools => {
            for (const tool of tools) {
                this._allTools.set(tool.id, ( new Tool(revive(tool))));
            }
        });
    }
    async $countTokensForInvocation(callId, input, token) {
        const fn = this._tokenCountFuncs.get(callId);
        if (!fn) {
            throw ( new Error(`Tool invocation call ${callId} not found`));
        }
        return await fn(input, token);
    }
    async invokeTool(extension, toolIdOrInfo, options, token) {
        const toolId = typeof toolIdOrInfo === "string" ? toolIdOrInfo : toolIdOrInfo.name;
        const callId = generateUuid();
        if (options.tokenizationOptions) {
            this._tokenCountFuncs.set(callId, options.tokenizationOptions.countTokens);
        }
        try {
            if (options.toolInvocationToken && !isToolInvocationContext(options.toolInvocationToken)) {
                throw ( new Error(`Invalid tool invocation token`));
            }
            if ((toolId === InternalEditToolId || toolId === ExtensionEditToolId) && !isProposedApiEnabled(extension, "chatParticipantPrivate")) {
                throw ( new Error(`Invalid tool: ${toolId}`));
            }
            const result = await this._proxy.$invokeTool({
                toolId,
                callId,
                parameters: options.input,
                tokenBudget: options.tokenizationOptions?.tokenBudget,
                context: options.toolInvocationToken,
                chatRequestId: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.chatRequestId : undefined,
                chatInteractionId: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.chatInteractionId : undefined,
                subAgentInvocationId: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.subAgentInvocationId : undefined,
                chatStreamToolCallId: isProposedApiEnabled(extension, "chatParticipantAdditions") ? options.chatStreamToolCallId : undefined,
                preToolUseResult: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.preToolUseResult : undefined,
                traceparent: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.traceparent : undefined,
                tracestate: isProposedApiEnabled(extension, "chatParticipantPrivate") ? options.tracestate : undefined
            }, token);
            const dto = result instanceof SerializableObjectWithBuffers ? result.value : result;
            return LanguageModelToolResult.to(revive(dto));
        } finally {
            this._tokenCountFuncs.delete(callId);
        }
    }
    $onDidChangeTools(tools) {
        const oldTools = ( new Set(( this._allTools.keys())));
        for (const tool of tools) {
            oldTools.delete(tool.id);
            const existing = this._allTools.get(tool.id);
            if (existing) {
                existing.update(tool);
            } else {
                this._allTools.set(tool.id, ( new Tool(revive(tool))));
            }
        }
        for (const id of oldTools) {
            this._allTools.delete(id);
        }
    }
    getTools(extension) {
        const hasParticipantAdditions = isProposedApiEnabled(extension);
        return ( Array.from(( this._allTools.values())).map(
            tool => hasParticipantAdditions ? tool.apiObjectWithChatParticipantAdditions : tool.apiObject
        )).filter(tool => {
            switch (tool.name) {
            case InternalEditToolId:
            case ExtensionEditToolId:
            case InternalFetchWebPageToolId:
            case SearchExtensionsToolId:
                return isProposedApiEnabled(extension);
            default:
                return true;
            }
        });
    }
    async $invokeTool(dto, token) {
        const item = this._registeredTools.get(dto.toolId);
        if (!item) {
            throw ( new Error(`Unknown tool ${dto.toolId}`));
        }
        const options = {
            input: dto.parameters,
            toolInvocationToken: revive(dto.context)
        };
        if (isProposedApiEnabled(item.extension)) {
            options.chatRequestId = dto.chatRequestId;
            options.chatInteractionId = dto.chatInteractionId;
            options.chatSessionResource = URI.revive(dto.context?.sessionResource);
            options.workingDirectory = URI.revive(dto.context?.workingDirectory);
            options.subAgentInvocationId = dto.subAgentInvocationId;
            options.traceparent = dto.traceparent;
            options.tracestate = dto.tracestate;
        }
        if (isProposedApiEnabled(item.extension) && dto.modelId) {
            options.model = await this.getModel(dto.modelId, item.extension);
        }
        if (isProposedApiEnabled(item.extension) && dto.chatStreamToolCallId) {
            options.chatStreamToolCallId = dto.chatStreamToolCallId;
        }
        if (dto.tokenBudget !== undefined) {
            options.tokenizationOptions = {
                tokenBudget: dto.tokenBudget,
                countTokens: this._tokenCountFuncs.get(dto.callId) || ((value, token = CancellationToken.None) => this._proxy.$countTokensForInvocation(dto.callId, value, token))
            };
        }
        let progress;
        if (isProposedApiEnabled(item.extension)) {
            let lastProgress;
            progress = {
                report: value => {
                    if (value.increment !== undefined) {
                        lastProgress = (lastProgress ?? 0) + value.increment;
                    }
                    this._proxy.$acceptToolProgress(dto.callId, {
                        message: MarkdownString.fromStrict(value.message),
                        progress: lastProgress === undefined ? undefined : lastProgress / 100
                    });
                }
            };
        }
        const extensionResult = await raceCancellation(Promise.resolve(item.tool.invoke(options, token, progress)), token);
        if (!extensionResult) {
            throw ( new CancellationError());
        }
        return LanguageModelToolResult.from(extensionResult, item.extension);
    }
    async getModel(modelId, extension) {
        let model;
        if (modelId) {
            model = await this._languageModels.getLanguageModelByIdentifier(extension, modelId);
        }
        if (!model) {
            model = await this._languageModels.getDefaultLanguageModel(extension);
            if (!model) {
                throw ( new Error("Language model unavailable"));
            }
        }
        return model;
    }
    async $handleToolStream(toolId, context, token) {
        const item = this._registeredTools.get(toolId);
        if (!item) {
            throw ( new Error(`Unknown tool ${toolId}`));
        }
        if (!item.tool.handleToolStream) {
            return undefined;
        }
        checkProposedApiEnabled(item.extension, "chatParticipantAdditions");
        const options = {
            rawInput: context.rawInput,
            chatRequestId: context.chatRequestId,
            chatSessionResource: context.chatSessionResource,
            chatInteractionId: context.chatInteractionId
        };
        const result = await item.tool.handleToolStream(options, token);
        if (!result) {
            return undefined;
        }
        return {
            invocationMessage: MarkdownString.fromStrict(result.invocationMessage)
        };
    }
    async $prepareToolInvocation(toolId, context, token) {
        const item = this._registeredTools.get(toolId);
        if (!item) {
            throw ( new Error(`Unknown tool ${toolId}`));
        }
        const options = {
            input: context.parameters,
            chatRequestId: context.chatRequestId,
            chatSessionResource: context.chatSessionResource,
            chatInteractionId: context.chatInteractionId,
            workingDirectory: URI.revive(context.workingDirectory),
            forceConfirmationReason: context.forceConfirmationReason
        };
        if (context.forceConfirmationReason) {
            checkProposedApiEnabled(item.extension, "chatParticipantPrivate");
        }
        if (item.tool.prepareInvocation) {
            const result = await item.tool.prepareInvocation(options, token);
            if (!result) {
                return undefined;
            }
            if (result.pastTenseMessage || result.presentation) {
                checkProposedApiEnabled(item.extension, "chatParticipantPrivate");
            }
            if (result.confirmationMessages?.approveCombination !== undefined) {
                checkProposedApiEnabled(item.extension, "toolInvocationApproveCombination");
            }
            const approveCombination = result.confirmationMessages?.approveCombination;
            const approveCombinationLabel = approveCombination ? MarkdownString.fromStrict(approveCombination.message) : undefined;
            const approveCombinationKey = approveCombinationLabel ? await computeCombinationKey(toolId, context.parameters) : undefined;
            return {
                confirmationMessages: result.confirmationMessages ? {
                    title: typeof result.confirmationMessages.title === "string" ? result.confirmationMessages.title : MarkdownString.from(result.confirmationMessages.title),
                    message: typeof result.confirmationMessages.message === "string" ? result.confirmationMessages.message : MarkdownString.from(result.confirmationMessages.message),
                    approveCombination: approveCombinationLabel && approveCombinationKey ? {
                        label: approveCombinationLabel,
                        key: approveCombinationKey,
                        arguments: approveCombination.arguments
                    } : undefined
                } : undefined,
                invocationMessage: MarkdownString.fromStrict(result.invocationMessage),
                pastTenseMessage: MarkdownString.fromStrict(result.pastTenseMessage),
                presentation: result.presentation
            };
        }
        return undefined;
    }
    registerTool(extension, id, tool) {
        this._registeredTools.set(id, {
            extension,
            tool
        });
        this._proxy.$registerTool(id, typeof tool.handleToolStream === "function");
        return toDisposable(() => {
            this._registeredTools.delete(id);
            this._proxy.$unregisterTool(id);
        });
    }
    registerToolDefinition(extension, definition, tool) {
        checkProposedApiEnabled(extension, "languageModelToolSupportsModel");
        const id = definition.name;
        const dto = {
            id,
            displayName: definition.displayName,
            toolReferenceName: definition.toolReferenceName,
            userDescription: definition.userDescription,
            modelDescription: definition.description,
            inputSchema: definition.inputSchema,
            source: {
                type: "extension",
                label: extension.displayName ?? extension.name,
                extensionId: extension.identifier
            },
            icon: IconPath.from(definition.icon),
            models: definition.models,
            toolSet: definition.toolSet
        };
        this._registeredTools.set(id, {
            extension,
            tool
        });
        this._proxy.$registerToolWithDefinition(extension.identifier, dto, typeof tool.handleToolStream === "function");
        return toDisposable(() => {
            this._registeredTools.delete(id);
            this._proxy.$unregisterTool(id);
        });
    }
}

export { ExtHostLanguageModelTools };
