import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { Type } from "typebox";

const STITCH_URL = process.env.STITCH_MCP_URL ?? "https://stitch.googleapis.com/mcp";
const REQUEST_TIMEOUT_MS = Number(process.env.STITCH_MCP_TIMEOUT_MS ?? 30_000);

type JsonObject = Record<string, unknown>;

type McpTool = {
	name: string;
	description?: string;
	inputSchema?: JsonObject;
};

let stitchApiKey: string | undefined = process.env.GOOGLE_STITCH_API_KEY;
let client: Client | undefined;
let transport: StreamableHTTPClientTransport | undefined;
let toolCache: McpTool[] | undefined;

async function getClient(): Promise<Client> {
	if (client) return client;

	if (!stitchApiKey) throw new Error("Stitch API key not configured. Set GOOGLE_STITCH_API_KEY env var or run /stitch-mcp-set-key <key>.");

	transport = new StreamableHTTPClientTransport(new URL(STITCH_URL), {
		requestInit: {
			headers: {
				Accept: "application/json, text/event-stream",
				"X-Goog-Api-Key": stitchApiKey,
			},
		},
	});

	client = new Client(
		{ name: "pi-stitch-mcp", version: "0.1.0" },
		{ capabilities: {} },
	);

	await client.connect(transport);
	return client;
}

async function listTools(): Promise<McpTool[]> {
	if (toolCache) return toolCache;

	const c = await getClient();
	const tools: McpTool[] = [];
	let cursor: string | undefined;

	do {
		const res = await c.listTools(cursor ? { cursor } : undefined, {
			timeout: REQUEST_TIMEOUT_MS,
		});
		tools.push(...(res.tools as McpTool[]));
		cursor = res.nextCursor;
	} while (cursor);

	toolCache = tools;
	return tools;
}

function resetClient() {
	client = undefined;
	transport = undefined;
	toolCache = undefined;
}

function initMcpTool(pi: ExtensionAPI, tool: McpTool) {
	const safeName = `stitch_${tool.name.replace(/[^a-zA-Z0-9_]/g, "_")}`;
	pi.registerTool({
		name: safeName,
		label: `Stitch: ${tool.name}`,
		description: tool.description ?? `Call Stitch MCP tool ${tool.name}`,
		promptSnippet: `Call Stitch MCP tool ${tool.name}`,
		promptGuidelines: [`Use ${safeName} for Stitch MCP tool ${tool.name}.`],
		parameters: Type.Object({
			arguments: Type.Optional(
				Type.Any({
					description: `Arguments matching this MCP input schema: ${JSON.stringify(tool.inputSchema ?? {})}`,
				}),
			),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
			try {
				const c = await getClient();
				const result = await c.callTool(
					{
						name: tool.name,
						arguments: (params.arguments ?? {}) as JsonObject,
					},
					undefined,
					{ signal, timeout: REQUEST_TIMEOUT_MS },
				);
				return stitchResult(result);
			} catch (error) {
				return stitchResultErr(error);
			}
		},
	});
}

function stitchResult(result: unknown) {
	const r = result as { content?: unknown[]; isError?: boolean };
	const parts: string[] = [];
	if (Array.isArray(r.content)) {
		for (const item of r.content) {
			const c = item as Record<string, unknown>;
			if (c.type === "text" && typeof c.text === "string") {
				parts.push(c.text);
			} else {
				parts.push(JSON.stringify(c, null, 2));
			}
		}
	}
	if (parts.length === 0) parts.push(JSON.stringify(result, null, 2));
	return {
		content: [{ type: "text" as const, text: parts.join("\n\n") }],
		details: { result: result as JsonObject, isError: r.isError },
	};
}

function stitchResultErr(error: unknown) {
	const msg = error instanceof Error ? error.message : String(error);
	return {
		content: [{ type: "text" as const, text: `Stitch MCP error: ${msg}` }],
		details: { result: { error: msg }, isError: true },
	};
}

export default function stitchMcpExtension(pi: ExtensionAPI) {
	pi.registerCommand("stitch-mcp-set-key", {
		description: "Set the Google Stitch API key. Usage: /stitch-mcp-set-key <your-api-key>",
		handler: async (args, ctx) => {
			const key = args?.trim();
			if (!key || key.length < 5) {
				ctx.ui.notify("Stitch MCP: Provide a valid API key.", "error");
				return;
			}
			stitchApiKey = key;
			resetClient();
			ctx.ui.notify("Stitch MCP API key set. Use /stitch-mcp-reset to reconnect.", "info");
		},
	});

	pi.registerCommand("stitch-mcp-register", {
		description:
			"Register Stitch MCP tools as native pi tools. Usage: /stitch-mcp-register or /stitch-mcp-register tool1,tool2",
		handler: async (args, ctx) => {
			try {
				const tools = await listTools();
				const selected = args?.trim()
					? args
							.split(",")
							.map((s) => s.trim())
							.filter(Boolean)
					: tools.map((t) => t.name);
				let count = 0;
				for (const tool of tools.filter((t) => selected.includes(t.name))) {
					initMcpTool(pi, tool);
					count++;
				}
				ctx.ui.notify(`Registered ${count} Stitch MCP tool(s).`, "info");
			} catch (error) {
				ctx.ui.notify(
					`Stitch MCP error: ${error instanceof Error ? error.message : String(error)}`,
					"error",
				);
			}
		},
	});

	pi.registerCommand("stitch-mcp-reset", {
		description: "Reset Stitch MCP client connection",
		handler: async (_args, ctx) => {
			try {
				await transport?.close();
			} catch {
				// ignore
			}
			resetClient();
			ctx.ui.notify("Stitch MCP connection reset.", "info");
		},
	});

	pi.registerTool({
		name: "stitch_mcp_list_tools",
		label: "Stitch MCP: List Tools",
		description: "List tools exposed by the Stitch remote MCP server.",
		promptSnippet: "List Stitch MCP tools available on the remote server",
		promptGuidelines: [
			"Use stitch_mcp_list_tools to discover available Stitch actions before calling stitch_mcp_call_tool.",
		],
		parameters: Type.Object({}),
		async execute(_toolCallId, _params, signal, _onUpdate, _ctx) {
			try {
				const tools = await listTools();
				return {
					content: [
						{
							type: "text" as const,
							text: JSON.stringify(
								tools.map((t) => ({
									name: t.name,
									description: t.description,
									inputSchema: t.inputSchema,
								})),
								null,
								2,
							),
						},
					],
					details: { result: { count: tools.length }, isError: false },
				};
			} catch (error) {
				return stitchResultErr(error);
			}
		},
	});

	pi.registerTool({
		name: "stitch_mcp_call_tool",
		label: "Stitch MCP: Call Tool",
		description:
			"Call any tool exposed by the Stitch remote MCP server by name with JSON arguments.",
		promptSnippet: "Call a Stitch MCP tool by name with JSON arguments",
		promptGuidelines: [
			"Use stitch_mcp_call_tool only after knowing the Stitch MCP tool name and argument schema.",
			"For Stitch operations, call stitch_mcp_list_tools first if the required tool name or arguments are unknown.",
		],
		parameters: Type.Object({
			name: Type.String({ description: "Stitch MCP tool name" }),
			arguments: Type.Optional(
				Type.Any({ description: "JSON arguments for the MCP tool" }),
			),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
			try {
				const c = await getClient();
				const result = await c.callTool(
					{
						name: params.name,
						arguments: (params.arguments ?? {}) as JsonObject,
					},
					undefined,
					{ signal, timeout: REQUEST_TIMEOUT_MS },
				);
				return stitchResult(result);
			} catch (error) {
				return stitchResultErr(error);
			}
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		ctx.ui.setStatus("stitch-mcp", "Stitch MCP");
		try {
			const tools = await listTools();
			ctx.ui.notify(`Stitch MCP connected (${tools.length} tools).`, "info");
		} catch (error) {
			ctx.ui.notify(
				`Stitch MCP not connected: ${error instanceof Error ? error.message : String(error)}. Set key with /stitch-mcp-set-key or GOOGLE_STITCH_API_KEY env var.`,
				"warning",
			);
		}
	});

	pi.on("session_shutdown", async () => {
		try {
			await transport?.close();
		} catch {
			// ignore
		}
		resetClient();
	});
}