
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/languageModelToolsService.js';
import { MainContext } from './extHost.protocol.js';
import { LanguageModelToolResult, LanguageModelToolDescription, MarkdownString } from './extHostTypeConverters.js';

class ExtHostLanguageModelTools {
    constructor(mainContext) {
        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, 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(toolId, options, token) {
        const callId = generateUuid();
        if (options.tokenizationOptions) {
            this._tokenCountFuncs.set(callId, options.tokenizationOptions.countTokens);
        }
        if (options.toolInvocationToken && !isToolInvocationContext(options.toolInvocationToken)) {
            throw ( new Error(`Invalid tool invocation token`));
        }
        try {
            const result = await this._proxy.$invokeTool({
                toolId,
                callId,
                parameters: options.input,
                tokenBudget: options.tokenizationOptions?.tokenBudget,
                context: options.toolInvocationToken,
            }, token);
            return LanguageModelToolResult.to(result);
        }
        finally {
            this._tokenCountFuncs.delete(callId);
        }
    }
    $onDidChangeTools(tools) {
        this._allTools.clear();
        for (const tool of tools) {
            this._allTools.set(tool.id, tool);
        }
    }
    get tools() {
        return ( Array.from(( this._allTools.values()))
            .map(tool => LanguageModelToolDescription.to(tool)));
    }
    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: dto.context };
        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))
            };
        }
        const extensionResult = await raceCancellation(Promise.resolve(item.tool.invoke(options, token)), token);
        if (!extensionResult) {
            throw ( new CancellationError());
        }
        return LanguageModelToolResult.from(extensionResult);
    }
    async $prepareToolInvocation(toolId, input, token) {
        const item = this._registeredTools.get(toolId);
        if (!item) {
            throw ( new Error(`Unknown tool ${toolId}`));
        }
        if (!item.tool.prepareInvocation) {
            return undefined;
        }
        const options = { input };
        const result = await item.tool.prepareInvocation(options, token);
        if (!result) {
            return undefined;
        }
        return {
            confirmationMessages: result.confirmationMessages ? {
                title: result.confirmationMessages.title,
                message: typeof result.confirmationMessages.message === 'string' ? result.confirmationMessages.message : MarkdownString.from(result.confirmationMessages.message),
            } : undefined,
            invocationMessage: typeof result.invocationMessage === 'string' ?
                result.invocationMessage :
                (result.invocationMessage ?
                    MarkdownString.from(result.invocationMessage) :
                    undefined),
        };
    }
    registerTool(extension, id, tool) {
        this._registeredTools.set(id, { extension, tool });
        this._proxy.$registerTool(id);
        return toDisposable(() => {
            this._registeredTools.delete(id);
            this._proxy.$unregisterTool(id);
        });
    }
}

export { ExtHostLanguageModelTools };
