/**
 * connection-manager.ts — MCP connection lifecycle and tool cache.
 *
 * Purpose: manage lazy MCP server connections with per-session lifetime,
 * expose a typed connection cache, and coordinate tool discovery.
 *
 * Consumer: mcp-client extension tools (mcp_discover, mcp_call),
 * lifecycle hooks (session_shutdown, session_switch).
 */
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { buildHttpTransportOpts } from "./auth.js";
import {
	getServerConfig,
	type McpServerConfig,
	readMcpConfigs,
} from "./config.js";

export type { AuthConfig, HttpTransportOptions } from "./auth.js";
export type { McpServerConfig, OAuthConfig } from "./config.js";

export interface McpToolDefinition {
	name: string;
	description: string;
	inputSchema?: unknown;
}

export interface ConnectionStatus {
	connected: boolean;
	tools: string[];
	error?: string;
}

export interface RegisterToolParams {
	name: string;
	label: string;
	description: string;
	inputSchemaRaw: unknown;
	execute: (
		id: string,
		params: Record<string, unknown>,
		signal?: AbortSignal,
	) => Promise<{
		content: { type: "text"; text: string }[];
		details: Record<string, unknown>;
	}>;
}

export type RegisterToolFn = (params: RegisterToolParams) => void;

const SAFE_CHILD_ENV_KEYS = new Set([
	"PATH",
	"HOME",
	"USER",
	"LOGNAME",
	"SHELL",
	"LANG",
	"LC_ALL",
	"LC_CTYPE",
	"LC_MESSAGES",
	"LC_NUMERIC",
	"LC_TIME",
	"TMPDIR",
	"TMP",
	"TEMP",
	"TZ",
	"TERM",
	"COLORTERM",
]);

interface ActiveConnection {
	client: Client;
	transport: Transport;
}

/**
 * Manages MCP client connections for a single agent session.
 *
 * Each instance holds its own connection map, config cache, and tool cache.
 * Call closeAll() / disconnectAll() on session shutdown or switch.
 *
 * Purpose: eliminate module-level mutable state so multiple concurrent
 * sessions each get an isolated connection pool.
 *
 * Consumer: mcp-client extension (one instance per extension activation).
 */
export class McpConnectionManager {
	private readonly connections = new Map<string, ActiveConnection>();
	private configCache: McpServerConfig[] | null = null;
	private readonly autoRegisteredServers = new Set<string>();
	private readonly toolCache = new Map<string, McpToolDefinition[]>();

	/** Read (and cache) the full ordered list of configured MCP servers. */
	readConfigs(): McpServerConfig[] {
		if (this.configCache) return this.configCache;
		this.configCache = readMcpConfigs();
		return this.configCache;
	}

	/** Invalidate the config cache so the next readConfigs() re-reads from disk. */
	invalidateConfigCache(): void {
		this.configCache = null;
	}

	/** Return the config for a single server by name (case-insensitive fallback). */
	getServerConfig(name: string): McpServerConfig | undefined {
		return getServerConfig(name, this.readConfigs());
	}

	/**
	 * Return true if the given server name is currently connected.
	 *
	 * Purpose: let the mcp_servers tool show live connection status without
	 * triggering a connection.
	 *
	 * Consumer: formatServerList in the extension wrapper.
	 */
	isConnected(name: string): boolean {
		return this.connections.has(name);
	}

	/** Return cached tools for a server, or undefined if not yet discovered. */
	getCachedTools(serverName: string): McpToolDefinition[] | undefined {
		return this.toolCache.get(serverName);
	}

	/** Store discovered tools in the cache. */
	setCachedTools(serverName: string, tools: McpToolDefinition[]): void {
		this.toolCache.set(serverName, tools);
	}

	/**
	 * Return a live MCP Client for the named server, connecting lazily on first call.
	 *
	 * Purpose: provide a single canonical connect path so every tool (discover,
	 * call, auto-registered) always gets the same cached client.
	 *
	 * Consumer: mcp_discover, mcp_call, registerToolsForServer execute functions.
	 */
	async getOrConnect(name: string, signal?: AbortSignal): Promise<Client> {
		const config = this.getServerConfig(name);
		if (!config) {
			throw new Error(
				`Unknown MCP server: "${name}". Use mcp_servers to list available servers.`,
			);
		}
		const existing = this.connections.get(config.name);
		if (existing) return existing.client;

		const client = new Client({ name: "sf", version: "1.0.0" });
		let transport: Transport;

		if (config.transport === "stdio" && config.command) {
			transport = new StdioClientTransport({
				command: config.command,
				args: config.args,
				env: this.buildChildEnv(config.env),
				cwd: config.cwd,
				stderr: "pipe",
			});
		} else if (config.transport === "http" && config.url) {
			const resolvedUrl = config.url.replace(
				/\$\{([^}]+)\}/g,
				(_: string, varName: string) => process.env[varName] ?? "",
			);
			const httpOpts = buildHttpTransportOpts({
				headers: config.headers,
				oauth: config.oauth,
			});
			transport = new StreamableHTTPClientTransport(
				new URL(resolvedUrl),
				httpOpts,
			);
		} else {
			throw new Error(
				`Server "${config.name}" has unsupported transport: ${config.transport}`,
			);
		}

		try {
			await client.connect(transport, { signal, timeout: 30000 });
		} catch (err) {
			try {
				await transport.close();
			} catch {
				/* best-effort */
			}
			try {
				await client.close();
			} catch {
				/* best-effort */
			}
			throw err;
		}

		this.connections.set(config.name, { client, transport });
		return client;
	}

	/**
	 * Close all active connections and clear the tool cache.
	 *
	 * Purpose: ensure clean shutdown on session end so no dangling stdio child
	 * processes or HTTP keep-alive connections survive the session.
	 *
	 * Consumer: session_shutdown and session_switch lifecycle hooks.
	 */
	async closeAll(): Promise<void> {
		const closing = Array.from(this.connections.entries()).map(
			async ([name, conn]) => {
				try {
					await conn.transport.close();
				} catch {
					/* best-effort */
				}
				try {
					await conn.client.close();
				} catch {
					/* best-effort */
				}
				this.connections.delete(name);
			},
		);
		await Promise.allSettled(closing);
		this.toolCache.clear();
		this.autoRegisteredServers.clear();
	}

	/**
	 * Alias for closeAll — named for the /mcp reload command surface.
	 *
	 * Purpose: allow reload command to disconnect all servers so the next
	 * mcp_discover or mcp_call lazily reconnects with fresh config.
	 *
	 * Consumer: /mcp reload command handler.
	 */
	async disconnectAll(): Promise<void> {
		await this.closeAll();
	}

	/**
	 * Return the live connection status for a named server.
	 * Safe to call even when the server has never been connected.
	 *
	 * Purpose: provide non-destructive status inspection for the status command.
	 *
	 * Consumer: /mcp status command handler.
	 */
	getConnectionStatus(name: string): ConnectionStatus {
		const conn = this.connections.get(name);
		const cached = this.toolCache.get(name);
		return {
			connected: !!conn,
			tools: cached ? cached.map((t) => t.name) : [],
			error: undefined,
		};
	}

	/**
	 * Register MCP tools discovered for a server as first-class agent tools.
	 * No-op if the server's tools were already registered in this session.
	 *
	 * Purpose: surface each MCP tool by its real name so the LLM can call
	 * tools directly (e.g. serena_find_symbol) without the mcp_call indirection.
	 *
	 * Consumer: mcp_discover execute handler, after listTools() succeeds.
	 *
	 * @param serverName  Canonical server name from the config.
	 * @param tools       Tool list returned by client.listTools().
	 * @param registerTool  Extension-provided callback that registers one tool.
	 */
	registerToolsForServer(
		serverName: string,
		tools: McpToolDefinition[],
		registerTool: RegisterToolFn,
	): void {
		if (this.autoRegisteredServers.has(serverName)) return;
		this.autoRegisteredServers.add(serverName);
		for (const tool of tools) {
			const piToolName = `${serverName}_${tool.name}`;
			const description =
				tool.description || `MCP tool: ${tool.name} on ${serverName}`;
			try {
				registerTool({
					name: piToolName,
					label: `${serverName}:${tool.name}`,
					description,
					inputSchemaRaw: tool.inputSchema,
					execute: async (_id, params, signal) => {
						const client = await this.getOrConnect(serverName, signal);
						const result = await client.callTool(
							{ name: tool.name, arguments: params },
							undefined,
							{ signal, timeout: 60000 },
						);
						const contentItems = result.content as {
							type: string;
							text?: string;
						}[];
						const raw = contentItems
							.map((c) =>
								c.type === "text" ? (c.text ?? "") : JSON.stringify(c),
							)
							.join("\n");
						return {
							content: [{ type: "text" as const, text: raw }],
							details: { server: serverName, tool: tool.name },
						};
					},
				});
			} catch {
				// Non-fatal — tool registration can fail if schema is unconvertible
			}
		}
	}

	/**
	 * Build a sanitised child environment for stdio transport processes.
	 *
	 * Purpose: prevent leaking host secrets or unwanted env vars into MCP
	 * server child processes by allow-listing safe keys and merging config-
	 * provided env on top.
	 *
	 * Consumer: getOrConnect() when creating a StdioClientTransport.
	 */
	buildChildEnv(configEnv?: Record<string, string>): Record<string, string> {
		const safe: Record<string, string> = {};
		for (const key of SAFE_CHILD_ENV_KEYS) {
			const val = process.env[key];
			if (val !== undefined) safe[key] = val;
		}
		return { ...safe, ...this.resolveEnv(configEnv ?? {}) };
	}

	private resolveEnv(env: Record<string, string>): Record<string, string> {
		const resolved: Record<string, string> = {};
		for (const [key, value] of Object.entries(env)) {
			if (typeof value === "string") {
				resolved[key] = value.replace(
					/\$\{([^}]+)\}/g,
					(_match: string, varName: string) => process.env[varName] ?? "",
				);
			} else {
				resolved[key] = value;
			}
		}
		return resolved;
	}
}

/** Export SAFE_CHILD_ENV_KEYS for tests that need to verify env filtering. */
export { SAFE_CHILD_ENV_KEYS };
