/*
 * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.
 */

import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { localhostHostValidation } from "@modelcontextprotocol/sdk/server/middleware/hostHeaderValidation.js";
import express from "express";
import { SDKOptions } from "../../../lib/config.js";
import { LocalContext } from "../../cli.js";
import {
  ConsoleLoggerLevel,
  createConsoleLogger,
} from "../../console-logger.js";
import { MCPScope } from "../../scopes.js";
import { createMCPServer } from "../../server.js";
import {
  assertSafeSSEBind,
  isLoopbackHost,
  requireBearerToken,
} from "./sse-security.js";
import {
  installTlsBlocklist,
  unpatchedRuntimeWarning,
} from "./tls-blocklist.js";

interface StartCommandFlags {
  readonly transport: "stdio" | "sse";
  readonly port: number;
  readonly host: string;
  readonly "auth-token"?: string | undefined;
  readonly tool?: string[];
  readonly scope?: MCPScope[];
  readonly "api-key"?: string | undefined;
  readonly "server-url"?: string;
  readonly "server-index"?: SDKOptions["serverIdx"];
  readonly "log-level": ConsoleLoggerLevel;
  readonly env?: [string, string][];
}

export async function main(this: LocalContext, flags: StartCommandFlags) {
  // INC-289: block the compromised eMudhra intermediates before any request.
  installTlsBlocklist();
  const runtimeWarning = unpatchedRuntimeWarning();
  if (runtimeWarning) {
    process.stderr.write(runtimeWarning);
  }

  flags.env?.forEach(([key, value]) => {
    process.env[key] = value;
  });

  switch (flags.transport) {
    case "stdio":
      await startStdio(flags);
      break;
    case "sse":
      await startSSE(flags);
      break;
    default:
      throw new Error(`Invalid transport: ${flags.transport}`);
  }
}

async function startStdio(flags: StartCommandFlags) {
  const logger = createConsoleLogger(flags["log-level"]);
  const transport = new StdioServerTransport();
  const server = createMCPServer({
    logger,
    allowedTools: flags.tool,
    scopes: flags.scope,
    ...{ apiKey: flags["api-key"] ?? "" },
    serverURL: flags["server-url"],
    serverIdx: flags["server-index"],
  });
  await server.connect(transport);

  const abort = async () => {
    await server.close();
    process.exit(0);
  };
  process.on("SIGTERM", abort);
  process.on("SIGINT", abort);
}

async function startSSE(flags: StartCommandFlags) {
  const logger = createConsoleLogger(flags["log-level"]);
  const host = flags.host;
  const authToken = flags["auth-token"];
  assertSafeSSEBind(host, authToken);

  const app = express();
  if (isLoopbackHost(host)) {
    app.use(localhostHostValidation());
  }
  if (authToken) {
    app.use(requireBearerToken(authToken));
  }

  const sessions = new Map<
    string,
    { transport: SSEServerTransport; server: ReturnType<typeof createMCPServer> }
  >();
  const controller = new AbortController();

  app.get("/sse", async (_req, res) => {
    const server = createMCPServer({
      logger,
      allowedTools: flags.tool,
      scopes: flags.scope,
      ...{ apiKey: flags["api-key"] ?? "" },
      serverURL: flags["server-url"],
      serverIdx: flags["server-index"],
    });
    const transport = new SSEServerTransport("/message", res);
    sessions.set(transport.sessionId, { transport, server });

    res.on("close", () => {
      sessions.delete(transport.sessionId);
      void server.close();
    });

    try {
      await server.connect(transport);
    } catch (err) {
      logger.error("Failed to open SSE session", {
        error: err instanceof Error ? err.message : String(err),
      });
      sessions.delete(transport.sessionId);
      void server.close();
      if (!res.headersSent) {
        res.status(500).send("Internal Server Error");
      }
    }
  });

  app.post("/message", async (req, res) => {
    const sessionId = typeof req.query["sessionId"] === "string"
      ? req.query["sessionId"]
      : "";
    const session = sessions.get(sessionId);
    if (!session) {
      res.status(404).send("Unknown session");
      return;
    }

    try {
      await session.transport.handlePostMessage(req, res);
    } catch (err) {
      logger.error("Failed to handle SSE message", {
        error: err instanceof Error ? err.message : String(err),
      });
      if (!res.headersSent) {
        res.status(500).send("Internal Server Error");
      }
    }
  });

  const httpServer = app.listen(flags.port, host, () => {
    const ha = httpServer.address();
    const host = typeof ha === "string" ? ha : `${ha?.address}:${ha?.port}`;
    logger.info("MCP HTTP server started", { host });
  });

  let closing = false;
  controller.signal.addEventListener("abort", async () => {
    if (closing) {
      logger.info("Received second signal. Forcing shutdown.");
      process.exit(1);
    }
    closing = true;

    logger.info("Shutting down MCP servers");

    await Promise.all([...sessions.values()].map(({ server }) => server.close()));

    logger.info("Shutting down HTTP server");

    const timer = setTimeout(() => {
      logger.info("Forcing shutdown");
      process.exit(1);
    }, 5000);

    httpServer.close(() => {
      clearTimeout(timer);
      logger.info("Graceful shutdown complete");
      process.exit(0);
    });
  });

  const abort = () => controller.abort();
  process.on("SIGTERM", abort);
  process.on("SIGINT", abort);
}