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";

export default function (pi: ExtensionAPI) {
  pi.on("session_start", async (_event, ctx) => {
    const token = process.env.NANOCLAW_MCP_TOKEN;
    const serverUrl = process.env.NANOCLAW_MCP_URL || "http://localhost:3002/mcp";

    if (!token) {
      ctx.ui.notify("NANOCLAW_MCP_TOKEN environment variable is missing. MCP tools won't be loaded.", "warning");
      return;
    }

    ctx.ui.notify(`Connecting to nanoclaw MCP server at ${serverUrl}...`, "info");
    try {
      const transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
        requestInit: {
          headers: {
            "Authorization": `Bearer ${token}`
          }
        }
      });

      const client = new Client(
        { name: "pi-mcp-client", version: "1.0.0" },
        { capabilities: { tools: {} } }
      );

      await client.connect(transport);
      const { tools } = await client.listTools();
      
      for (const tool of tools) {
        pi.registerTool({
          name: tool.name,
          description: tool.description || `MCP Tool: ${tool.name}`,
          parameters: Type.Any(),
          async execute(id, params) {
            try {
              const result = await client.callTool({
                name: tool.name,
                arguments: params as any
              });
              
              if (result.isError) {
                throw new Error(result.content.map((c: any) => c.text).join("\n"));
              }
              
              return {
                content: result.content.map((c: any) => ({
                  type: "text",
                  text: c.text || JSON.stringify(c)
                }))
              };
            } catch (err: any) {
              throw new Error(`MCP Tool Error: ${err.message}`);
            }
          }
        });
      }
      
      ctx.ui.notify(`Connected to nanoclaw MCP server. Loaded ${tools.length} tools: ${tools.map(t => t.name).join(', ')}`, "success");
    } catch (error: any) {
      ctx.ui.notify(`Failed to connect to nanoclaw MCP server: ${error.message}`, "error");
    }
  });
}