import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { type Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import express, { Request, Response } from "express";
import { randomUUID } from "node:crypto";
import { z } from "zod";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import {
  CallToolResult,
  GetPromptResult,
  isInitializeRequest,
  PrimitiveSchemaDefinition,
  ReadResourceResult,
  Resource,
  ResourceLink,
} from "@modelcontextprotocol/sdk/types.js";
import { InMemoryEventStore } from "@modelcontextprotocol/sdk/examples/shared/inMemoryEventStore.js";

import cors from "cors";
import { text } from "node:stream/consumers";
import path from "node:path";
import fs from "node:fs/promises";

// const isProduction = process.env.NODE_ENV !== 'development';
const exampleUiPath = __filename.endsWith(".ts")
  ? path.join(__dirname, "dist", "example-ui.html")
  : path.join(__dirname, "example-server.html");
const exampleUiHtml = await fs.readFile(exampleUiPath, "utf-8");

// Create an MCP server with implementation details
const getServer = async () => {
  const server = new McpServer(
    {
      name: "example-mcp-ui-server",
      version: "1.0.0",
      icons: [
        { src: "./mcp.svg", sizes: ["512x512"], mimeType: "image/svg+xml" },
      ],
      websiteUrl: "https://github.com/modelcontextprotocol/typescript-sdk",
    },
    { capabilities: { logging: {} } },
  );

  const exampleUiTemplateResource: Resource = {
    name: "example-ui-template",
    uri: "ui://example-ui-template",
    title: "Greeting Template",
    description: "A simple greeting UI",
    mimeType: "text/html",
  };

  server.registerResource(
    exampleUiTemplateResource.name,
    exampleUiTemplateResource.uri,
    exampleUiTemplateResource,
    async (): Promise<ReadResourceResult> => {
      return {
        contents: [
          {
            uri: exampleUiTemplateResource.uri,
            mimeType: exampleUiTemplateResource.mimeType,
            text: exampleUiHtml,
          },
        ],
      };
    },
  );

  server.registerTool(
    "create-example-ui",
    {
      title: "Example UI",
      description: "A tool that returns a simple UI",
      inputSchema: {
        message: z.string().describe("Name to greet"),
      },
      _meta: {
        "ui/resourceUri": exampleUiTemplateResource.uri,
      },
    },
    async ({ message }): Promise<CallToolResult> => {
      const structuredContent = {
        message,
      };
      return {
        // content: [],
        content: [
          {
            type: "text",
            text: JSON.stringify(structuredContent),
          },
        ],
        structuredContent,
      };
    },
  );

  server.registerTool(
    "get-weather",
    {
      title: "Get Weather",
      description:
        "A tool that returns the current weather for a given location",
      inputSchema: {
        location: z.string().describe("Location to get the weather for"),
      },
    },
    async ({ location }): Promise<CallToolResult> => {
      // For simplicity, return a fixed weather report
      const temperature = 25; // Fixed temperature
      const condition = "sunny";
      const weatherReport = `The current weather in ${location} is ${condition} with a temperature of ${temperature}°C.`;
      const structuredContent = {
        temperature,
        condition,
      };
      return {
        content: [
          {
            type: "text",
            text: weatherReport,
          },
        ],
        structuredContent,
      };
    },
  );

  return server;
};

const MCP_PORT = process.env.MCP_PORT
  ? parseInt(process.env.MCP_PORT, 10)
  : 3001;
const MCP_HOST = process.env.MCP_HOST || "localhost";

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

// Allow CORS all domains, expose the Mcp-Session-Id header
app.use(
  cors({
    origin: "*", // Allow all origins
    exposedHeaders: ["Mcp-Session-Id"],
  }),
);

// Map to store transports by session ID
const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};

const mcpPostHandler = async (req: Request, res: Response) => {
  const sessionId = req.headers["mcp-session-id"] as string | undefined;
  if (sessionId) {
    console.log(`Received MCP request for session: ${sessionId}`);
  } else {
    console.log("Request body:", req.body);
  }

  try {
    let transport: StreamableHTTPServerTransport;
    if (sessionId && transports[sessionId]) {
      // Reuse existing transport
      transport = transports[sessionId];
    } else if (!sessionId && isInitializeRequest(req.body)) {
      // New initialization request
      const eventStore = new InMemoryEventStore();
      transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => randomUUID(),
        eventStore, // Enable resumability
        onsessioninitialized: (sessionId) => {
          // Store the transport by session ID when session is initialized
          // This avoids race conditions where requests might come in before the session is stored
          console.log(`Session initialized with ID: ${sessionId}`);
          transports[sessionId] = transport;
        },
      });

      // Set up onclose handler to clean up transport when closed
      transport.onclose = () => {
        const sid = transport.sessionId;
        if (sid && transports[sid]) {
          console.log(
            `Transport closed for session ${sid}, removing from transports map`,
          );
          delete transports[sid];
        }
      };

      // Connect the transport to the MCP server BEFORE handling the request
      // so responses can flow back through the same transport
      const server = await getServer();
      await server.connect(transport);

      await transport.handleRequest(req, res, req.body);
      return; // Already handled
    } else {
      // Invalid request - no session ID or not initialization request
      res.status(400).json({
        jsonrpc: "2.0",
        error: {
          code: -32000,
          message: "Bad Request: No valid session ID provided",
        },
        id: null,
      });
      return;
    }

    // Handle the request with existing transport - no need to reconnect
    // The existing transport is already connected to the server
    await transport.handleRequest(req, res, req.body);
  } catch (error) {
    console.error("Error handling MCP request:", error);
    if (!res.headersSent) {
      res.status(500).json({
        jsonrpc: "2.0",
        error: {
          code: -32603,
          message: "Internal server error",
        },
        id: null,
      });
    }
  }
};

app.post("/mcp", mcpPostHandler);

// Handle GET requests for SSE streams (using built-in support from StreamableHTTP)
const mcpGetHandler = async (req: Request, res: Response) => {
  const sessionId = req.headers["mcp-session-id"] as string | undefined;
  if (!sessionId || !transports[sessionId]) {
    res.status(400).send("Invalid or missing session ID");
    return;
  }

  // Check for Last-Event-ID header for resumability
  const lastEventId = req.headers["last-event-id"] as string | undefined;
  if (lastEventId) {
    console.log(`Client reconnecting with Last-Event-ID: ${lastEventId}`);
  } else {
    console.log(`Establishing new SSE stream for session ${sessionId}`);
  }

  const transport = transports[sessionId];
  await transport.handleRequest(req, res);
};

app.get("/mcp", mcpGetHandler);

// Handle DELETE requests for session termination (according to MCP spec)
const mcpDeleteHandler = async (req: Request, res: Response) => {
  const sessionId = req.headers["mcp-session-id"] as string | undefined;
  if (!sessionId || !transports[sessionId]) {
    res.status(400).send("Invalid or missing session ID");
    return;
  }

  console.log(`Received session termination request for session ${sessionId}`);

  try {
    const transport = transports[sessionId];
    await transport.handleRequest(req, res);
  } catch (error) {
    console.error("Error handling session termination:", error);
    if (!res.headersSent) {
      res.status(500).send("Error processing session termination");
    }
  }
};

app.delete("/mcp", mcpDeleteHandler);

app.listen(MCP_PORT, (error) => {
  if (error) {
    console.error("Failed to start server:", error);
    process.exit(1);
  }
  console.log(`MCP Streamable HTTP Server listening on port ${MCP_PORT}`);
});

// Handle server shutdown
process.on("SIGINT", async () => {
  console.log("Shutting down server...");

  // Close all active transports to properly clean up resources
  for (const sessionId in transports) {
    try {
      console.log(`Closing transport for session ${sessionId}`);
      await transports[sessionId].close();
      delete transports[sessionId];
    } catch (error) {
      console.error(`Error closing transport for session ${sessionId}:`, error);
    }
  }
  console.log("Server shutdown complete");
  process.exit(0);
});