/**
 * MCP Client Extension — Native MCP server integration for pi
 *
 * Provides on-demand access to MCP servers configured in project files
 * (.mcp.json, .sf/mcp.json) using the @modelcontextprotocol/sdk Client
 * directly — no external CLI dependency required.
 *
 * Three tools:
 *   mcp_servers   — List available MCP servers from config files
 *   mcp_discover  — Get tool signatures for a specific server (lazy connect)
 *   mcp_call      — Call a tool on an MCP server (lazy connect)
 */

import { existsSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { Client } from "@modelcontextprotocol/sdk/client";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { Type } from "@sinclair/typebox";
import type { ExtensionAPI } from "@singularity-forge/pi-coding-agent";
import {
	DEFAULT_MAX_BYTES,
	DEFAULT_MAX_LINES,
	formatSize,
	truncateHead,
} from "@singularity-forge/pi-coding-agent";
import { Text } from "@singularity-forge/pi-tui";
import type { McpHttpAuthConfig } from "./auth.js";
import { buildHttpTransportOpts } from "./auth.js";

// ─── Types ────────────────────────────────────────────────────────────────────

interface McpServerConfig {
	name: string;
	transport: "stdio" | "http" | "unknown";
	command?: string;
	args?: string[];
	env?: Record<string, string>;
	url?: string;
	cwd?: string;
	/** Static headers for HTTP transport (supports ${VAR} env resolution). */
	headers?: Record<string, string>;
	/** OAuth config for HTTP transport. */
	oauth?: McpHttpAuthConfig["oauth"];
}

interface McpToolSchema {
	name: string;
	description: string;
	inputSchema?: Record<string, unknown>;
}

interface ManagedConnection {
	client: Client;
	transport: StdioClientTransport | StreamableHTTPClientTransport;
}

// ─── Connection Manager ───────────────────────────────────────────────────────

const connections = new Map<string, ManagedConnection>();
let configCache: McpServerConfig[] | null = null;
/** Servers whose MCP tools have been auto-registered as first-class pi tools. */
const autoRegisteredServers = new Set<string>();
const toolCache = new Map<string, McpToolSchema[]>();

function readConfigs(): McpServerConfig[] {
	if (configCache) return configCache;

	const servers: McpServerConfig[] = [];
	const seen = new Set<string>();
	// Search order matters: first hit wins (seen-guard below), so put
	// project-local configs first — a project can override or shadow a
	// globally-registered server by re-declaring the same name.
	const sfHome = process.env.SF_HOME || join(homedir(), ".sf");
	const configPaths = [
		join(process.cwd(), ".mcp.json"),
		join(process.cwd(), ".sf", "mcp.json"),
		join(sfHome, "mcp.json"), // global: ~/.sf/mcp.json
		join(sfHome, "agent", "mcp.json"), // global: ~/.sf/agent/mcp.json (legacy alt)
		join(homedir(), ".mcp.json"), // user-global: ~/.mcp.json (Claude Code, npx, etc.)
	];

	for (const configPath of configPaths) {
		try {
			if (!existsSync(configPath)) continue;
			const raw = readFileSync(configPath, "utf-8");
			const data = JSON.parse(raw) as Record<string, unknown>;
			const mcpServers = (data.mcpServers ?? data.servers) as
				| Record<string, Record<string, unknown>>
				| undefined;
			if (!mcpServers || typeof mcpServers !== "object") continue;

			for (const [name, config] of Object.entries(mcpServers)) {
				if (seen.has(name)) continue;
				seen.add(name);

				const hasCommand = typeof config.command === "string";
				const hasUrl = typeof config.url === "string";
				const transport: McpServerConfig["transport"] = hasCommand
					? "stdio"
					: hasUrl
						? "http"
						: "unknown";

				const hasHeaders =
					hasUrl && config.headers && typeof config.headers === "object";
				const hasOAuth =
					hasUrl && config.oauth && typeof config.oauth === "object";

				servers.push({
					name,
					transport,
					...(hasCommand && {
						command: config.command as string,
						args: Array.isArray(config.args)
							? (config.args as string[])
							: undefined,
						env:
							config.env && typeof config.env === "object"
								? (config.env as Record<string, string>)
								: undefined,
						cwd: typeof config.cwd === "string" ? config.cwd : undefined,
					}),
					...(hasUrl && { url: config.url as string }),
					headers: hasHeaders
						? (config.headers as Record<string, string>)
						: undefined,
					oauth: hasOAuth
						? (config.oauth as McpHttpAuthConfig["oauth"])
						: undefined,
				});
			}
		} catch {
			// Non-fatal — config file may not exist or be malformed
		}
	}

	configCache = servers;
	return servers;
}

function getServerConfig(name: string): McpServerConfig | undefined {
	const trimmed = name.trim();
	return readConfigs().find(
		(s) => s.name === trimmed || s.name.toLowerCase() === trimmed.toLowerCase(),
	);
}

/** Resolve ${VAR} references in env values against process.env. */
function 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, varName) => process.env[varName] ?? "",
			);
		} else {
			resolved[key] = value;
		}
	}
	return resolved;
}

// ─── JSON Schema → TypeBox converter ─────────────────────────────────────────

// eslint-disable-next-line @typescript-eslint/no-explicit-any
function jsonSchemaPropToTypeBox(schema: Record<string, unknown>): any {
	if (!schema || typeof schema !== "object") return Type.Any();
	const t = schema.type as string;
	if (t === "string") return Type.String({ description: schema.description as string | undefined });
	if (t === "number" || t === "integer") return Type.Number({ description: schema.description as string | undefined });
	if (t === "boolean") return Type.Boolean({ description: schema.description as string | undefined });
	if (t === "array") return Type.Array(Type.Any());
	if (t === "object") {
		const props = schema.properties as Record<string, Record<string, unknown>> | undefined;
		if (props) {
			const entries: Record<string, unknown> = {};
			for (const [k, v] of Object.entries(props)) {
				entries[k] = jsonSchemaPropToTypeBox(v);
			}
			return Type.Object(entries as Parameters<typeof Type.Object>[0]);
		}
	}
	return Type.Any();
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
function jsonSchemaToTypeBox(schema: Record<string, unknown> | undefined): any {
	if (!schema || typeof schema !== "object") return Type.Object({});
	const obj = schema as Record<string, unknown>;
	const props = obj.properties as Record<string, Record<string, unknown>> | undefined;
	if (!props) return Type.Object({});
	const entries: Record<string, unknown> = {};
	for (const [k, v] of Object.entries(props)) {
		entries[k] = jsonSchemaPropToTypeBox(v);
	}
	return Type.Object(entries as Parameters<typeof Type.Object>[0]);
}

// ─── Dynamic MCP tool auto-registration ───────────────────────────────────────

function registerMcpToolsForServer(pi: ExtensionAPI, serverName: string, tools: McpToolSchema[]) {
	if (autoRegisteredServers.has(serverName)) return;
	autoRegisteredServers.add(serverName);

	for (const tool of tools) {
		const piToolName = `${serverName}_${tool.name}`;
		const description = tool.description || `MCP tool: ${tool.name} on ${serverName}`;
		// Build parameter TypeBox type from MCP inputSchema
		const paramType = tool.inputSchema
			? jsonSchemaToTypeBox(tool.inputSchema)
			: Type.Object({});

		try {
			pi.registerTool({
				name: piToolName,
				label: `${serverName}:${tool.name}`,
				description,
				parameters: paramType,
				async execute(_id, params) {
					// Delegate to the internal mcp_call logic directly via the client
					const client = await getOrConnect(serverName);
					const result = await client.callTool(
						{ name: tool.name, arguments: params as Record<string, unknown> },
						undefined,
						{ timeout: 60000 },
					);
					const contentItems = result.content as Array<{ type: string; text?: string }>;
					const raw = contentItems
						.map((c) => (c.type === "text" ? (c.text ?? "") : JSON.stringify(c)))
						.join("\n");
					const truncation = truncateHead(raw, {
						maxLines: DEFAULT_MAX_LINES,
						maxBytes: DEFAULT_MAX_BYTES,
					});
					let finalText = truncation.content;
					if (truncation.truncated) {
						finalText += `\n\n[Output truncated: ${truncation.outputLines}/${truncation.totalLines} lines]`;
					}
					return {
						content: [{ type: "text", text: finalText }],
						details: { server: serverName, tool: tool.name },
					};
				},
			});
		}
		catch {
			// Non-fatal — tool registration can fail if schema is unconvertible
		}
	}
}

async function getOrConnect(
	name: string,
	signal?: AbortSignal,
): Promise<Client> {
	const config = getServerConfig(name);
	if (!config)
		throw new Error(
			`Unknown MCP server: "${name}". Use mcp_servers to list available servers.`,
		);

	// Always use config.name as the canonical cache key so that variant
	// casing / whitespace still hits the same connection.
	const existing = connections.get(config.name);
	if (existing) return existing.client;

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

	if (config.transport === "stdio" && config.command) {
		transport = new StdioClientTransport({
			command: config.command,
			args: config.args,
			env: config.env
				? ({ ...process.env, ...resolveEnv(config.env) } as Record<
						string,
						string
					>)
				: undefined,
			cwd: config.cwd,
			stderr: "pipe",
		});
	} else if (config.transport === "http" && config.url) {
		const resolvedUrl = config.url.replace(
			/\$\{([^}]+)\}/g,
			(_, varName) => 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}`,
		);
	}

	await client.connect(transport, { signal, timeout: 30000 });
	connections.set(config.name, { client, transport });
	return client;
}

async function closeAll(): Promise<void> {
	const closing = Array.from(connections.entries()).map(
		async ([name, conn]) => {
			try {
				await conn.client.close();
			} catch {
				// Best-effort cleanup
			}
			connections.delete(name);
		},
	);
	await Promise.allSettled(closing);
	toolCache.clear();
}

// ─── Formatters ───────────────────────────────────────────────────────────────

function formatServerList(servers: McpServerConfig[]): string {
	if (servers.length === 0)
		return "No MCP servers configured. Add servers to .mcp.json or .sf/mcp.json.";

	const lines: string[] = [`${servers.length} MCP servers configured:\n`];

	for (const s of servers) {
		const connected = connections.has(s.name) ? "✓" : "○";
		const cached = toolCache.get(s.name);
		const toolCount = cached ? ` — ${cached.length} tools` : "";
		lines.push(`${connected} ${s.name} (${s.transport})${toolCount}`);
	}

	lines.push(
		"\nUse mcp_discover to see full tool schemas for a specific server.",
	);
	lines.push("Use mcp_call to invoke a tool: mcp_call(server, tool, args).");
	return lines.join("\n");
}

function formatToolList(serverName: string, tools: McpToolSchema[]): string {
	const lines: string[] = [`${serverName} — ${tools.length} tools:\n`];

	for (const tool of tools) {
		lines.push(`## ${tool.name}`);
		if (tool.description) lines.push(tool.description);
		if (tool.inputSchema) {
			lines.push("```json");
			lines.push(JSON.stringify(tool.inputSchema, null, 2));
			lines.push("```");
		}
		lines.push("");
	}

	lines.push(
		`Call with: mcp_call(server="${serverName}", tool="<tool_name>", args={...})`,
	);
	return lines.join("\n");
}

// ─── Status helper (consumed by /sf mcp) ─────────────────────────────────────

/**
 * Return the live connection status for a named MCP server.
 * Safe to call even when the server has never been connected.
 */
export function getConnectionStatus(name: string): {
	connected: boolean;
	tools: string[];
	error?: string;
} {
	const conn = connections.get(name);
	const cached = toolCache.get(name);
	return {
		connected: !!conn,
		tools: cached ? cached.map((t) => t.name) : [],
		error: undefined,
	};
}

// ─── Test-exported helpers ────────────────────────────────────────────────────

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",
]);

export function _buildMcpChildEnvForTest(
	env: Record<string, string>,
): Record<string, string> {
	const safe: Record<string, string> = {};
	for (const key of SAFE_CHILD_ENV_KEYS) {
		if (process.env[key] !== undefined) safe[key] = process.env[key]!;
	}
	return { ...safe, ...resolveEnv(env) };
}

export function _buildMcpTrustConfirmOptionsForTest(signal: AbortSignal): {
	timeout: number;
	signal: AbortSignal;
} {
	return { timeout: 120_000, signal };
}

// ─── Extension ────────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	// ── mcp_servers ──────────────────────────────────────────────────────────

	pi.registerTool({
		name: "mcp_servers",
		label: "MCP Servers",
		description:
			"List all available MCP servers configured in project files (.mcp.json, .sf/mcp.json). " +
			"Shows server names, transport type, and connection status. After mcp_discover, each server's " +
			"tools are auto-registered as first-class pi tools (e.g. serena_find_symbol).",
		promptSnippet: "List available MCP servers from project configuration",
		promptGuidelines: [
			"Call mcp_servers to see what MCP servers are available before trying to use one.",
			"After mcp_discover(server), the server's tools appear as real pi tools.",
			"MCP servers provide external integrations (Twitter, Linear, Railway, etc.) via the Model Context Protocol.",
			"After listing, use mcp_discover(server) to get tool schemas, then mcp_call(server, tool, args) to invoke.",
		],
		parameters: Type.Object({
			refresh: Type.Optional(
				Type.Boolean({
					description: "Force refresh the server list (default: use cache)",
				}),
			),
		}),

		async execute(_id, params) {
			if (params.refresh) configCache = null;

			const servers = readConfigs();
			return {
				content: [{ type: "text", text: formatServerList(servers) }],
				details: {
					serverCount: servers.length,
					cached: !params.refresh && configCache !== null,
				},
			};
		},

		renderCall(args, theme) {
			let text = theme.fg("toolTitle", theme.bold("mcp_servers"));
			if (args.refresh) text += theme.fg("warning", " (refresh)");
			return new Text(text, 0, 0);
		},

		renderResult(result, { isPartial }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Reading MCP config..."), 0, 0);
			const d = result.details as { serverCount: number } | undefined;
			return new Text(
				theme.fg("success", `${d?.serverCount ?? 0} servers configured`),
				0,
				0,
			);
		},
	});

	// ── mcp_discover ─────────────────────────────────────────────────────────

	pi.registerTool({
		name: "mcp_discover",
		label: "MCP Discover",
		description:
			"Get detailed tool signatures and JSON schemas for a specific MCP server. " +
			"Connects to the server on first call (lazy connection). " +
			"After discovery, each MCP tool is auto-registered as a first-class pi tool " +
			"(e.g. serena_find_symbol) — the LLM can call them directly without mcp_call.",
		promptSnippet:
			"Discover MCP server tools and register them as first-class pi tools",
		promptGuidelines: [
			"Call mcp_discover(server) to connect to an MCP server and surface its tools.",
			"After discovery, the LLM sees each tool by its real name (e.g. serena_search_for_pattern).",
			"Call tools directly by their names instead of going through mcp_call.",
		],
		parameters: Type.Object({
			server: Type.String({
				description:
					"MCP server name (from mcp_servers output), e.g. 'railway', 'twitter-mcp', 'linear'",
			}),
		}),

		async execute(_id, params, signal) {
			try {
				// Return cached tools if available
				const cached = toolCache.get(params.server);
				if (cached) {
					const text = formatToolList(params.server, cached);
					const truncation = truncateHead(text, {
						maxLines: DEFAULT_MAX_LINES,
						maxBytes: DEFAULT_MAX_BYTES,
					});
					let finalText = truncation.content;
					if (truncation.truncated) {
						finalText += `\n\n[Truncated: ${truncation.outputLines}/${truncation.totalLines} lines (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`;
					}
					return {
						content: [{ type: "text", text: finalText }],
						details: {
							server: params.server,
							toolCount: cached.length,
							cached: true,
						},
					};
				}

				const client = await getOrConnect(params.server, signal);
				const result = await client.listTools(undefined, {
					signal,
					timeout: 30000,
				});
				const tools: McpToolSchema[] = (result.tools ?? []).map((t) => ({
					name: t.name,
					description: t.description ?? "",
					inputSchema: t.inputSchema as Record<string, unknown> | undefined,
				}));
				toolCache.set(params.server, tools);

				// Auto-register each MCP tool as a first-class pi tool.
				// After this, the LLM sees e.g. serena_find_symbol directly instead
				// of going through the generic mcp_call indirection.
				registerMcpToolsForServer(pi, params.server, tools);

				const text = formatToolList(params.server, tools);
				const truncation = truncateHead(text, {
					maxLines: DEFAULT_MAX_LINES,
					maxBytes: DEFAULT_MAX_BYTES,
				});
				let finalText = truncation.content;
				if (truncation.truncated) {
					finalText += `\n\n[Truncated: ${truncation.outputLines}/${truncation.totalLines} lines (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`;
				}

				return {
					content: [{ type: "text", text: finalText }],
					details: {
						server: params.server,
						toolCount: tools.length,
						cached: false,
					},
				};
			} catch (err: unknown) {
				const msg = err instanceof Error ? err.message : String(err);
				throw new Error(
					`Failed to discover tools for "${params.server}": ${msg}`,
				);
			}
		},

		renderCall(args, theme) {
			let text = theme.fg("toolTitle", theme.bold("mcp_discover "));
			text += theme.fg("accent", args.server);
			return new Text(text, 0, 0);
		},

		renderResult(result, { isPartial }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Discovering tools..."), 0, 0);
			const d = result.details as
				| { server: string; toolCount: number }
				| undefined;
			return new Text(
				theme.fg("success", `${d?.toolCount ?? 0} tools`) +
					theme.fg("dim", ` · ${d?.server}`),
				0,
				0,
			);
		},
	});

	// ── mcp_call ─────────────────────────────────────────────────────────────

	pi.registerTool({
		name: "mcp_call",
		label: "MCP Call",
		description:
			"Call a tool on an MCP server. Provide the server name, tool name, and arguments. " +
			"Connects to the server on first call (lazy connection). " +
			"Use mcp_discover first to see available tools and their required arguments.",
		promptSnippet: "Call a tool on an MCP server",
		promptGuidelines: [
			"Always use mcp_discover first to understand the tool's parameters before calling mcp_call.",
			"Arguments are passed as a JSON object matching the tool's input schema.",
		],
		parameters: Type.Object({
			server: Type.String({
				description: "MCP server name, e.g. 'railway', 'twitter-mcp'",
			}),
			tool: Type.String({
				description: "Tool name on that server, e.g. 'railway_list_projects'",
			}),
			args: Type.Optional(
				Type.Object(
					{},
					{
						additionalProperties: true,
						description:
							"Tool arguments as key-value pairs matching the tool's input schema",
					},
				),
			),
		}),

		async execute(_id, params, signal) {
			try {
				const client = await getOrConnect(params.server, signal);
				const result = await client.callTool(
					{ name: params.tool, arguments: params.args ?? {} },
					undefined,
					{ signal, timeout: 60000 },
				);

				// Serialize result content to text
				const contentItems = result.content as Array<{
					type: string;
					text?: string;
				}>;
				const raw = contentItems
					.map((c) => (c.type === "text" ? (c.text ?? "") : JSON.stringify(c)))
					.join("\n");

				const truncation = truncateHead(raw, {
					maxLines: DEFAULT_MAX_LINES,
					maxBytes: DEFAULT_MAX_BYTES,
				});
				let finalText = truncation.content;
				if (truncation.truncated) {
					finalText += `\n\n[Output truncated: ${truncation.outputLines}/${truncation.totalLines} lines (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`;
				}

				return {
					content: [{ type: "text", text: finalText }],
					details: {
						server: params.server,
						tool: params.tool,
						charCount: finalText.length,
						truncated: truncation.truncated,
					},
				};
			} catch (err: unknown) {
				const msg = err instanceof Error ? err.message : String(err);
				throw new Error(
					`MCP call failed: ${params.server}.${params.tool}\n${msg}`,
				);
			}
		},

		renderCall(args, theme) {
			let text = theme.fg("toolTitle", theme.bold("mcp_call "));
			text += theme.fg("accent", `${args.server}.${args.tool}`);
			if (args.args && Object.keys(args.args).length > 0) {
				const preview = Object.entries(args.args)
					.slice(0, 3)
					.map(([k, v]) => {
						const val = typeof v === "string" ? v : JSON.stringify(v);
						return `${k}:${val.length > 30 ? val.slice(0, 30) + "…" : val}`;
					})
					.join(" ");
				text += " " + theme.fg("muted", preview);
			}
			return new Text(text, 0, 0);
		},

		renderResult(result, { isPartial, expanded }, theme) {
			if (isPartial)
				return new Text(theme.fg("warning", "Calling MCP tool..."), 0, 0);

			const d = result.details as
				| {
						server: string;
						tool: string;
						charCount: number;
						truncated: boolean;
				  }
				| undefined;

			let text = theme.fg("success", `✓ ${d?.server}.${d?.tool}`);
			text += theme.fg(
				"dim",
				` · ${(d?.charCount ?? 0).toLocaleString()} chars`,
			);
			if (d?.truncated) text += theme.fg("warning", " · truncated");

			if (expanded) {
				const content = result.content[0];
				if (content?.type === "text") {
					const preview = content.text.split("\n").slice(0, 15).join("\n");
					text += "\n\n" + theme.fg("dim", preview);
				}
			}

			return new Text(text, 0, 0);
		},
	});

	// ── Lifecycle ─────────────────────────────────────────────────────────────

	pi.on("session_start", async (_event, ctx) => {
		const servers = readConfigs();
		if (servers.length > 0) {
			ctx.ui.notify(
				`MCP client ready — ${servers.length} server(s) configured`,
				"info",
			);
		}
	});

	pi.on("session_shutdown", async () => {
		await closeAll();
	});

	pi.on("session_switch", async () => {
		await closeAll();
		configCache = null;
	});
}
