import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import express from "express";
import pino from "pino";
import { z } from "zod";

import { outlineMcpFactory } from "./outline";

const sessionIdSchema = z.string().uuid();

async function main() {
  const logger = pino({
    level: "debug",
  });

  const transports: Record<string, StreamableHTTPServerTransport> = {};

  const app = express();
  app.use(express.json());

  app.post("/mcp", async (req, res) => {
    logger.debug(
      {
        body: JSON.stringify(req.body),
      },
      "Received MCP request"
    );
    const sessionIdHeader = req.headers["mcp-session-id"] as string | undefined;
    let transport: StreamableHTTPServerTransport;

    if (sessionIdHeader) {
      const safeSessionId = sessionIdSchema.safeParse(sessionIdHeader);
      if (!safeSessionId.success) {
        logger.error("Invalid session ID format");
        res.status(400).json({
          jsonrpc: "2.0",
          error: {
            code: -32000,
            message: "Invalid session ID format",
          },
          id: null,
        });
        return;
      }

      const sessionId = safeSessionId.data;
      if (transports[sessionId]) {
        transport = transports[sessionId];
      } else {
        logger.error("Session not found");
        res.status(400).json({
          jsonrpc: "2.0",
          error: {
            code: -32000,
            message: "Session not found",
          },
          id: null,
        });
        return;
      }
    } else if (isInitializeRequest(req.body)) {
      transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => crypto.randomUUID(),
        onsessioninitialized: (sessionId) => {
          transports[sessionId] = transport;
        },
      });

      transport.onclose = () => {
        if (transport.sessionId) {
          delete transports[transport.sessionId];
        }
      };
      const outlineMcpServer = outlineMcpFactory(logger.child({ sessionId: transport.sessionId }));

      await outlineMcpServer.connect(transport);
    } else {
      logger.error("Bad Request: No valid session ID provided");
      res.status(400).json({
        jsonrpc: "2.0",
        error: {
          code: -32000,
          message: "Bad Request: No valid session ID provided",
        },
        id: null,
      });
      return;
    }

    await transport.handleRequest(req, res, req.body);
  });

  app.listen(3000, () => {
    logger.info("Server is running on port 3000");
  });
}

main();