server.ts

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import type {
  CallToolResult,
  ReadResourceResult,
} from "@modelcontextprotocol/sdk/types.js";
import fs from "node:fs/promises";
import path from "node:path";
import { z } from "zod";
import {
  RESOURCE_MIME_TYPE,
  registerAppResource,
  registerAppTool,
} from "@modelcontextprotocol/ext-apps/server";
import {
  generateCustomers,
  generateSegmentSummaries,
} from "./src/data-generator.js";
import { SEGMENTS, type Customer, type SegmentSummary } from "./src/types.js";

// Works both from source (server.ts) and compiled (dist/server.js)
const DIST_DIR = import.meta.filename.endsWith(".ts")
  ? path.join(import.meta.dirname, "dist")
  : import.meta.dirname;

// Schemas - types are derived from these using z.infer
const GetCustomerDataInputSchema = z.object({
  segment: z
    .enum(["All", ...SEGMENTS])
    .optional()
    .describe("Filter by segment (default: All)"),
});

const CustomerSchema = z.object({
  id: z.string(),
  name: z.string(),
  segment: z.string(),
  annualRevenue: z.number(),
  employeeCount: z.number(),
  accountAge: z.number(),
  engagementScore: z.number(),
  supportTickets: z.number(),
  nps: z.number(),
});

const SegmentSummarySchema = z.object({
  name: z.string(),
  count: z.number(),
  color: z.string(),
});

const GetCustomerDataOutputSchema = z.object({
  customers: z.array(CustomerSchema),
  segments: z.array(SegmentSummarySchema),
});

// Cache generated data for session consistency
let cachedCustomers: Customer[] | null = null;
let cachedSegments: SegmentSummary[] | null = null;

function getCustomerData(segmentFilter?: string): {
  customers: Customer[];
  segments: SegmentSummary[];
} {
  // Generate data on first call
  if (!cachedCustomers) {
    cachedCustomers = generateCustomers(250);
    cachedSegments = generateSegmentSummaries(cachedCustomers);
  }

  // Filter by segment if specified
  let customers = cachedCustomers;
  if (segmentFilter && segmentFilter !== "All") {
    customers = cachedCustomers.filter((c) => c.segment === segmentFilter);
  }

  return {
    customers,
    segments: cachedSegments!,
  };
}

/**
 * Creates a new MCP server instance with tools and resources registered.
 * Each HTTP session needs its own server instance because McpServer only supports one transport.
 */
export function createServer(): McpServer {
  const server = new McpServer({
    name: "Customer Segmentation Server",
    version: "1.0.0",
  });

  // Register the get-customer-data tool and its associated UI resource
  {
    const resourceUri = "ui://customer-segmentation/mcp-app.html";

    registerAppTool(
      server,
      "get-customer-data",
      {
        title: "Get Customer Data",
        description:
          "Returns customer data with segment information for visualization. Optionally filter by segment.",
        inputSchema: GetCustomerDataInputSchema.shape,
        outputSchema: GetCustomerDataOutputSchema.shape,
        _meta: { ui: { resourceUri } },
      },
      async ({ segment }): Promise<CallToolResult> => {
        const data = getCustomerData(segment);

        return {
          content: [{ type: "text", text: JSON.stringify(data) }],
          structuredContent: data,
        };
      },
    );

    registerAppResource(
      server,
      resourceUri,
      resourceUri,
      {
        mimeType: RESOURCE_MIME_TYPE,
        description: "Customer Segmentation Explorer UI",
      },
      async (): Promise<ReadResourceResult> => {
        const html = await fs.readFile(
          path.join(DIST_DIR, "mcp-app.html"),
          "utf-8",
        );

        return {
          contents: [
            {
              uri: resourceUri,
              mimeType: RESOURCE_MIME_TYPE,
              text: html,
            },
          ],
        };
      },
    );
  }

  return server;
}