import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";

export function createWebMcpModule(
  { title }: { title: string; description: string },
  callback: (mcpServer: McpServer) => void
): Client {
  const version = "unset";
  const [clientTransport, serverTransport] =
    InMemoryTransport.createLinkedPair();

  const client: Client = new Client({
    name: title,
    version,
  });
  client.connect(clientTransport);

  const mcpServer: McpServer = new McpServer({
    name: title,
    version,
  });

  callback(mcpServer);

  mcpServer.connect(serverTransport);

  return client;
}

export type McpServerConfig = {
  title: string;
  description: string;
  serverUrl: string;
};

export async function createMcpServerModule(
  config: McpServerConfig
): Promise<Client> {
  const { title, serverUrl } = config;

  // 添加重试逻辑
  const maxRetries = 3;
  let lastError: any;

  for (let attempt = 1; attempt <= maxRetries; attempt++) {
    try {
      console.log(
        `正在连接到MCP服务器 (尝试 ${attempt}/${maxRetries}): ${serverUrl}`
      );

      // 创建MCP客户端连接到外部服务器
      const client = new Client(
        {
          name: title,
          version: "1.0.0",
        },
        {
          capabilities: {
            tools: {},
            resources: {},
          },
        }
      );

      // 根据URL协议选择连接方式
      let transport;
      if (serverUrl.startsWith("ws://") || serverUrl.startsWith("wss://")) {
        // WebSocket连接
        const { WebSocketClientTransport } = await import(
          "@modelcontextprotocol/sdk/client/websocket.js"
        );
        transport = new WebSocketClientTransport(new URL(serverUrl));
      } else if (
        serverUrl.startsWith("http://") ||
        serverUrl.startsWith("https://")
      ) {
        // StreamableHTTP连接 - 使用默认配置，SDK 自动处理 Session ID
        const { StreamableHTTPClientTransport } = await import(
          "@modelcontextprotocol/sdk/client/streamableHttp.js"
        );

        // 使用默认配置，让 SDK 自动处理所有细节
        transport = new StreamableHTTPClientTransport(new URL(serverUrl));
      } else {
        throw new Error(
          `不支持的协议: ${serverUrl}，支持 ws://, wss://, http://, https://`
        );
      }

      // 连接到服务器 - Session ID 会在此过程中自动处理
      await client.connect(transport);

      console.log(`MCP 客户端 "${title}" 已成功连接到: ${serverUrl}`);
      return client;
    } catch (error: any) {
      lastError = error;
      console.warn(`连接尝试 ${attempt} 失败:`, error?.message || error);

      if (attempt < maxRetries) {
        // 等待一段时间后重试
        const delay = attempt * 1000; // 递增延迟
        console.log(`等待 ${delay}ms 后重试...`);
        await new Promise((resolve) => setTimeout(resolve, delay));
      }
    }
  }

  console.error("创建MCP服务器连接失败 (所有重试都失败):", lastError);
  throw new Error(`连接到MCP服务器失败: ${lastError?.message || lastError}`);
}