ax-frontend-app/src/services/webMcpService.ts

import { blockService } from './blockService';
import { documentService } from './documentService';
import { exportRecordService, listExportRecords } from './exportRecordService';
import { exportToWord } from './exportService';
import type { BlockType, CreateBlockRequest, UpdateBlockRequest } from '../types/editor';
import type { DocumentListFilters } from '../types/document';
import { useUIStore } from '../stores/uiStore';
import { finishWebMcpActivity, startWebMcpActivity } from './webMcpActivityService';
import { BLOCK_CREATED_EVENT, BLOCK_DELETED_EVENT, BLOCK_UPDATED_EVENT } from './blockService';

export interface WebMcpToolResult {
	ok: boolean;
	data?: unknown;
	error?: string;
}

export interface WebMcpToolInputSchema {
	type: 'object';
	properties: Record<string, Record<string, unknown>>;
	required?: string[];
	additionalProperties?: boolean;
}

export interface WebMcpTool {
	name: string;
	title: string;
	description: string;
	inputSchema: WebMcpToolInputSchema;
	requiresConfirmation?: boolean;
	readOnlyHint?: boolean;
	untrustedContentHint?: boolean;
	execute: (input: Record<string, unknown>) => Promise<WebMcpToolResult>;
}

interface ModelContext {
	registerTool?: (tool: Omit<WebMcpTool, 'requiresConfirmation'>) => void | Promise<void>;
	provideContext?: (context: { tools: ReadonlyArray<Omit<WebMcpTool, 'requiresConfirmation'>> }) =>
		| void
		| Promise<void>;
}

interface WebMcpNavigator extends Navigator {
	modelContext?: ModelContext;
}

export interface WebMcpRegistrationResult {
	supported: boolean;
	registeredTools: string[];
	error?: string;
}

export interface WebMcpExecutionOptions {
	confirmed?: boolean;
}

const DEFAULT_USER_ID = import.meta.env.VITE_WEBMCP_USER_ID || 'default-user';
let registrationPromise: Promise<WebMcpRegistrationResult> | undefined;

const text = (data: unknown): WebMcpToolResult => ({ ok: true, data });

const isRecord = (value: unknown): value is Record<string, unknown> =>
	typeof value === 'object' && value !== null && !Array.isArray(value);

const failed = (error: unknown): WebMcpToolResult => ({
	ok: false,
	error: error instanceof Error ? error.message : 'WebMCP 工具执行失败',
});

const requiredString = (input: Record<string, unknown>, key: string): string => {
	const value = input[key];
	if (typeof value !== 'string' || !value.trim()) {
		throw new Error(`${key} 不能为空`);
	}
	const hasUnsafeCharacter = Array.from(value).some((character) => {
		const codePoint = character.codePointAt(0) ?? 0;
		return codePoint < 32 || codePoint === 127 || character === '/' || character === '\\';
	});
	if (value.length > 128 || hasUnsafeCharacter) {
		throw new Error(`${key} 无效`);
	}
	return value.trim();
};

const requiredContent = (input: Record<string, unknown>, key: string): string | object | unknown[] => {
	const value = input[key];
	if (typeof value === 'string') {
		if (!value.trim()) throw new Error(`${key} 不能为空`);
		if (value.length > 200_000) throw new Error(`${key} 长度过长`);
		return value;
	}
	if (Array.isArray(value) || isRecord(value)) {
		if (JSON.stringify(value).length > 200_000) throw new Error(`${key} 内容过大`);
		return value;
	}
	throw new Error(`${key} 必须是字符串、对象或数组`);
};

const optionalString = (input: Record<string, unknown>, key: string): string | undefined => {
	const value = input[key];
	if (value === undefined || value === null || value === '') return undefined;
	if (typeof value !== 'string') throw new Error(`${key} 必须是字符串`);
	if (value.length > 128) throw new Error(`${key} 无效`);
	return value.trim();
};

const numberOr = (input: Record<string, unknown>, key: string, fallback: number): number => {
	const value = input[key];
	if (value === undefined) return fallback;
	if (typeof value !== 'number' || !Number.isFinite(value)) throw new Error(`${key} 必须是数字`);
	return value;
};

const run = async (operation: () => Promise<unknown>): Promise<WebMcpToolResult> => {
	try {
		return text(await operation());
	} catch (error) {
		return failed(error);
	}
};

const objectSchema = (
	properties: Record<string, Record<string, unknown>>,
	required: string[] = []
): WebMcpToolInputSchema => ({
	type: 'object',
	properties,
	...(required.length ? { required } : {}),
	additionalProperties: false,
});

const validateToolInput = (tool: WebMcpTool, input: Record<string, unknown>): string | undefined => {
	const schema = tool.inputSchema;
	for (const key of schema.required || []) {
		if (input[key] === undefined || input[key] === null || input[key] === '') return `${key} 不能为空`;
	}
	if (schema.additionalProperties === false) {
		const unknownKey = Object.keys(input).find((key) => !(key in schema.properties));
		if (unknownKey) return `不支持的参数: ${unknownKey}`;
	}
	for (const [key, definition] of Object.entries(schema.properties)) {
		const value = input[key];
		if (value === undefined || value === null) continue;
		if (definition.type === 'string' && typeof value !== 'string') return `${key} 必须是字符串`;
		if (definition.type === 'string' && typeof value === 'string') {
			if (definition.minLength !== undefined && value.length < Number(definition.minLength)) return `${key} 长度过短`;
			if (definition.maxLength !== undefined && value.length > Number(definition.maxLength)) return `${key} 长度过长`;
		}
		if (definition.type === 'number' && (typeof value !== 'number' || !Number.isFinite(value))) return `${key} 必须是数字`;
		if (definition.type === 'number' && typeof value === 'number') {
			if (definition.minimum !== undefined && value < Number(definition.minimum)) return `${key} 不能小于 ${definition.minimum}`;
			if (definition.maximum !== undefined && value > Number(definition.maximum)) return `${key} 不能大于 ${definition.maximum}`;
		}
		if (definition.enum && Array.isArray(definition.enum) && !definition.enum.includes(value)) return `${key} 参数值无效`;
	}
	return undefined;
};

const documentIdProperty = { type: 'string', description: '文档 ID' };
const blockIdProperty = { type: 'string', description: '块 ID' };

const tools: WebMcpTool[] = [
	{
		name: 'open_document',
		title: '打开文档',
		description: '在当前网站编辑器中打开指定文档。只读操作。',
		readOnlyHint: true,
		inputSchema: objectSchema({ documentId: documentIdProperty }, ['documentId']),
		execute: (input) => run(async () => {
			const documentId = requiredString(input, 'documentId');
			useUIStore.getState().openDocumentPreview(documentId);
			return { documentId, message: '文档已打开' };
		}),
	},
	{
		name: 'list_documents',
		title: '列出文档',
		description: '列出当前用户的文档。只读操作。',
		readOnlyHint: true,
		inputSchema: objectSchema({
			sessionId: { type: 'string', description: '可选，会话 ID' },
			page: { type: 'number', minimum: 1 },
			pageSize: { type: 'number', minimum: 1, maximum: 100 },
		}),
		execute: (input) =>
			run(() => {
				const filters: DocumentListFilters = {
					userId: DEFAULT_USER_ID,
					page: numberOr(input, 'page', 1),
					pageSize: numberOr(input, 'pageSize', 20),
					sessionId: optionalString(input, 'sessionId'),
					sortBy: 'updated_at',
					sortOrder: 'desc',
				};
				return documentService.list(filters);
			}),
	},
	{
		name: 'get_document',
		title: '读取文档',
		description: '读取指定文档的元数据和全部内容块。只读操作。',
		readOnlyHint: true,
		untrustedContentHint: true,
		inputSchema: objectSchema({ documentId: documentIdProperty }, ['documentId']),
		execute: (input) => run(() => documentService.get(requiredString(input, 'documentId'), { includeBlocks: true })),
	},
	{
		name: 'search_document',
		title: '搜索文档',
		description: '在文档中搜索包含关键词的内容块。只读操作。',
		readOnlyHint: true,
		untrustedContentHint: true,
		inputSchema: objectSchema(
			{ documentId: documentIdProperty, query: { type: 'string', description: '搜索关键词', maxLength: 2000 }, type: { type: 'string', maxLength: 64 } },
			['documentId', 'query']
		),
		execute: (input) =>
			run(() => blockService.searchBlocks(requiredString(input, 'documentId'), requiredString(input, 'query'), optionalString(input, 'type'))),
	},
	{
		name: 'get_block',
		title: '读取文档块',
		description: '读取指定文档块。只读操作。',
		readOnlyHint: true,
		untrustedContentHint: true,
		inputSchema: objectSchema({ documentId: documentIdProperty, blockId: blockIdProperty }, ['documentId', 'blockId']),
		execute: (input) => run(() => blockService.getBlock(requiredString(input, 'documentId'), requiredString(input, 'blockId'))),
	},
	{
		name: 'get_document_toc',
		title: '读取文档目录',
		description: '获取文档目录树。只读操作。',
		readOnlyHint: true,
		untrustedContentHint: true,
		inputSchema: objectSchema({ documentId: documentIdProperty }, ['documentId']),
		execute: (input) => run(() => blockService.getTOC(requiredString(input, 'documentId'))),
	},
	{
		name: 'get_document_stats',
		title: '读取文档统计',
		description: '获取文档内容块统计信息。只读操作。',
		readOnlyHint: true,
		inputSchema: objectSchema({ documentId: documentIdProperty }, ['documentId']),
		execute: (input) => run(() => blockService.getStats(requiredString(input, 'documentId'))),
	},
	{
		name: 'list_export_records',
		title: '列出导出记录',
		description: '列出当前用户的导出记录。只读操作。',
		readOnlyHint: true,
		inputSchema: objectSchema({ page: { type: 'number', minimum: 1 }, pageSize: { type: 'number', minimum: 1, maximum: 100 } }),
		execute: (input) =>
			run(() => listExportRecords({
				userId: DEFAULT_USER_ID,
				page: numberOr(input, 'page', 1),
				pageSize: numberOr(input, 'pageSize', 20),
				sortOrder: 'desc',
			})),
	},
	{
		name: 'insert_block',
		title: '插入内容块',
		description: '向文档插入内容块。需要用户确认。',
		requiresConfirmation: true,
		inputSchema: objectSchema(
			{
				documentId: documentIdProperty,
				type: { type: 'string', enum: ['heading', 'paragraph', 'table', 'image', 'toc'] },
				content: { type: 'string', description: '块内容', maxLength: 200000 },
				level: { type: 'number', minimum: 0, maximum: 6 },
				afterBlockId: { type: 'string', description: '插入到此块之后，可选' },
				clientBlockId: { type: 'string', description: '可选，客户端幂等块 ID' },
			},
			['documentId', 'type', 'content']
		),
		execute: (input) => run(async () => {
			const type = requiredString(input, 'type') as BlockType;
			const level = input.level === undefined ? 0 : numberOr(input, 'level', 0);
			if (type === 'heading' && (level < 1 || level > 6)) {
				throw new Error('heading 的 level 必须是 1-6');
			}
			if (type !== 'heading' && level !== 0) {
				throw new Error('非 heading 块的 level 必须是 0');
			}
			const documentId = requiredString(input, 'documentId');
			const response = await blockService.createBlock(documentId, {
				type,
				content: requiredContent(input, 'content') as CreateBlockRequest['content'],
				level,
				after_block_id: optionalString(input, 'afterBlockId'),
				client_block_id: optionalString(input, 'clientBlockId'),
			});
			const created = await blockService.getBlock(documentId, response.blockId);
			window.dispatchEvent(new CustomEvent(BLOCK_CREATED_EVENT, {
				detail: { documentId, block: created.block },
			}));
			return { ...response, block: created.block, message: 'Block created successfully' };
		}),
	},
	{
		name: 'update_block',
		title: '修改内容块',
		description: '更新文档块内容。需要用户确认。',
		requiresConfirmation: true,
		inputSchema: objectSchema({ documentId: documentIdProperty, blockId: blockIdProperty, content: { type: 'string', description: '新的块内容' } }, ['documentId', 'blockId', 'content']),
		execute: (input) => run(() => {
			const documentId = requiredString(input, 'documentId');
			const blockId = requiredString(input, 'blockId');
			const updates: UpdateBlockRequest = { content: requiredString(input, 'content') };
			return blockService.updateBlock(documentId, blockId, updates).then((result) => {
				window.dispatchEvent(new CustomEvent(BLOCK_UPDATED_EVENT, {
					detail: { documentId, blockId, updates },
				}));
				return result;
			});
		}),
	},
	{
		name: 'delete_block',
		title: '删除内容块',
		description: '删除指定文档块。需要用户确认。',
		requiresConfirmation: true,
		inputSchema: objectSchema({ documentId: documentIdProperty, blockId: blockIdProperty }, ['documentId', 'blockId']),
		execute: (input) => run(() => {
			const documentId = requiredString(input, 'documentId');
			const blockId = requiredString(input, 'blockId');
			return blockService.deleteBlock(documentId, blockId).then(() => {
				window.dispatchEvent(new CustomEvent(BLOCK_DELETED_EVENT, {
					detail: { documentId, blockId },
				}));
				return { message: 'Block deleted successfully' };
			});
		}),
	},
	{
		name: 'export_document',
		title: '导出文档',
		description: '将文档导出为 Word 文件。需要用户确认。',
		requiresConfirmation: true,
		inputSchema: objectSchema({ documentId: documentIdProperty }, ['documentId']),
		execute: (input) => run(() => exportToWord({ documentId: requiredString(input, 'documentId') })),
	},
	{
		name: 'download_export_record',
		title: '下载导出记录',
		description: '下载指定导出记录的文件。需要用户确认。',
		requiresConfirmation: true,
		inputSchema: objectSchema({ recordId: { type: 'string', description: '导出记录 ID' } }, ['recordId']),
		execute: (input) => run(async () => {
			const recordId = requiredString(input, 'recordId');
			await exportRecordService.download(recordId, DEFAULT_USER_ID);
			return { recordId, message: '下载已开始' };
		}),
	},
];

const browserTools = (): ReadonlyArray<Omit<WebMcpTool, 'requiresConfirmation'>> =>
	tools.map((tool) => ({
		name: tool.name,
		title: tool.title,
		description: tool.description,
		inputSchema: tool.inputSchema,
		readOnlyHint: tool.readOnlyHint,
		untrustedContentHint: tool.untrustedContentHint,
		execute: (input) => {
			if (tool.requiresConfirmation) {
				const activityId = startWebMcpActivity(tool.name, input);
				const result = { ok: false, error: '该 WebMCP 工具需要当前页面用户确认' };
				finishWebMcpActivity(activityId, result);
				return Promise.resolve(result);
			}
			return executeWebMcpTool(tool.name, input);
		},
	}));

export const getWebMcpTools = (): ReadonlyArray<WebMcpTool> => tools;
export const getWebMcpTool = (name: string): WebMcpTool | undefined => tools.find((tool) => tool.name === name);

export const executeWebMcpTool = async (
	name: string,
	input: unknown,
	options: WebMcpExecutionOptions = {}
): Promise<WebMcpToolResult> => {
	if (!isRecord(input)) {
		return { ok: false, error: 'WebMCP 工具参数必须是对象' };
	}
	const activityId = startWebMcpActivity(name, input);
	const tool = getWebMcpTool(name);
	if (!tool) {
		const result = { ok: false, error: `未知的 WebMCP 工具: ${name}` };
		finishWebMcpActivity(activityId, result);
		return result;
	}
	if (tool.requiresConfirmation && !options.confirmed) {
		const result = { ok: false, error: '该 WebMCP 工具需要当前页面用户确认' };
		finishWebMcpActivity(activityId, result);
		return result;
	}
	const validationError = validateToolInput(tool, input);
	if (validationError) {
		const result = { ok: false, error: validationError };
		finishWebMcpActivity(activityId, result);
		return result;
	}
	try {
		const result = await tool.execute(input);
		finishWebMcpActivity(activityId, result);
		return result;
	} catch (error) {
		const result = { ok: false, error: error instanceof Error ? error.message : '工具执行失败' };
		finishWebMcpActivity(activityId, result);
		return result;
	}
};

export const registerWebMcpTools = (): Promise<WebMcpRegistrationResult> => {
	if (registrationPromise) return registrationPromise;
	registrationPromise = (async () => {
		const modelContext = (navigator as WebMcpNavigator).modelContext;
		if (!modelContext) return { supported: false, registeredTools: [], error: '当前浏览器未提供 navigator.modelContext' };
		try {
			const definitions = browserTools();
			if (modelContext.provideContext) await modelContext.provideContext({ tools: definitions });
			else if (modelContext.registerTool) for (const tool of definitions) await modelContext.registerTool(tool);
			else return { supported: true, registeredTools: [], error: 'WebMCP API 不支持工具注册' };
			return { supported: true, registeredTools: tools.map((tool) => tool.name) };
		} catch (error) {
			console.warn('[WebMCP] 工具注册失败，编辑器将继续运行', error);
			return { supported: true, registeredTools: [], error: error instanceof Error ? error.message : '工具注册失败' };
		}
	})();
	return registrationPromise;
};

export default tools;