/*
 * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.
 * @generated-id: 37494f674015
 */

import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import express from "express";
import { LocalContext } from "../../cli.js";
import { resolveMcpApps } from "../../apps/config.js";
import {
  ConsoleLoggerLevel,
  createConsoleLogger,
} from "../../console-logger.js";
import { MCPServerFlags } from "../../flags.js";
import { createMCPServer } from "../../server.js";
import { buildAnnotationFilter } from "../../tools.js";

import { landingPageExpress } from "../../../landing-page.js";

interface StartCommandFlags extends MCPServerFlags {
  readonly transport: "stdio" | "sse";
  readonly port: number;
  readonly "log-level": ConsoleLoggerLevel;
  readonly env?: [string, string][];
}

interface Session {
  mcpServer: ReturnType<typeof createMCPServer>["server"];
  transport: SSEServerTransport;
  cleanup: () => Promise<void>;
}

export async function main(this: LocalContext, flags: StartCommandFlags) {
  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,
    dynamic: flags.mode === "dynamic",
    annotationFilter: buildAnnotationFilter(flags["tool-annotations"]),
    scopes: flags.scope,
    security: {
      cloudinaryAuth: flags["api-key"] != null && flags["api-secret"] != null
        ? {
          api_key: flags["api-key"] ?? "",
          api_secret: flags["api-secret"] ?? "",
        }
        : void 0,
      oauth2: flags.oauth2 ?? "",
    },
    cloud_name: flags["cloud-name"],
    serverURL: flags["server-url"],
    serverIdx: flags["server-index"],
    region: flags.region,
    host: flags["api-host"],
    mcpApps: resolveMcpApps(flags["mcp-apps"]),
  });
  await server.connect(transport);

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

async function startSSE(cliFlags: StartCommandFlags) {
  const logger = createConsoleLogger(cliFlags["log-level"]);
  const app = express();
  const sessions = new Map<string, Session>();
  const controller = new AbortController();

  app.get("/sse", async (req, res) => {
    const sessionId = crypto.randomUUID();

    logger.info("SSE connection initiated", {
      sessionId,
      userAgent: req.headers["user-agent"],
    });

    // Track early disconnection
    let connectionClosed = false;
    res.on("close", () => {
      if (!connectionClosed) {
        logger.info("Connection closed early", { sessionId });
      }
      connectionClosed = true;
    });

    // Merge CLI flags with header overrides for security credentials
    const flags: StartCommandFlags = {
      ...cliFlags,
      // Security fields can be overridden via headers
      "api-key": (req.headers["api_key"] as string) ?? cliFlags["api-key"],
      "api-secret": (req.headers["api_secret"] as string)
        ?? cliFlags["api-secret"],
      "oauth2": (req.headers["oauth2"] as string) ?? cliFlags["oauth2"],
    };

    // Create a new MCP server for this connection with its auth
    const { server: mcpServer } = createMCPServer({
      logger,
      allowedTools: flags.tool,
      dynamic: flags.mode === "dynamic",
      annotationFilter: buildAnnotationFilter(flags["tool-annotations"]),
      scopes: flags.scope,
      security: {
        cloudinaryAuth: flags["api-key"] != null && flags["api-secret"] != null
          ? {
            api_key: flags["api-key"] ?? "",
            api_secret: flags["api-secret"] ?? "",
          }
          : void 0,
        oauth2: flags.oauth2 ?? "",
      },
      cloud_name: flags["cloud-name"],
      serverURL: flags["server-url"],
      serverIdx: flags["server-index"],
      region: flags.region,
      host: flags["api-host"],
      mcpApps: resolveMcpApps(flags["mcp-apps"]),
    });

    // Message path includes session ID for routing
    const transport = new SSEServerTransport(`/message/${sessionId}`, res);

    const cleanup = async () => {
      await mcpServer.close();
    };

    sessions.set(sessionId, { mcpServer, transport, cleanup });
    logger.info("Session created", { sessionId, totalSessions: sessions.size });
    controller.signal.addEventListener("abort", cleanup);

    // Log server-side errors for debugging (e.g. tools/list handler failures)
    mcpServer.server.onerror = (error) => {
      logger.error("MCP protocol error", {
        sessionId,
        error: error instanceof Error ? error.message : String(error),
        stack: error instanceof Error ? error.stack : undefined,
      });
    };

    // Wrap transport.send to log error responses for debugging
    const originalSend = transport.send.bind(transport);
    transport.send = async (message) => {
      if ("error" in message && message.error) {
        logger.error("MCP sending error response", {
          sessionId,
          id: message.id,
          errorCode: message.error.code,
          errorMessage: message.error.message,
        });
      }
      return originalSend(message);
    };

    await mcpServer.connect(transport);

    mcpServer.server.onclose = async () => {
      connectionClosed = true;
      controller.signal.removeEventListener("abort", cleanup);
      sessions.delete(sessionId);
      logger.info("Session closed", {
        sessionId,
        totalSessions: sessions.size,
      });
      res.end();
    };
  });

  app.post("/message/:sessionId", async (req, res) => {
    const session = sessions.get(req.params.sessionId);
    if (!session) {
      res.status(404).send("Session not found");
      return;
    }

    await session.transport.handlePostMessage(req, res);
  });

  app.get("/", landingPageExpress);

  const httpServer = app.listen(cliFlags.port, "0.0.0.0", () => {
    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 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);
}