/**
 * Created by liujiangyong on 2025/08/20.
 */

'use strict';

const { McpServer } = require("@modelcontextprotocol/sdk/server/mcp.js");
const { StreamableHTTPServerTransport } = require("@modelcontextprotocol/sdk/server/streamableHttp.js");
const { SSEServerTransport } = require("@modelcontextprotocol/sdk/server/sse.js");
const { isInitializeRequest } = require("@modelcontextprotocol/sdk/types.js");
const { randomUUID } = require("node:crypto");
const { z } = require("zod");

module.exports = function (app, router) {
   // 创建 MCP server
   function createServer() {
      const server = new McpServer({
         name: "demo-server",
         version: "1.0.0",
      });

      // 注册一个加法工具
      server.registerTool(
         "add",
         {
            title: "Addition Tool",
            description: "Add two numbers",
            inputSchema: { a: z.number(), b: z.number() },
         },
         async ({ a, b }) => ({
            content: [{ type: "text", text: String(a + b) }],
         })
      );

      // 注册一个获取 AI Query 记录的工具
      server.registerTool("getAiQueryRecord",
         {
            title: "getAiQueryRecord",
            description: "获取 AI Query 记录",
            inputSchema: {
               page: z.number().default(1),
               pageSize: z.number().default(10)
            }
         },
         async ({ page, pageSize }) => {
            const models = app.fs.dc.models;
            const records = await models.AiQueryRecord.findAll({
               limit: pageSize,
               offset: (page - 1) * pageSize,
               order: [['id', 'DESC']],
               raw: true,
            });
            return {
               content: [
                  {
                     type: "text", text: JSON.stringify({ data: records })
                  },
               ],
               structuredContent: { data: records }
            }
         }
      );

      return server;
   }

   // 用于存储每个会话的 transport 实例
   const transports = { streamable: {}, sse: {} };

   //
   // ---- Streamable HTTP 端点 ----
   //
   router.post("/mcp", async (ctx) => {
      const sessionId = ctx.get("mcp-session-id");
      let transport;

      if (sessionId && transports.streamable[sessionId]) {
         // 已有会话
         transport = transports.streamable[sessionId];
      } else if (!sessionId && isInitializeRequest(ctx.request.body)) {
         // 新会话
         const server = createServer();
         transport = new StreamableHTTPServerTransport({
            sessionIdGenerator: () => randomUUID(),
            onsessioninitialized: (sid) => {
               transports.streamable[sid] = transport;
            },
         });

         transport.onclose = () => {
            if (transport.sessionId) delete transports.streamable[transport.sessionId];
         };

         await server.connect(transport);
      } else {
         ctx.status = 400;
         ctx.body = { error: "Bad request: invalid session" };
         return;
      }

      await transport.handleRequest(ctx.req, ctx.res, ctx.request.body);
      ctx.respond = false;
   });

   router.get("/mcp", async (ctx) => {
      const sessionId = ctx.get("mcp-session-id");
      const transport = transports.streamable[sessionId];
      if (!transport) {
         ctx.status = 400;
         ctx.body = "Invalid session";
         return;
      }
      await transport.handleRequest(ctx.req, ctx.res);
      ctx.respond = false;
   });

   router.delete("/mcp", async (ctx) => {
      const sessionId = ctx.get("mcp-session-id");
      const transport = transports.streamable[sessionId];
      if (transport) {
         await transport.close();
         delete transports.streamable[sessionId];
         ctx.body = { ok: true };
      } else {
         ctx.status = 400;
         ctx.body = "Invalid session";
      }
   });

   //
   // ---- SSE 端点 ----
   //
   // router.get("/sse", async (ctx) => {
   //    const server = createServer();
   //    const transport = new SSEServerTransport("/messages", ctx.res);
   //    transports.sse[transport.sessionId] = transport;

   //    ctx.res.on("close", () => {
   //       delete transports.sse[transport.sessionId];
   //    });

   //    await server.connect(transport);
   //    ctx.respond = false;
   // });

   // router.post("/messages", async (ctx) => {
   //    const sessionId = ctx.query.sessionId;
   //    const transport = transports.sse[sessionId];
   //    if (transport) {
   //       await transport.handlePostMessage(ctx.req, ctx.res, ctx.request.body);
   //       ctx.respond = false;
   //    } else {
   //       ctx.status = 400;
   //       ctx.body = "Invalid session";
   //    }
   // });

   // 将路由中间件挂载到应用上
   app.use(router.routes()).use(router.allowedMethods());

   console.log('✅ MCP server is set up /mcp (Streamable) endpoints');
}