import {
	McpClient,
	type BuiltTool,
	type McpServerConfig as NativeMcpServerConfig,
	type McpToolCallSettledEvent as NativeMcpToolCallSettledEvent,
} from '@n8n/agents';
import type { Result } from '@n8n/utils/result';
import { UserError } from 'n8n-workflow';

import {
	addSafeMcpTools,
	createClaimedToolNames,
	isSafeMcpIdentifierName,
} from '../agent/mcp-tool-name-validation';
import type { McpToolNameValidationError } from '../agent/mcp-tool-name-validation';
import type { McpDescriptionTruncation } from '../agent/sanitize-mcp-descriptions';
import { sanitizeMcpToolSchemas } from '../agent/sanitize-mcp-schemas';
import type { McpSchemaSanitizationError } from '../agent/sanitize-mcp-schemas';
import type { Logger } from '../logger';
import { createToolRegistry, createToolRegistryFromTools } from '../tool-registry';
import type { InstanceAiToolRegistry, McpServerConfig } from '../types';

type McpToolRegistry = InstanceAiToolRegistry;

/** Per-server connection failures recorded while listing tools for one config. */
type McpConnectionFailure = { server: McpServerConfig; error: string };

/** Result of `getRegularTools`: the loaded tool registry plus any per-server
 * connection failures for the requested config. Failures travel with the call
 * result (not shared mutable state) so concurrent runs can't read each other's
 * failures. */
interface McpRegularToolsResult {
	tools: McpToolRegistry;
	connectionFailures: McpConnectionFailure[];
}

export interface McpToolCallSettledEvent {
	server: McpServerConfig;
	toolName: string;
	success: boolean;
}

export interface McpClientManagerOptions {
	onToolCallSettled?: (event: McpToolCallSettledEvent) => void;
	/**
	 * Invoked once per MCP server that fails to connect during `getRegularTools`.
	 * The server's tools are skipped; the run continues with the remaining
	 * servers' tools. Hosts surface these as non-fatal warnings.
	 */
	onConnectionFailed?: (event: { server: McpServerConfig; error: string }) => void;
}

/**
 * SSRF policy gate for outbound MCP URLs. The cli's `SsrfProtectionService`
 * satisfies this structurally; we keep the local shape narrow to avoid pulling
 * `n8n-core` into this package just for one type.
 */
export interface SsrfUrlValidator {
	validateUrl(url: string | URL): Promise<Result<void, Error>>;
}

function buildNativeMcpConfigs(
	configs: McpServerConfig[],
	requireApproval: boolean,
	onToolCallSettled?: McpClientManagerOptions['onToolCallSettled'],
): NativeMcpServerConfig[] {
	const servers: NativeMcpServerConfig[] = [];
	for (const server of configs) {
		const baseConfig = {
			name: server.name,
			toolFilter: server.toolFilter,
			requireApproval,
			...(onToolCallSettled
				? {
						onToolCallSettled: ({ toolName, success }: NativeMcpToolCallSettledEvent) =>
							onToolCallSettled({ server, toolName, success }),
					}
				: {}),
		};
		if (server.url) {
			servers.push({
				...baseConfig,
				url: server.url,
				transport: server.transport,
				fetch: server.fetch,
			});
		} else if (server.command) {
			servers.push({
				...baseConfig,
				command: server.command,
				args: server.args,
				env: server.env,
			});
		}
	}
	return servers;
}

function toolsToRegistry(tools: BuiltTool[]): McpToolRegistry {
	return createToolRegistryFromTools(tools);
}

function warnSkippedMcpSchema(logger: Logger, source: string) {
	return (error: McpSchemaSanitizationError) => {
		logger.warn('Skipped MCP tool with unsupported schema', {
			toolName: error.details.toolName,
			source,
			path: error.details.path,
			depth: error.details.depth,
			maxDepth: error.details.maxDepth,
			limitType: error.details.limitType,
			limit: error.details.limit,
			reason: error.message,
		});
	};
}

function warnTruncatedMcpDescription(logger: Logger, source: string) {
	return (truncation: McpDescriptionTruncation) => {
		logger.warn('Truncated overlong MCP description', {
			toolName: truncation.toolName,
			source,
			path: truncation.path,
			originalLength: truncation.originalLength,
			limit: truncation.limit,
		});
	};
}

function warnSkippedMcpTool(logger: Logger) {
	return (error: McpToolNameValidationError) => {
		logger.warn('Skipped MCP tool with unsafe name', {
			toolName: error.toolName,
			source: error.source,
			reason: error.message,
		});
	};
}

function getSafeMcpServers(
	configs: McpServerConfig[],
	logger: Logger,
	source: string,
): McpServerConfig[] {
	return configs.filter((config) => {
		if (isSafeMcpIdentifierName(config.name)) return true;

		logger.warn('Skipped MCP server with unsafe name', {
			serverName: config.name,
			source,
		});
		return false;
	});
}

/**
 * Owns the lifecycle of MCP client connections used by Instance AI.
 *
 * External MCP servers configured by the admin are merged into the
 * orchestrator's toolset. Tool listings are cached by config hash; clients are
 * tracked in one map so `disconnect()` can clean them up.
 */
export class McpClientManager {
	private regularToolsByKey = new Map<string, McpRegularToolsResult>();

	private inFlightRegularByKey = new Map<string, Promise<McpRegularToolsResult>>();

	private clientsByKey = new Map<string, McpClient>();

	constructor(
		private readonly ssrfValidator?: SsrfUrlValidator,
		private readonly options: McpClientManagerOptions = {},
	) {}

	async getRegularTools(
		configs: McpServerConfig[],
		logger: Logger,
		requireApproval = true,
	): Promise<McpRegularToolsResult> {
		const safeConfigs = getSafeMcpServers(configs, logger, 'external MCP');
		if (safeConfigs.length === 0) return { tools: createToolRegistry(), connectionFailures: [] };

		// Approval mode is part of the cache key: the same servers wrapped with vs
		// without an approval gate are distinct tool sets.
		const key = JSON.stringify({ configs: safeConfigs, requireApproval });
		return await this.getOrLoad(
			this.regularToolsByKey,
			this.inFlightRegularByKey,
			key,
			async () => {
				await this.validateConfigs(safeConfigs);
				return await this.connectAndListTools(
					safeConfigs,
					key,
					requireApproval,
					logger,
					'external MCP',
				);
			},
		);
	}

	async disconnect(): Promise<void> {
		const clients = [...this.clientsByKey.values()];
		this.clientsByKey.clear();
		this.regularToolsByKey.clear();
		this.inFlightRegularByKey.clear();
		await Promise.all(clients.map(async (client) => await client.close()));
	}

	private async getOrLoad<T>(
		cache: Map<string, T>,
		inFlight: Map<string, Promise<T>>,
		key: string,
		produce: () => Promise<T>,
	): Promise<T> {
		const cached = cache.get(key);
		if (cached !== undefined) return cached;

		const pending = inFlight.get(key);
		if (pending) return await pending;

		const promise = (async () => {
			const value = await produce();
			cache.set(key, value);
			return value;
		})();

		inFlight.set(key, promise);
		try {
			return await promise;
		} finally {
			inFlight.delete(key);
		}
	}

	private async validateConfigs(configs: McpServerConfig[]): Promise<void> {
		for (const server of configs) {
			if (!server.url) continue;

			let parsed: URL;
			try {
				parsed = new URL(server.url);
			} catch {
				throw new UserError(`MCP server "${server.name}": invalid URL "${server.url}"`);
			}

			if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
				throw new UserError(
					`MCP server "${server.name}": only http(s) URLs are allowed, got "${parsed.protocol}"`,
				);
			}

			if (this.ssrfValidator) {
				const result = await this.ssrfValidator.validateUrl(server.url);
				if (!result.ok) {
					throw new UserError(
						`MCP server "${server.name}": URL blocked by SSRF policy - ${result.error.message}`,
					);
				}
			}
		}
	}

	private async connectAndListTools(
		configs: McpServerConfig[],
		clientKey: string,
		requireApproval: boolean,
		logger: Logger,
		source: string,
	): Promise<McpRegularToolsResult> {
		const client = new McpClient(
			buildNativeMcpConfigs(configs, requireApproval, this.options.onToolCallSettled),
		);
		this.clientsByKey.set(clientKey, client);

		const registry = toolsToRegistry(await client.listTools());
		// The SDK client is the single source of truth for per-server connection
		// failures. Map each back to the originating host config and surface it
		// via the manager-level observer + logger so a single misconfigured or
		// unhealthy MCP server surfaces as a non-fatal warning instead of
		// aborting the run. Failures are returned with the result (not stored on
		// a shared field) so concurrent runs with different configs can't read
		// each other's failures.
		const connectionFailures: McpConnectionFailure[] = [...client.getConnectionFailures()].map(
			(f) => {
				const server =
					configs.find((c) => c.name === f.server) ??
					({
						name: f.server,
					} as McpServerConfig);
				return { server, error: f.error };
			},
		);
		for (const failure of connectionFailures) {
			logger.warn('Skipped MCP server that failed to connect', {
				serverName: failure.server.name,
				source,
				error: failure.error,
			});
			this.options.onConnectionFailed?.(failure);
		}

		const sanitizedTools = sanitizeMcpToolSchemas(registry, {
			onError: warnSkippedMcpSchema(logger, source),
			onDescriptionTruncated: warnTruncatedMcpDescription(logger, source),
		});

		const safeTools: McpToolRegistry = createToolRegistry();
		addSafeMcpTools(safeTools, sanitizedTools, {
			source,
			claimedToolNames: createClaimedToolNames([]),
			warn: warnSkippedMcpTool(logger),
		});
		return { tools: safeTools, connectionFailures };
	}
}