import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import {
  EventStore,
  StreamableHTTPServerTransport,
} from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import {
  isInitializeRequest,
  JSONRPCMessageSchema,
} from "@modelcontextprotocol/sdk/types.js";
import fs from "fs";
import http from "http";
import https from "https";
import { randomUUID } from "node:crypto";

import { AuthConfig, AuthenticationMiddleware } from "./authentication.js";
import { InMemoryEventStore } from "./InMemoryEventStore.js";

const DEFAULT_KEEP_ALIVE_TIMEOUT = 300_000;

/**
 * `false` disables the resumability event store entirely (no replay-on-
 * reconnect, no retained state). Omitted/`undefined` creates a fresh,
 * bounded `InMemoryEventStore` per session - see `eventStoreMaxEvents`.
 * Pass an `EventStore` instance to use a shared or custom-backed store.
 */
export type EventStoreOption = EventStore | false;

const resolveEventStore = (
  eventStore: EventStoreOption | undefined,
  maxEvents: number | undefined,
): EventStore | undefined => {
  if (eventStore === false) {
    return undefined;
  }

  return eventStore ?? new InMemoryEventStore({ maxEvents });
};

export interface CorsOptions {
  allowedHeaders?: string | string[]; // Allow string[] or '*' for wildcard
  credentials?: boolean;
  exposedHeaders?: string[];
  maxAge?: number;
  methods?: string[];
  origin?: ((origin: string) => boolean) | string | string[];
}

export type SSEServer = {
  close: () => Promise<void>;
};

type ServerLike = {
  close: Server["close"];
  connect: Server["connect"];
};

const DEFAULT_MAX_BODY_SIZE = 10_485_760; // 10 MiB

/**
 * Caps how many bytes of a request body the stream endpoint buffers before it
 * gives up. `false` (or `0` from the CLI) disables the cap entirely, restoring
 * unbounded buffering - only do that behind a gateway that already limits body
 * size. Omitted/`undefined` uses `DEFAULT_MAX_BODY_SIZE`.
 */
export type MaxBodySizeOption = false | number;

/**
 * "Too large" is kept distinct from the `null` that every other unusable body
 * resolves to, so the caller can answer 413 rather than let it fall through to
 * the generic 400. The limit travels with the signal so the response can name
 * it.
 */
type BodyResult =
  | { readonly body: unknown; readonly tooLarge?: never }
  | { readonly limit: number; readonly tooLarge: true };

const getBody = (
  request: http.IncomingMessage,
  maxBodySize: MaxBodySizeOption = DEFAULT_MAX_BODY_SIZE,
) => {
  return new Promise<BodyResult>((resolve) => {
    if (maxBodySize !== false) {
      // A client that declares its size up front can be rejected before a
      // single byte of body is read. The streaming check below is still
      // needed for chunked bodies and for clients that under-declare.
      const declaredSize = Number(request.headers["content-length"]);

      if (Number.isFinite(declaredSize) && declaredSize > maxBodySize) {
        resolve({ limit: maxBodySize, tooLarge: true });

        return;
      }
    }

    const bodyParts: Buffer[] = [];
    let body: string;
    let size = 0;
    request
      .on("data", (chunk) => {
        if (maxBodySize !== false) {
          size += chunk.length;
          if (size > maxBodySize) {
            resolve({ limit: maxBodySize, tooLarge: true });
            return;
          }
        }
        bodyParts.push(chunk);
      })
      .on("end", () => {
        body = Buffer.concat(bodyParts).toString();
        try {
          resolve({ body: JSON.parse(body) });
        } catch (error) {
          console.error("[mcp-proxy] error parsing body", error);
          resolve({ body: null });
        }
      })
      .on("error", (error) => {
        console.error("[mcp-proxy] error reading body", error);
        resolve({ body: null });
      })
      .on("close", () => {
        resolve({ body: null });
      });
  });
};

// Helper function to create JSON RPC error responses
const createJsonRpcErrorResponse = (code: number, message: string) => {
  return JSON.stringify({
    error: { code, message },
    id: null,
    jsonrpc: "2.0",
  });
};

/**
 * Answers an over-sized request with a 413 and only then tears the connection
 * down. Destroying the socket outright - which is all the size check can do on
 * its own - leaves the client with a bare ECONNRESET and no way to tell a size
 * limit from a crash.
 */
const sendPayloadTooLarge = ({
  maxBodySize,
  req,
  res,
}: {
  readonly maxBodySize: number;
  readonly req: http.IncomingMessage;
  readonly res: http.ServerResponse;
}) => {
  console.error(
    `[mcp-proxy] request body too large (exceeds ${maxBodySize} bytes)`,
  );

  // Stop consuming immediately so a client that keeps sending applies TCP
  // backpressure instead of growing this process's buffers.
  req.pause();

  if (res.headersSent) {
    req.destroy();

    return;
  }

  res.setHeader("Connection", "close");
  res.setHeader("Content-Type", "application/json");

  // Destroy only once the response has flushed, otherwise the socket can go
  // away before the client ever sees the 413.
  res.writeHead(413).end(
    createJsonRpcErrorResponse(
      -32600,
      `Payload Too Large: request body exceeds ${maxBodySize} bytes`,
    ),
    () => {
      req.destroy();
    },
  );
};

type SessionUnauthorizedResponseOptions = {
  readonly body?: unknown;
  readonly oauth?: AuthConfig["oauth"];
  readonly res: http.ServerResponse;
};

const getRequestId = (body: unknown): unknown => {
  if (
    typeof body !== "object" ||
    body === null ||
    Array.isArray(body) ||
    !("id" in body)
  ) {
    return null;
  }

  return body.id;
};

const isJsonRpcMessage = (message: unknown): boolean => {
  return JSONRPCMessageSchema.safeParse(message).success;
};

const isJsonRpcBody = (body: unknown): boolean => {
  return Array.isArray(body)
    ? body.every(isJsonRpcMessage)
    : isJsonRpcMessage(body);
};

// Helper function to get WWW-Authenticate header value
const getWWWAuthenticateHeader = (
  oauth?: AuthConfig["oauth"],
  options?: {
    error?: string;
    error_description?: string;
    error_uri?: string;
    scope?: string;
  },
): string | undefined => {
  if (!oauth) {
    return undefined;
  }

  const params: string[] = [];

  // Add realm if configured
  if (oauth.realm) {
    params.push(`realm="${oauth.realm}"`);
  }

  // Add resource_metadata if configured
  if (oauth.protectedResource?.resource) {
    params.push(
      `resource_metadata="${oauth.protectedResource.resource}/.well-known/oauth-protected-resource"`,
    );
  }

  // Add error from options or config (options takes precedence)
  const error = options?.error || oauth.error;
  if (error) {
    params.push(`error="${error}"`);
  }

  // Add error_description from options or config (options takes precedence)
  const error_description =
    options?.error_description || oauth.error_description;
  if (error_description) {
    // Escape quotes in error description
    const escaped = error_description.replace(/"/g, '\\"');
    params.push(`error_description="${escaped}"`);
  }

  // Add error_uri from options or config (options takes precedence)
  const error_uri = options?.error_uri || oauth.error_uri;
  if (error_uri) {
    params.push(`error_uri="${error_uri}"`);
  }

  // Add scope from options or config (options takes precedence)
  const scope = options?.scope || oauth.scope;
  if (scope) {
    params.push(`scope="${scope}"`);
  }

  // Return undefined if no parameters were added
  if (params.length === 0) {
    return undefined;
  }

  return `Bearer ${params.join(", ")}`;
};

const sendSessionUnauthorizedResponse = ({
  body,
  oauth,
  res,
}: SessionUnauthorizedResponseOptions): void => {
  const message = "Unauthorized: No valid session ID provided";

  res.setHeader("Content-Type", "application/json");

  const wwwAuthHeader = getWWWAuthenticateHeader(oauth, {
    error: "invalid_token",
    error_description: message,
  });
  if (wwwAuthHeader) {
    res.setHeader("WWW-Authenticate", wwwAuthHeader);
  }

  res.writeHead(401).end(
    JSON.stringify({
      error: {
        code: -32000,
        message,
      },
      id: getRequestId(body),
      jsonrpc: "2.0",
    }),
  );
};

// Helper function to detect scope challenge errors
const isScopeChallengeError = (
  error: unknown,
): error is {
  data: {
    error: string;
    errorDescription?: string;
    requiredScopes: string[];
  };
  name: string;
} => {
  return (
    typeof error === "object" &&
    error !== null &&
    "name" in error &&
    error.name === "InsufficientScopeError" &&
    "data" in error &&
    typeof error.data === "object" &&
    error.data !== null &&
    "error" in error.data &&
    error.data.error === "insufficient_scope"
  );
};

// Helper function to handle Response errors and send appropriate HTTP response
const handleResponseError = async (
  error: unknown,
  res: http.ServerResponse,
): Promise<boolean> => {
  // Check if it's a Response-like object (duck typing)
  // The instanceof check may fail due to different Response implementations across module boundaries
  const isResponseLike =
    error &&
    typeof error === "object" &&
    "status" in error &&
    "headers" in error &&
    "statusText" in error;

  if (isResponseLike || error instanceof Response) {
    const responseError = error as Response;

    // Convert Headers to http.OutgoingHttpHeaders format
    const fixedHeaders: http.OutgoingHttpHeaders = {};
    responseError.headers.forEach((value, key) => {
      if (fixedHeaders[key]) {
        if (Array.isArray(fixedHeaders[key])) {
          (fixedHeaders[key] as string[]).push(value);
        } else {
          fixedHeaders[key] = [fixedHeaders[key] as string, value];
        }
      } else {
        fixedHeaders[key] = value;
      }
    });

    // Read the body from the Response object
    const body = await responseError.text();

    res.writeHead(responseError.status, responseError.statusText, fixedHeaders);
    res.end(body);

    return true;
  }

  return false;
};

// Helper function to clean up server resources
const cleanupServer = async <T extends ServerLike>(
  server: T,
  onClose?: (server: T) => Promise<void>,
) => {
  if (onClose) {
    await onClose(server);
  }

  try {
    await server.close();
  } catch (error) {
    console.error("[mcp-proxy] error closing server", error);
  }
};

// Helper function to apply CORS headers
const applyCorsHeaders = (
  req: http.IncomingMessage,
  res: http.ServerResponse,
  corsOptions?: boolean | CorsOptions,
) => {
  if (!req.headers.origin) {
    return;
  }

  // Default CORS configuration for backward compatibility
  const defaultCorsOptions: CorsOptions = {
    allowedHeaders:
      "Content-Type, Authorization, Accept, Mcp-Session-Id, Mcp-Protocol-Version, Last-Event-Id",
    credentials: true,
    exposedHeaders: ["Mcp-Session-Id"],
    methods: ["GET", "POST", "OPTIONS"],
    origin: "*",
  };

  let finalCorsOptions: CorsOptions;

  if (corsOptions === false) {
    // CORS disabled
    return;
  } else if (corsOptions === true || corsOptions === undefined) {
    // Use default CORS settings
    finalCorsOptions = defaultCorsOptions;
  } else {
    // Merge user options with defaults
    finalCorsOptions = {
      ...defaultCorsOptions,
      ...corsOptions,
    };
  }

  try {
    const origin = new URL(req.headers.origin);

    // Handle origin
    let allowedOrigin = "*";
    if (finalCorsOptions.origin) {
      if (typeof finalCorsOptions.origin === "string") {
        allowedOrigin = finalCorsOptions.origin;
      } else if (Array.isArray(finalCorsOptions.origin)) {
        allowedOrigin = finalCorsOptions.origin.includes(origin.origin)
          ? origin.origin
          : "false";
      } else if (typeof finalCorsOptions.origin === "function") {
        allowedOrigin = finalCorsOptions.origin(origin.origin)
          ? origin.origin
          : "false";
      }
    }

    if (allowedOrigin !== "false") {
      res.setHeader("Access-Control-Allow-Origin", allowedOrigin);
    }

    // Handle credentials
    if (finalCorsOptions.credentials !== undefined) {
      res.setHeader(
        "Access-Control-Allow-Credentials",
        finalCorsOptions.credentials.toString(),
      );
    }

    // Handle methods
    if (finalCorsOptions.methods) {
      res.setHeader(
        "Access-Control-Allow-Methods",
        finalCorsOptions.methods.join(", "),
      );
    }

    // Handle allowed headers
    if (finalCorsOptions.allowedHeaders) {
      const allowedHeaders =
        typeof finalCorsOptions.allowedHeaders === "string"
          ? finalCorsOptions.allowedHeaders
          : finalCorsOptions.allowedHeaders.join(", ");
      res.setHeader("Access-Control-Allow-Headers", allowedHeaders);
    }

    // Handle exposed headers
    if (finalCorsOptions.exposedHeaders) {
      res.setHeader(
        "Access-Control-Expose-Headers",
        finalCorsOptions.exposedHeaders.join(", "),
      );
    }

    // Handle max age
    if (finalCorsOptions.maxAge !== undefined) {
      res.setHeader(
        "Access-Control-Max-Age",
        finalCorsOptions.maxAge.toString(),
      );
    }
  } catch (error) {
    console.error("[mcp-proxy] error parsing origin", error);
  }
};

const handleStreamRequest = async <T extends ServerLike>({
  activeTransports,
  authenticate,
  authMiddleware,
  createServer,
  enableJsonResponse,
  endpoint,
  eventStore,
  eventStoreMaxEvents,
  maxBodySize,
  oauth,
  onClose,
  onConnect,
  req,
  res,
  stateless,
}: {
  activeTransports: Record<
    string,
    { server: T; transport: StreamableHTTPServerTransport }
  >;
  authenticate?: (request: http.IncomingMessage) => Promise<unknown>;
  authMiddleware: AuthenticationMiddleware;
  createServer: (request: http.IncomingMessage) => Promise<T>;
  enableJsonResponse?: boolean;
  endpoint: string;
  eventStore?: EventStoreOption;
  eventStoreMaxEvents?: number;
  maxBodySize?: MaxBodySizeOption;
  oauth?: AuthConfig["oauth"];
  onClose?: (server: T) => Promise<void>;
  onConnect?: (server: T) => Promise<void>;
  req: http.IncomingMessage;
  res: http.ServerResponse;
  stateless?: boolean;
}) => {
  if (
    req.method === "POST" &&
    new URL(req.url!, "http://localhost").pathname === endpoint
  ) {
    let body: unknown;
    try {
      // In stateless mode, ignore session ID header entirely (like Python MCP SDK)
      const sessionId = stateless
        ? undefined
        : Array.isArray(req.headers["mcp-session-id"])
          ? req.headers["mcp-session-id"][0]
          : req.headers["mcp-session-id"];

      let transport: StreamableHTTPServerTransport;

      let server: T;

      const bodyResult = await getBody(req, maxBodySize);

      if (bodyResult.tooLarge) {
        sendPayloadTooLarge({ maxBodySize: bodyResult.limit, req, res });

        return true;
      }

      body = bodyResult.body;

      // Per-request authentication for all requests
      // Store authResult to update existing sessions with fresh auth context
      let authResult: unknown;
      if (authenticate) {
        try {
          authResult = await authenticate(req);

          // Check for both falsy AND { authenticated: false } pattern
          if (
            !authResult ||
            (typeof authResult === "object" &&
              "authenticated" in authResult &&
              !authResult.authenticated)
          ) {
            // Extract error message if available
            const errorMessage =
              authResult &&
              typeof authResult === "object" &&
              "error" in authResult &&
              typeof authResult.error === "string"
                ? authResult.error
                : "Unauthorized: Authentication failed";

            res.setHeader("Content-Type", "application/json");

            // Add WWW-Authenticate header if OAuth config is available
            const wwwAuthHeader = getWWWAuthenticateHeader(oauth, {
              error: "invalid_token",
              error_description: errorMessage,
            });
            if (wwwAuthHeader) {
              res.setHeader("WWW-Authenticate", wwwAuthHeader);
            }

            res.writeHead(401).end(
              JSON.stringify({
                error: {
                  code: -32000,
                  message: errorMessage,
                },
                id: (body as { id?: unknown })?.id ?? null,
                jsonrpc: "2.0",
              }),
            );
            return true;
          }
        } catch (error) {
          // Check if error is a Response object with headers already set
          if (await handleResponseError(error, res)) {
            return true;
          }

          // Extract error details from thrown errors
          const errorMessage =
            error instanceof Error
              ? error.message
              : "Unauthorized: Authentication error";
          console.error("Authentication error:", error);
          res.setHeader("Content-Type", "application/json");

          // Add WWW-Authenticate header if OAuth config is available
          const wwwAuthHeader = getWWWAuthenticateHeader(oauth, {
            error: "invalid_token",
            error_description: errorMessage,
          });
          if (wwwAuthHeader) {
            res.setHeader("WWW-Authenticate", wwwAuthHeader);
          }

          res.writeHead(401).end(
            JSON.stringify({
              error: {
                code: -32000,
                message: errorMessage,
              },
              id: (body as { id?: unknown })?.id ?? null,
              jsonrpc: "2.0",
            }),
          );
          return true;
        }
      }

      if (sessionId) {
        const activeTransport = activeTransports[sessionId];
        if (!activeTransport) {
          if (authenticate && isJsonRpcBody(body)) {
            sendSessionUnauthorizedResponse({ body, oauth, res });

            return true;
          }

          res.setHeader("Content-Type", "application/json");
          res
            .writeHead(404)
            .end(createJsonRpcErrorResponse(-32001, "Session not found"));

          return true;
        }

        transport = activeTransport.transport;
        server = activeTransport.server;

        // Update session's auth context with fresh authentication result
        if (
          authResult &&
          typeof server === "object" &&
          server !== null &&
          "updateAuth" in server &&
          typeof (server as { updateAuth?: unknown }).updateAuth === "function"
        ) {
          (server as { updateAuth: (auth: unknown) => void }).updateAuth(
            authResult,
          );
        }
      } else if (!sessionId && isInitializeRequest(body)) {
        // Create a new transport for the session
        transport = new StreamableHTTPServerTransport({
          enableJsonResponse,
          eventStore: resolveEventStore(eventStore, eventStoreMaxEvents),
          onsessioninitialized: (_sessionId) => {
            // add only when the id Session id is generated (skip in stateless mode)
            if (!stateless && _sessionId) {
              activeTransports[_sessionId] = {
                server,
                transport,
              };
            }
          },
          sessionIdGenerator: stateless ? undefined : randomUUID,
        });

        // Handle the server close event
        let isCleaningUp = false;

        transport.onclose = async () => {
          const sid = transport.sessionId;

          if (isCleaningUp) {
            return;
          }

          isCleaningUp = true;

          if (!stateless && sid && activeTransports[sid]) {
            await cleanupServer(server, onClose);
            delete activeTransports[sid];
          } else if (stateless) {
            // In stateless mode, always call onClose when transport closes
            await cleanupServer(server, onClose);
          }
        };

        try {
          server = await createServer(req);
        } catch (error) {
          // Check if error is a Response object with headers already set
          if (await handleResponseError(error, res)) {
            return true;
          }

          // Detect authentication errors and return HTTP 401
          const errorMessage =
            error instanceof Error ? error.message : String(error);
          const isAuthError =
            errorMessage.includes("Authentication") ||
            errorMessage.includes("Invalid JWT") ||
            errorMessage.includes("Token") ||
            errorMessage.includes("Unauthorized");

          if (isAuthError) {
            res.setHeader("Content-Type", "application/json");

            // Add WWW-Authenticate header if OAuth config is available
            const wwwAuthHeader = getWWWAuthenticateHeader(oauth, {
              error: "invalid_token",
              error_description: errorMessage,
            });
            if (wwwAuthHeader) {
              res.setHeader("WWW-Authenticate", wwwAuthHeader);
            }

            res.writeHead(401).end(
              JSON.stringify({
                error: {
                  code: -32000,
                  message: errorMessage,
                },
                id: (body as { id?: unknown })?.id ?? null,
                jsonrpc: "2.0",
              }),
            );
            return true;
          }

          res.writeHead(500).end("Error creating server");

          return true;
        }

        server.connect(transport);

        if (onConnect) {
          await onConnect(server);
        }

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

        return true;
      } else if (stateless && !sessionId && !isInitializeRequest(body)) {
        // In stateless mode, handle non-initialize requests by creating a new transport
        transport = new StreamableHTTPServerTransport({
          enableJsonResponse,
          eventStore: resolveEventStore(eventStore, eventStoreMaxEvents),
          onsessioninitialized: () => {
            // No session tracking in stateless mode
          },
          sessionIdGenerator: undefined,
        });

        try {
          server = await createServer(req);
        } catch (error) {
          // Check if error is a Response object with headers already set
          if (await handleResponseError(error, res)) {
            return true;
          }

          // Detect authentication errors and return HTTP 401
          const errorMessage =
            error instanceof Error ? error.message : String(error);
          const isAuthError =
            errorMessage.includes("Authentication") ||
            errorMessage.includes("Invalid JWT") ||
            errorMessage.includes("Token") ||
            errorMessage.includes("Unauthorized");

          if (isAuthError) {
            res.setHeader("Content-Type", "application/json");

            // Add WWW-Authenticate header if OAuth config is available
            const wwwAuthHeader = getWWWAuthenticateHeader(oauth, {
              error: "invalid_token",
              error_description: errorMessage,
            });
            if (wwwAuthHeader) {
              res.setHeader("WWW-Authenticate", wwwAuthHeader);
            }

            res.writeHead(401).end(
              JSON.stringify({
                error: {
                  code: -32000,
                  message: errorMessage,
                },
                id: (body as { id?: unknown })?.id ?? null,
                jsonrpc: "2.0",
              }),
            );
            return true;
          }

          res.writeHead(500).end("Error creating server");

          return true;
        }

        server.connect(transport);

        if (onConnect) {
          await onConnect(server);
        }

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

        return true;
      } else {
        if (authenticate && isJsonRpcBody(body)) {
          sendSessionUnauthorizedResponse({ body, oauth, res });

          return true;
        }

        // Error if the server is not created but the request is not an initialize request
        res.setHeader("Content-Type", "application/json");

        res
          .writeHead(400)
          .end(
            createJsonRpcErrorResponse(
              -32000,
              "Bad Request: No valid session ID provided",
            ),
          );

        return true;
      }

      // Handle the request if the server is already created
      await transport.handleRequest(req, res, body);

      return true;
    } catch (error) {
      // Check for scope challenge errors
      if (isScopeChallengeError(error)) {
        const response = authMiddleware.getScopeChallengeResponse(
          error.data.requiredScopes,
          error.data.errorDescription,
          (body as { id?: unknown })?.id,
        );

        res.writeHead(response.statusCode, response.headers);
        res.end(response.body);
        return true;
      }

      console.error("[mcp-proxy] error handling request", error);

      res.setHeader("Content-Type", "application/json");

      res
        .writeHead(500)
        .end(createJsonRpcErrorResponse(-32603, "Internal Server Error"));
    }
    return true;
  }

  if (
    req.method === "GET" &&
    new URL(req.url!, "http://localhost").pathname === endpoint
  ) {
    const sessionId = req.headers["mcp-session-id"] as string | undefined;
    const activeTransport:
      | {
          server: T;
          transport: StreamableHTTPServerTransport;
        }
      | undefined = sessionId ? activeTransports[sessionId] : undefined;

    if (!sessionId) {
      // Return METHOD_NOT_ALLOWED so stateless clients' transport stops reconnecting
      if (stateless) {
        res.writeHead(405, { Allow: "POST" }).end("Method Not Allowed");

        return true;
      }

      if (authenticate) {
        sendSessionUnauthorizedResponse({ oauth, res });

        return true;
      }

      res.writeHead(400).end("No sessionId");

      return true;
    }

    if (!activeTransport) {
      if (authenticate) {
        sendSessionUnauthorizedResponse({ oauth, res });

        return true;
      }

      res.writeHead(400).end("No active transport");

      return true;
    }

    const lastEventId = req.headers["last-event-id"] as string | undefined;

    if (lastEventId) {
      console.log(
        `[mcp-proxy] client reconnecting with Last-Event-ID ${lastEventId} for session ID ${sessionId}`,
      );
    } else {
      console.log(
        `[mcp-proxy] establishing new SSE stream for session ID ${sessionId}`,
      );
    }

    await activeTransport.transport.handleRequest(req, res);

    return true;
  }

  if (
    req.method === "DELETE" &&
    new URL(req.url!, "http://localhost").pathname === endpoint
  ) {
    console.log("[mcp-proxy] received delete request");

    const sessionId = req.headers["mcp-session-id"] as string | undefined;

    if (!sessionId) {
      if (authenticate) {
        sendSessionUnauthorizedResponse({ oauth, res });

        return true;
      }

      res.writeHead(400).end("Invalid or missing sessionId");

      return true;
    }

    console.log("[mcp-proxy] received delete request for session", sessionId);

    const activeTransport = activeTransports[sessionId];

    if (!activeTransport) {
      if (authenticate) {
        sendSessionUnauthorizedResponse({ oauth, res });

        return true;
      }

      res.writeHead(400).end("No active transport");
      return true;
    }

    try {
      // handleRequest for DELETE calls transport.close() internally,
      // which triggers the transport.onclose callback that already
      // handles server cleanup. No need to call cleanupServer again.
      await activeTransport.transport.handleRequest(req, res);
    } catch (error) {
      console.error("[mcp-proxy] error handling delete request", error);

      if (!res.headersSent) {
        res.writeHead(500).end("Error handling delete request");
      }
    }

    return true;
  }

  return false;
};

const handleSSERequest = async <T extends ServerLike>({
  activeTransports,
  createServer,
  endpoint,
  onClose,
  onConnect,
  req,
  res,
}: {
  activeTransports: Record<string, SSEServerTransport>;
  createServer: (request: http.IncomingMessage) => Promise<T>;
  endpoint: string;
  onClose?: (server: T) => Promise<void>;
  onConnect?: (server: T) => Promise<void>;
  req: http.IncomingMessage;
  res: http.ServerResponse;
}) => {
  if (
    req.method === "GET" &&
    new URL(req.url!, "http://localhost").pathname === endpoint
  ) {
    const transport = new SSEServerTransport("/messages", res);

    let server: T;

    try {
      server = await createServer(req);
    } catch (error) {
      if (await handleResponseError(error, res)) {
        return true;
      }

      res.writeHead(500).end("Error creating server");

      return true;
    }

    activeTransports[transport.sessionId] = transport;

    let closed = false;
    let isCleaningUp = false;

    res.on("close", async () => {
      closed = true;

      // Prevent recursive cleanup
      if (isCleaningUp) {
        return;
      }

      isCleaningUp = true;
      await cleanupServer(server, onClose);

      delete activeTransports[transport.sessionId];
    });

    try {
      await server.connect(transport);

      await transport.send({
        jsonrpc: "2.0",
        method: "notifications/message",
        params: { data: "SSE Connection established", level: "info" },
      });

      if (onConnect) {
        await onConnect(server);
      }
    } catch (error) {
      if (!closed) {
        console.error("[mcp-proxy] error connecting to server", error);

        if (!res.headersSent) {
          res.writeHead(500).end("Error connecting to server");
        }
      }
    }

    return true;
  }

  if (req.method === "POST" && req.url?.startsWith("/messages")) {
    const sessionId = new URL(req.url, "https://example.com").searchParams.get(
      "sessionId",
    );

    if (!sessionId) {
      res.writeHead(400).end("No sessionId");

      return true;
    }

    const activeTransport: SSEServerTransport | undefined =
      activeTransports[sessionId];

    if (!activeTransport) {
      res.writeHead(400).end("No active transport");

      return true;
    }

    // `maxBodySize` deliberately does not reach here: the SDK reads and parses
    // this body itself, so the proxy never buffers it and has nothing to cap.
    // The SSE endpoint is therefore bounded by whatever limit the SDK applies,
    // not by the stream endpoint's. Front this with a gateway limit if you need
    // the two to match.
    await activeTransport.handlePostMessage(req, res);

    return true;
  }

  return false;
};

export const startHTTPServer = async <T extends ServerLike>({
  apiKey,
  authenticate,
  cors,
  createServer,
  enableJsonResponse,
  eventStore,
  eventStoreMaxEvents,
  host = "::",
  keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT,
  maxBodySize,
  oauth,
  onClose,
  onConnect,
  onUnhandledRequest,
  port,
  sseEndpoint = "/sse",
  sslCa,
  sslCert,
  sslKey,
  stateless,
  streamEndpoint = "/mcp",
}: {
  apiKey?: string;
  authenticate?: (request: http.IncomingMessage) => Promise<unknown>;
  cors?: boolean | CorsOptions;
  createServer: (request: http.IncomingMessage) => Promise<T>;
  enableJsonResponse?: boolean;
  /**
   * Event store for the streamable HTTP transport's resumability support.
   * Pass `false` to disable resumability entirely (recommended for
   * request/response-only deployments that don't need replay-on-reconnect).
   * Omit to get a fresh, bounded `InMemoryEventStore` per session - see
   * `eventStoreMaxEvents`. Pass an `EventStore` instance to bring your own
   * (e.g. a persistent, cross-process store); it will be shared across all
   * sessions handled by this server.
   */
  eventStore?: EventStoreOption;
  /**
   * Caps how many events the auto-created per-session `InMemoryEventStore`
   * retains (oldest evicted first) before it's overridden by an explicit
   * `eventStore`. Bounds memory for long-lived sessions. Default: 1000.
   */
  eventStoreMaxEvents?: number;
  host?: string;
  keepAliveTimeout?: number;
  /**
   * Caps how many bytes of a request body the stream endpoint buffers,
   * bounding the memory a single request can consume. A request over the cap
   * is answered with `413 Payload Too Large` and the connection is closed.
   * Default: 10485760 (10 MiB). Pass `false` to disable the cap entirely
   * (unbounded buffering - only safe behind a gateway that already limits body
   * size). Does not apply to the SSE endpoint, whose POST bodies are read by
   * the MCP SDK.
   */
  maxBodySize?: MaxBodySizeOption;
  oauth?: AuthConfig["oauth"];
  onClose?: (server: T) => Promise<void>;
  onConnect?: (server: T) => Promise<void>;
  onUnhandledRequest?: (
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ) => Promise<void>;
  port: number;
  sseEndpoint?: null | string;
  sslCa?: null | string;
  sslCert?: null | string;
  sslKey?: null | string;
  stateless?: boolean;
  streamEndpoint?: null | string;
}): Promise<SSEServer> => {
  const activeSSETransports: Record<string, SSEServerTransport> = {};

  const activeStreamTransports: Record<
    string,
    {
      server: T;
      transport: StreamableHTTPServerTransport;
    }
  > = {};

  const authMiddleware = new AuthenticationMiddleware({ apiKey, oauth });

  /**
   * @author https://dev.classmethod.jp/articles/mcp-sse/
   */
  const requestListener: http.RequestListener = async (req, res) => {
    // Apply CORS headers
    applyCorsHeaders(req, res, cors);

    if (req.method === "OPTIONS") {
      res.writeHead(204);
      res.end();
      return;
    }

    if (req.method === "GET" && req.url === `/ping`) {
      res.writeHead(200).end("pong");
      return;
    }

    // Determine whether the request targets an MCP protocol endpoint (SSE
    // or HTTP Stream). For those endpoints, onUnhandledRequest MUST NOT run
    // first — some consumers (e.g. fastmcp) use it as a catch-all 404 handler
    // and would otherwise short-circuit the MCP protocol handlers.
    // Use a fixed base because `host` may be "::" (IPv6 any), which is not a
    // valid URL authority. We only need pathname here.
    // A malformed request target (e.g. "//") makes `new URL` throw, which
    // would crash the process from this listener, so reject it with 400.
    let requestUrl: URL;
    try {
      requestUrl = new URL(req.url || "", "http://localhost");
    } catch {
      res.writeHead(400).end("Bad Request");
      return;
    }
    const isMcpEndpoint =
      (sseEndpoint && requestUrl.pathname === sseEndpoint) ||
      (streamEndpoint && requestUrl.pathname === streamEndpoint);

    // Let non-MCP routes (e.g. /health, /ready, OAuth metadata) be handled
    // before auth — API key auth protects MCP protocol endpoints, not custom routes.
    if (onUnhandledRequest && !isMcpEndpoint) {
      await onUnhandledRequest(req, res);
      if (res.writableEnded) {
        return;
      }
    }

    // Check authentication for MCP protocol endpoints
    if (!authMiddleware.validateRequest(req)) {
      const authResponse = authMiddleware.getUnauthorizedResponse();
      res.writeHead(401, authResponse.headers);
      res.end(authResponse.body);
      return;
    }

    if (
      sseEndpoint &&
      (await handleSSERequest({
        activeTransports: activeSSETransports,
        createServer,
        endpoint: sseEndpoint,
        onClose,
        onConnect,
        req,
        res,
      }))
    ) {
      return;
    }

    if (
      streamEndpoint &&
      (await handleStreamRequest({
        activeTransports: activeStreamTransports,
        authenticate,
        authMiddleware,
        createServer,
        enableJsonResponse,
        endpoint: streamEndpoint,
        eventStore,
        eventStoreMaxEvents,
        maxBodySize,
        oauth,
        onClose,
        onConnect,
        req,
        res,
        stateless,
      }))
    ) {
      return;
    }

    res.writeHead(404).end();
  };

  let httpServer: http.Server | https.Server;
  if (sslCa || sslCert || sslKey) {
    const options: https.ServerOptions = {};
    if (sslCa) {
      try {
        options.ca = fs.readFileSync(sslCa);
      } catch (error) {
        throw new Error(
          `Failed to read CA file '${sslCa}': ${(error as Error).message}`,
        );
      }
    }
    if (sslCert) {
      try {
        options.cert = fs.readFileSync(sslCert);
      } catch (error) {
        throw new Error(
          `Failed to read certificate file '${sslCert}': ${(error as Error).message}`,
        );
      }
    }
    if (sslKey) {
      try {
        options.key = fs.readFileSync(sslKey);
      } catch (error) {
        throw new Error(
          `Failed to read key file '${sslKey}': ${(error as Error).message}`,
        );
      }
    }
    httpServer = https.createServer(options, requestListener);
  } else {
    httpServer = http.createServer(requestListener);
  }

  // Keep stateful stream sessions from being torn down when Node closes
  // otherwise-idle HTTP keep-alive sockets after its 5 second default.
  httpServer.keepAliveTimeout = keepAliveTimeout;
  httpServer.headersTimeout = Math.max(
    httpServer.headersTimeout,
    keepAliveTimeout + 1000,
  );

  await new Promise((resolve) => {
    httpServer.listen(port, host, () => {
      resolve(undefined);
    });
  });

  return {
    close: async () => {
      for (const transport of Object.values(activeSSETransports)) {
        await transport.close();
      }

      for (const transport of Object.values(activeStreamTransports)) {
        await transport.transport.close();
      }

      return new Promise((resolve, reject) => {
        httpServer.close((error) => {
          if (error) {
            reject(error);

            return;
          }

          resolve();
        });
      });
    },
  };
};