/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { decodeBase64, VSBuffer } from '../../../../base/common/buffer.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { Codicon } from '../../../../base/common/codicons.js';
import { MarkdownString } from '../../../../base/common/htmlContent.js';
import { Lazy } from '../../../../base/common/lazy.js';
import { Disposable, DisposableMap, DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
import { equals } from '../../../../base/common/objects.js';
import { autorun } from '../../../../base/common/observable.js';
import { basename } from '../../../../base/common/resources.js';
import { isDefined, Mutable } from '../../../../base/common/types.js';
import { URI } from '../../../../base/common/uri.js';
import { localize } from '../../../../nls.js';
import { IFileService } from '../../../../platform/files/common/files.js';
import { IImageResizeService } from '../../../../platform/imageResize/common/imageResizeService.js';
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
import { mcpAppsEnabledConfig } from '../../../../platform/mcp/common/mcpManagement.js';
import { IProductService } from '../../../../platform/product/common/productService.js';
import { StorageScope } from '../../../../platform/storage/common/storage.js';
import { IWorkbenchContribution } from '../../../common/contributions.js';
import { ChatResponseResource, getAttachableImageExtension } from '../../chat/common/model/chatModel.js';
import { LanguageModelPartAudience } from '../../chat/common/languageModels.js';
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolConfirmationMessages, IToolData, IToolImpl, IToolInvocation, IToolInvocationPreparationContext, IToolResult, IToolResultInputOutputDetails, ToolDataSource, ToolProgress, ToolSet } from '../../chat/common/tools/languageModelToolsService.js';
import { IMcpRegistry } from './mcpRegistryTypes.js';
import { IMcpServer, IMcpService, IMcpTool, IMcpToolResourceLinkContents, McpResourceURI, McpToolResourceLinkMimeType, McpToolVisibility } from './mcpTypes.js';
import { mcpServerToSourceData } from './mcpTypesUtils.js';
import { ILifecycleService } from '../../../services/lifecycle/common/lifecycle.js';
import { McpServer } from './mcpServer.js';

interface ISyncedToolData {
	toolData: IToolData;
	store: DisposableStore;
}

export class McpLanguageModelToolContribution extends Disposable implements IWorkbenchContribution {

	public static readonly ID = 'workbench.contrib.mcp.languageModelTools';

	constructor(
		@ILanguageModelToolsService private readonly _toolsService: ILanguageModelToolsService,
		@IMcpService mcpService: IMcpService,
		@IInstantiationService private readonly _instantiationService: IInstantiationService,
		@IMcpRegistry private readonly _mcpRegistry: IMcpRegistry,
		@ILifecycleService private readonly lifecycleService: ILifecycleService,
	) {
		super();

		type Rec = { source?: ToolDataSource } & IDisposable;

		// Keep tools in sync with the tools service.
		const previous = this._register(new DisposableMap<IMcpServer, Rec>());
		this._register(autorun(reader => {
			const servers = mcpService.servers.read(reader);

			const toDelete = new Set(previous.keys());
			for (const server of servers) {
				const previousRec = previous.get(server);
				if (previousRec) {
					toDelete.delete(server);
					if (!previousRec.source || equals(previousRec.source, mcpServerToSourceData(server, reader))) {
						continue; // same definition, no need to update
					}

					previousRec.dispose();
				}

				const store = new DisposableStore();
				const rec: Rec = { dispose: () => store.dispose() };
				const toolSet = new Lazy(() => {
					const source = rec.source = mcpServerToSourceData(server);
					const referenceName = server.definition.label.toLowerCase().replace(/\s+/g, '-'); // see issue https://github.com/microsoft/vscode/issues/278152
					const toolSet = store.add(this._toolsService.createToolSet(
						source,
						server.definition.id,
						referenceName,
						{
							icon: Codicon.mcp,
							description: localize('mcp.toolset', "{0}: All Tools", server.definition.label)
						}
					));

					return { toolSet, source };
				});

				this._syncTools(server, toolSet, store);
				previous.set(server, rec);
			}

			for (const key of toDelete) {
				previous.deleteAndDispose(key);
			}
		}));
	}

	private _syncTools(server: IMcpServer, collectionData: Lazy<{ toolSet: ToolSet; source: ToolDataSource }>, store: DisposableStore) {
		const tools = new Map</* tool ID */string, ISyncedToolData>();

		const collectionObservable = this._mcpRegistry.collections.map(collections =>
			collections.find(c => c.id === server.collection.id));

		store.add(autorun(reader => {
			const toDelete = new Set(tools.keys());

			// toRegister is deferred until deleting tools that moving a tool between
			// servers (or deleting one instance of a multi-instance server) doesn't cause an error.
			const toRegister: (() => void)[] = [];
			const registerTool = (tool: IMcpTool, toolData: IToolData, store: DisposableStore) => {
				store.add(this._toolsService.registerTool(toolData, this._instantiationService.createInstance(McpToolImplementation, tool, server)));
				store.add(collectionData.value.toolSet.addTool(toolData));
			};

			// Don't bother cleaning up tools internally during shutdown. This just costs time for no benefit.
			if (this.lifecycleService.willShutdown) {
				return;
			}

			const collection = collectionObservable.read(reader);
			if (!collection) {
				tools.forEach(t => t.store.dispose());
				tools.clear();
				return;
			}

			for (const tool of server.tools.read(reader)) {
				// Skip app-only tools - they should not be registered with the language model tools service
				if (!(tool.visibility & McpToolVisibility.Model)) {
					continue;
				}

				const existing = tools.get(tool.id);
				const icons = tool.icons.getUrl(22);
				const toolData: IToolData = {
					id: tool.id,
					source: collectionData.value.source,
					icon: icons || Codicon.tools,
					// duplicative: https://github.com/modelcontextprotocol/modelcontextprotocol/pull/813
					displayName: tool.definition.annotations?.title || tool.definition.title || tool.definition.name,
					toolReferenceName: tool.referenceName,
					modelDescription: tool.definition.description ?? '',
					userDescription: tool.definition.description ?? '',
					inputSchema: tool.definition.inputSchema,
					canBeReferencedInPrompt: true,
					alwaysDisplayInputOutput: true,
					canRequestPreApproval: !tool.definition.annotations?.readOnlyHint,
					canRequestPostApproval: !!tool.definition.annotations?.openWorldHint,
					runsInWorkspace: collection?.scope === StorageScope.WORKSPACE || !!collection?.remoteAuthority,
					tags: ['mcp'],
				};

				if (existing) {
					if (!equals(existing.toolData, toolData)) {
						existing.toolData = toolData;
						existing.store.clear();
						// We need to re-register both the data and implementation, as the
						// implementation is discarded when the data is removed (#245921)
						registerTool(tool, toolData, existing.store);
					}
					toDelete.delete(tool.id);
				} else {
					const store = new DisposableStore();
					toRegister.push(() => registerTool(tool, toolData, store));
					tools.set(tool.id, { toolData, store });
				}
			}

			for (const id of toDelete) {
				const tool = tools.get(id);
				if (tool) {
					tool.store.dispose();
					tools.delete(id);
				}
			}

			for (const fn of toRegister) {
				fn();
			}

			// Important: flush tool updates when the server is fully registered so that
			// any consuming (e.g. autostarting) requests have the tools available immediately.
			this._toolsService.flushToolUpdates();
		}));

		store.add(toDisposable(() => {
			for (const tool of tools.values()) {
				tool.store.dispose();
			}
		}));
	}
}

class McpToolImplementation implements IToolImpl {
	constructor(
		private readonly _tool: IMcpTool,
		private readonly _server: IMcpServer,
		@IConfigurationService private readonly _configurationService: IConfigurationService,
		@IProductService private readonly _productService: IProductService,
		@IFileService private readonly _fileService: IFileService,
		@IImageResizeService private readonly _imageResizeService: IImageResizeService,
	) { }

	async prepareToolInvocation(context: IToolInvocationPreparationContext): Promise<IPreparedToolInvocation> {
		const tool = this._tool;
		const server = this._server;
		// ToDO: need to be revisited as the first tool invocation doesnt have sandbox info and we are optimistically assuming it is not sandboxed. We should ideally have the sandbox info.
		const sandboxEnabled = await McpServer.callOn(server, async (_handler, connection) => {
			return connection.definition.sandboxEnabled;
		});
		const isSandboxedServer = sandboxEnabled === true;

		const mcpToolWarning = localize(
			'mcp.tool.warning',
			"Note that MCP servers or malicious conversation content may attempt to misuse '{0}' through tools.",
			this._productService.nameShort
		);

		// duplicative: https://github.com/modelcontextprotocol/modelcontextprotocol/pull/813
		const title = tool.definition.annotations?.title || tool.definition.title || ('`' + tool.definition.name + '`');

		let confirm: IToolConfirmationMessages | undefined;
		if (!isSandboxedServer) {
			confirm = {};
			if (!tool.definition.annotations?.readOnlyHint) {
				confirm.title = new MarkdownString(localize('msg.title', "Run {0}", title));
				confirm.message = new MarkdownString(tool.definition.description, { supportThemeIcons: true });
				confirm.disclaimer = mcpToolWarning;
				confirm.allowAutoConfirm = true;
			}
			if (tool.definition.annotations?.openWorldHint) {
				confirm.confirmResults = true;
			}
		}

		const mcpUiEnabled = this._configurationService.getValue<boolean>(mcpAppsEnabledConfig);

		return {
			confirmationMessages: confirm,
			invocationMessage: new MarkdownString(localize('msg.run', "Running {0}", title)),
			pastTenseMessage: new MarkdownString(localize('msg.ran', "Ran {0} ", title)),
			originMessage: localize('msg.subtitle', "{0} (MCP Server)", server.definition.label),
			toolSpecificData: {
				kind: 'input',
				rawInput: context.parameters,
				mcpAppData: mcpUiEnabled && tool.uiResourceUri ? {
					resourceUri: tool.uiResourceUri,
					serverDefinitionId: server.definition.id,
					collectionId: server.collection.id,
				} : undefined,
			}
		};
	}

	async invoke(invocation: IToolInvocation, _countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken) {

		const result: IToolResult = {
			content: []
		};

		const callResult = await this._tool.callWithProgress(invocation.parameters as Record<string, unknown>, progress, { chatRequestId: invocation.chatRequestId, chatSessionResource: invocation.context?.sessionResource }, token);
		const details: Mutable<IToolResultInputOutputDetails> = {
			input: JSON.stringify(invocation.parameters, undefined, 2),
			output: [],
			isError: callResult.isError === true,
		};

		for (const item of callResult.content) {
			const audience = item.annotations?.audience?.map(a => {
				if (a === 'assistant') {
					return LanguageModelPartAudience.Assistant;
				} else if (a === 'user') {
					return LanguageModelPartAudience.User;
				} else {
					return undefined;
				}
			}).filter(isDefined);

			// Explicit user parts get pushed to progress to show in the status UI
			if (audience?.includes(LanguageModelPartAudience.User)) {
				if (item.type === 'text') {
					progress.report({ message: item.text });
				}
			}

			// Rewrite image resources to images so they are inlined nicely
			const addAsInlineData = async (mimeType: string, value: string, uri?: URI): Promise<VSBuffer | void> => {
				details.output.push({ type: 'embed', mimeType, value, uri, audience });
				if (isForModel) {
					let finalData: VSBuffer;
					try {
						const resized = await this._imageResizeService.resizeImage(decodeBase64(value).buffer, mimeType);
						finalData = VSBuffer.wrap(resized);
					} catch {
						finalData = decodeBase64(value);
					}
					result.content.push({ kind: 'data', value: { mimeType, data: finalData }, audience });
				}
			};

			const addAsLinkedResource = (uri: URI, mimeType?: string) => {
				const json: IMcpToolResourceLinkContents = { uri, underlyingMimeType: mimeType };
				result.content.push({
					kind: 'data',
					audience,
					value: {
						mimeType: McpToolResourceLinkMimeType,
						data: VSBuffer.fromString(JSON.stringify(json)),
					},
				});
			};

			const isForModel = !audience || audience.includes(LanguageModelPartAudience.Assistant);
			if (item.type === 'text') {
				details.output.push({ type: 'embed', isText: true, value: item.text });
				// structured content 'represents the result of the tool call', so take
				// that in place of any textual description when present.
				if (isForModel && !callResult.structuredContent) {
					result.content.push({
						kind: 'text',
						audience,
						value: item.text
					});
				}
			} else if (item.type === 'image' || item.type === 'audio') {
				// default to some image type if not given to hint
				await addAsInlineData(item.mimeType || 'image/png', item.data);
			} else if (item.type === 'resource_link') {
				const uri = McpResourceURI.fromServer(this._server.definition, item.uri);
				details.output.push({
					type: 'ref',
					uri,
					audience,
					mimeType: item.mimeType,
				});

				if (isForModel) {
					if (item.mimeType && getAttachableImageExtension(item.mimeType)) {
						result.content.push({
							kind: 'data',
							audience,
							value: {
								mimeType: item.mimeType,
								data: await this._fileService.readFile(uri).then(f => f.value).catch(() => VSBuffer.alloc(0)),
							}
						});
					} else {
						addAsLinkedResource(uri, item.mimeType);
					}
				}
			} else if (item.type === 'resource') {
				const uri = McpResourceURI.fromServer(this._server.definition, item.resource.uri);
				if (item.resource.mimeType && getAttachableImageExtension(item.resource.mimeType) && 'blob' in item.resource) {
					await addAsInlineData(item.resource.mimeType, item.resource.blob, uri);
				} else {
					details.output.push({
						type: 'embed',
						uri,
						isText: 'text' in item.resource,
						mimeType: item.resource.mimeType,
						value: 'blob' in item.resource ? item.resource.blob : item.resource.text,
						audience,
						asResource: true,
					});

					if (isForModel) {
						const permalink = invocation.context && ChatResponseResource.createUri(invocation.context.sessionResource, invocation.callId, result.content.length, basename(uri));
						addAsLinkedResource(permalink || uri, item.resource.mimeType);
					}
				}
			}
		}

		if (callResult.structuredContent) {
			details.output.push({ type: 'embed', isText: true, value: JSON.stringify(callResult.structuredContent, null, 2), audience: [LanguageModelPartAudience.Assistant] });
			result.content.push({ kind: 'text', value: JSON.stringify(callResult.structuredContent), audience: [LanguageModelPartAudience.Assistant] });
		}

		// Add raw MCP output for MCP App UI rendering if this tool has UI
		if (this._tool.uiResourceUri) {
			details.mcpOutput = callResult;
		}

		result.toolResultDetails = details;
		return result;
	}

}
