import { randomUUID } from "node:crypto";
import type { Express, Request, Response, RequestHandler } from "express";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import type { AuthenticatedDingTalkUser } from "./xft.js";
import { McpSessionLifecycle } from "./mcp-session-lifecycle.js";
import { DEFAULT_MCP_SESSION_MAX_PER_USER, MCP_SESSION_MAX_TOTAL } from "./config.js";

type TransportEntry =
  | { kind: "streamable"; transport: StreamableHTTPServerTransport; userId: string; lifecycle: McpSessionLifecycle }
  | { kind: "sse"; transport: SSEServerTransport; userId: string; lifecycle: McpSessionLifecycle };

interface McpTransportOptions {
  authenticate?: RequestHandler;
  // Production always resolves this from verified OAuth, never request parameters.
  requestUser: (req: Request) => AuthenticatedDingTalkUser;
  createServer: (user: AuthenticatedDingTalkUser, lifecycle: McpSessionLifecycle) => McpServer;
  now?: () => number;
  maxSessionsPerUser?: number;
}

export function installMcpTransports(app: Express, options: McpTransportOptions) {
  const now = options.now ?? Date.now;
  const maxSessionsPerUser = options.maxSessionsPerUser ?? DEFAULT_MCP_SESSION_MAX_PER_USER;
  const transports = new Map<string, TransportEntry>();
  let pendingTransportSessions = 0;
  const pendingTransportSessionsByUser = new Map<string, number>();
  const closeTransport = (id: string, entry: TransportEntry) => {
    transports.delete(id);
    void entry.transport.close().catch(() => undefined);
  };
  const cleanupTransports = () => {
    for (const [id, entry] of transports) {
      if (entry.lifecycle.expired()) {
        closeTransport(id, entry);
      }
    }
  };
  const sessionCapacityAvailable = (userId: string) => transports.size + pendingTransportSessions < MCP_SESSION_MAX_TOTAL
    && [...transports.values()].filter((entry) => entry.userId === userId).length
      + (pendingTransportSessionsByUser.get(userId) ?? 0) < maxSessionsPerUser;
  const reserveSession = (userId: string) => {
    pendingTransportSessions += 1;
    pendingTransportSessionsByUser.set(userId, (pendingTransportSessionsByUser.get(userId) ?? 0) + 1);
  };
  const releaseSessionReservation = (userId: string) => {
    pendingTransportSessions = Math.max(0, pendingTransportSessions - 1);
    const remaining = Math.max(0, (pendingTransportSessionsByUser.get(userId) ?? 1) - 1);
    if (remaining === 0) pendingTransportSessionsByUser.delete(userId);
    else pendingTransportSessionsByUser.set(userId, remaining);
  };
  const transportCleanupTimer = setInterval(cleanupTransports, 15_000);
  transportCleanupTimer.unref();

  const mcpHandler = async (req: Request, res: Response) => {
    let reservedUserId: string | undefined;
    try {
      cleanupTransports();
      const user = options.requestUser(req);
      const sessionId = req.get("mcp-session-id");
      let entry = sessionId ? transports.get(sessionId) : undefined;
      if (entry && entry.kind !== "streamable") {
        res.status(400).json({ error: "Invalid MCP session transport" });
        return;
      }
      if (entry && entry.userId !== user.userId) {
        res.status(403).json({ error: "MCP session belongs to another authenticated user" });
        return;
      }
      if (entry) entry.lifecycle.touch();
      if (!entry && req.method === "POST" && !sessionId && isInitializeRequest(req.body)) {
        if (!sessionCapacityAvailable(user.userId)) {
          res.set("Retry-After", "60").status(429).json({ error: "MCP session limit reached; close an old Agent session and retry" });
          return;
        }
        reserveSession(user.userId);
        reservedUserId = user.userId;
        const transport: StreamableHTTPServerTransport = new StreamableHTTPServerTransport({
          sessionIdGenerator: () => randomUUID(),
          onsessioninitialized: (id) => {
            transports.set(id, entry!);
            // The live entry now owns the slot; do not count its reservation twice.
            releaseSessionReservation(user.userId);
            reservedUserId = undefined;
          },
        });
        transport.onclose = () => { if (transport.sessionId) transports.delete(transport.sessionId); };
        entry = { kind: "streamable", transport, userId: user.userId, lifecycle: new McpSessionLifecycle(now) };
        await options.createServer(user, entry.lifecycle).connect(transport);
      }
      if (!entry || entry.kind !== "streamable") {
        res.status(sessionId ? 404 : 400).json({ error: "Missing or invalid MCP session" });
        return;
      }
      if (req.method === "GET") entry.lifecycle.trackStream(res);
      await entry.transport.handleRequest(req, res, req.body);
    } catch (error) { sendServerError(res, error); }
    finally { if (reservedUserId) releaseSessionReservation(reservedUserId); }
  };
  app.all("/mcp", ...(options.authenticate ? [options.authenticate, mcpHandler] : [mcpHandler]));

  const sseHandler = async (req: Request, res: Response) => {
    const user = options.requestUser(req);
    cleanupTransports();
    if (!sessionCapacityAvailable(user.userId)) {
      res.set("Retry-After", "60").status(429).send("MCP session limit reached");
      return;
    }
    const transport = new SSEServerTransport("/messages", res);
    const lifecycle = new McpSessionLifecycle(now);
    lifecycle.trackStream(res);
    transports.set(transport.sessionId, { kind: "sse", transport, userId: user.userId, lifecycle });
    res.on("close", () => transports.delete(transport.sessionId));
    await options.createServer(user, lifecycle).connect(transport);
  };
  app.get("/sse", ...(options.authenticate ? [options.authenticate, sseHandler] : [sseHandler]));
  const messagesHandler = async (req: Request, res: Response) => {
    const user = options.requestUser(req);
    const entry = transports.get(String(req.query.sessionId ?? ""));
    if (!entry || entry.kind !== "sse") {
      res.status(404).send("Invalid SSE session");
      return;
    }
    if (entry.userId !== user.userId) {
      res.status(403).send("Invalid SSE session");
      return;
    }
    entry.lifecycle.touch();
    await entry.transport.handlePostMessage(req, res, req.body);
  };
  app.post("/messages", ...(options.authenticate ? [options.authenticate, messagesHandler] : [messagesHandler]));
  return {
    cleanup: cleanupTransports,
    async close() {
      clearInterval(transportCleanupTimer);
      const entries = [...transports.values()];
      transports.clear();
      await Promise.all(entries.map((entry) => entry.transport.close()));
    },
  };
}

function sendServerError(res: Response, _error: unknown): void {
  console.error("[dingtalk-approval-mcp] MCP transport request failed");
  if (!res.headersSent) res.status(500).json({ error: "Internal server error" });
}