// @ts-nocheck – The MCP SDK's deeply recursive zod v4 types cause tsc to
// stack-overflow (SIGABRT) during type-checking.  Runtime tests are unaffected.
import { spawn } from "child_process";
import type http from "http";
import type { AddressInfo } from "net";
import path from "path";

import * as MCPClientModule from "@modelcontextprotocol/sdk/client/index.js";
import * as MCPClientSSEModule from "@modelcontextprotocol/sdk/client/sse.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import * as MCPClientStdioModule from "@modelcontextprotocol/sdk/client/stdio.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import * as MCPClientStreamableHTTPModule from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import * as MCPServerSSEModule from "@modelcontextprotocol/sdk/server/sse.js";
import * as MCPServerStdioModule from "@modelcontextprotocol/sdk/server/stdio.js";
import * as MCPServerStreamableHTTPModule from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import type { Tracer } from "@opentelemetry/api";
import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-base";
import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest";
import { z } from "zod/v3";

import { isPatched, MCPInstrumentation } from "../src";
import type { ExportedSpan } from "./collector";
import { startCollector, Telemetry } from "./collector";

describe("MCPInstrumentation", () => {
  const telemetry = new Telemetry();

  let collector: http.Server;
  let tracerProvider: NodeTracerProvider;
  let tracer: Tracer;
  let otlpEndpoint: string;

  const instrumentation = new MCPInstrumentation();
  instrumentation.disable();

  beforeAll(async () => {
    collector = await startCollector(telemetry);
    otlpEndpoint = `http://localhost:${(collector.address() as AddressInfo).port}`;
    const exporter = new OTLPTraceExporter({
      url: `${otlpEndpoint}/v1/traces`,
    });
    tracerProvider = new NodeTracerProvider({
      spanProcessors: [new BatchSpanProcessor(exporter)],
    });
    tracerProvider.register();
    tracer = tracerProvider.getTracer("mcp-test-client");

    instrumentation.setTracerProvider(tracerProvider);
    instrumentation.enable();

    instrumentation.manuallyInstrument({
      clientSSEModule: MCPClientSSEModule,
      serverSSEModule: MCPServerSSEModule,
      clientStdioModule: MCPClientStdioModule,
      serverStdioModule: MCPServerStdioModule,
      clientStreamableHTTPModule: MCPClientStreamableHTTPModule,
      serverStreamableHTTPModule: MCPServerStreamableHTTPModule,
    });
  });

  afterAll(async () => {
    instrumentation.disable();
    await tracerProvider.shutdown();

    await new Promise<void>((resolve, reject) => {
      collector.close((err) => {
        if (err) {
          reject(err);
        } else {
          resolve();
        }
      });
    });
  });

  beforeEach(() => {
    telemetry.clear();
  });
  it("is patched", () => {
    expect((MCPClientSSEModule as { openInferencePatched?: boolean }).openInferencePatched).toBe(
      true,
    );
    expect(isPatched("@modelcontextprotocol/sdk/client/sse")).toBe(true);
    expect((MCPServerSSEModule as { openInferencePatched?: boolean }).openInferencePatched).toBe(
      true,
    );
    expect(isPatched("@modelcontextprotocol/sdk/server/sse")).toBe(true);
    expect((MCPClientStdioModule as { openInferencePatched?: boolean }).openInferencePatched).toBe(
      true,
    );
    expect(isPatched("@modelcontextprotocol/sdk/client/stdio")).toBe(true);
    expect((MCPServerStdioModule as { openInferencePatched?: boolean }).openInferencePatched).toBe(
      true,
    );
    expect(isPatched("@modelcontextprotocol/sdk/server/stdio")).toBe(true);
    expect(
      (MCPClientStreamableHTTPModule as { openInferencePatched?: boolean }).openInferencePatched,
    ).toBe(true);
    expect(isPatched("@modelcontextprotocol/sdk/client/streamableHttp")).toBe(true);
    expect(
      (MCPServerStreamableHTTPModule as { openInferencePatched?: boolean }).openInferencePatched,
    ).toBe(true);
    expect(isPatched("@modelcontextprotocol/sdk/server/streamableHttp")).toBe(true);
  });
  it("propagates context - stdio", { timeout: 30_000 }, async () => {
    const transport = new StdioClientTransport({
      command: "tsx",
      args: [path.join(__dirname, "mcpserver.ts")],
      env: {
        ...process.env,
        MCP_TRANSPORT: "stdio",
        OTEL_EXPORTER_OTLP_ENDPOINT: otlpEndpoint,
      },
    });
    const client = new MCPClientModule.Client({
      name: "MCP",
      version: "0.1.0",
    });
    // @ts-expect-error TS2589: MCP SDK's deeply nested Zod generics exceed TypeScript's recursion limit
    client.setRequestHandler(z.object({ method: z.literal("whoami") }), () => {
      return tracer.startActiveSpan("whoami", (span) => {
        try {
          return {
            name: "OpenInference",
          };
        } finally {
          span.end();
        }
      });
    });
    await client.connect(transport);
    try {
      await tracer.startActiveSpan("root", async (span) => {
        try {
          const { tools } = await client.listTools();
          expect(tools).toHaveLength(1);
          expect(tools[0].name).toBe("hello");

          const res = (await client.callTool({ name: "hello" })) as {
            content: Array<{ type: string; text: string }>;
          };
          expect(res.content).toHaveLength(1);
          expect(res.content[0].text).toBe("Hello OpenInference!");
        } finally {
          span.end();
        }
      });
    } finally {
      await client.close();
    }
    await tracerProvider.forceFlush();
    // There does not seem to be a reliable way to wait for the child to exit completely with StdioClientTransport
    // so poll for the expected telemetry length instead.
    for (let i = 0; i < 500; i++) {
      if (telemetry.resourceSpans.length >= 2) {
        break;
      }
      await new Promise((resolve) => setTimeout(resolve, 50));
    }
    expect(telemetry.resourceSpans).toHaveLength(2);
    let rootSpan: ExportedSpan | undefined;
    let serverSpan: ExportedSpan | undefined;
    let whoamiSpan: ExportedSpan | undefined;
    for (const resourceSpan of telemetry.resourceSpans) {
      expect(resourceSpan.scopeSpans).toHaveLength(1);
      for (const scopeSpan of resourceSpan.scopeSpans) {
        if (scopeSpan.scope.name === "mcp-test-client") {
          for (const span of scopeSpan.spans) {
            if (span.name === "root") {
              rootSpan = span;
            } else if (span.name === "whoami") {
              whoamiSpan = span;
            }
          }
        } else if (scopeSpan.scope.name === "mcp-test-server") {
          serverSpan = scopeSpan.spans[0];
        }
      }
    }
    expect(rootSpan).toBeDefined();
    expect(rootSpan!.name).toBe("root");
    expect(serverSpan).toBeDefined();
    expect(serverSpan!.name).toBe("hello");
    expect(whoamiSpan).toBeDefined();
    expect(whoamiSpan!.name).toBe("whoami");
    expect(serverSpan!.traceId).toBe(rootSpan!.traceId);
    expect(serverSpan!.parentSpanId).toBe(rootSpan!.spanId);
    expect(whoamiSpan!.traceId).toBe(rootSpan!.traceId);
    expect(whoamiSpan!.parentSpanId).toBe(serverSpan!.spanId);
  });
  it("propagates context - sse", async () => {
    const child = spawn("tsx", [path.join(__dirname, "mcpserver.ts")], {
      env: {
        ...process.env,
        MCP_TRANSPORT: "sse",
        OTEL_EXPORTER_OTLP_ENDPOINT: otlpEndpoint,
      },
      stdio: ["ignore", "pipe", "inherit"],
    });
    const childClosedPromise = new Promise<void>((resolve) => {
      child.on("close", () => {
        resolve();
      });
    });
    const sseUrl = await new Promise<string>((resolve) => {
      child.stdout.on("data", (data: Buffer) => {
        const line = data.toString().trim();
        if (line.startsWith("Server running on ")) {
          const endpoint = line.substring("Server running on ".length);
          resolve(endpoint);
        }
      });
    });
    try {
      const transport = new SSEClientTransport(new URL(sseUrl));
      const client = new MCPClientModule.Client({
        name: "MCP",
        version: "0.1.0",
      });
      // @ts-expect-error TS2589: MCP SDK's deeply nested Zod generics exceed TypeScript's recursion limit
      client.setRequestHandler(z.object({ method: z.literal("whoami") }), () => {
        return tracer.startActiveSpan("whoami", (span) => {
          try {
            return {
              name: "OpenInference",
            };
          } finally {
            span.end();
          }
        });
      });
      await client.connect(transport);
      try {
        await tracer.startActiveSpan("root", async (span) => {
          try {
            const { tools } = await client.listTools();
            expect(tools).toHaveLength(1);
            expect(tools[0].name).toBe("hello");

            const res = (await client.callTool({ name: "hello" })) as {
              content: Array<{ type: string; text: string }>;
            };
            expect(res.content).toHaveLength(1);
            expect(res.content[0].text).toBe("Hello OpenInference!");
          } finally {
            span.end();
          }
        });
      } finally {
        await client.close();
      }
    } finally {
      child.kill();
    }
    await tracerProvider.forceFlush();
    await childClosedPromise;
    expect(telemetry.resourceSpans).toHaveLength(2);
    let rootSpan: ExportedSpan | undefined;
    let serverSpan: ExportedSpan | undefined;
    let whoamiSpan: ExportedSpan | undefined;
    for (const resourceSpan of telemetry.resourceSpans) {
      expect(resourceSpan.scopeSpans).toHaveLength(1);
      for (const scopeSpan of resourceSpan.scopeSpans) {
        if (scopeSpan.scope.name === "mcp-test-client") {
          for (const span of scopeSpan.spans) {
            if (span.name === "root") {
              rootSpan = span;
            } else if (span.name === "whoami") {
              whoamiSpan = span;
            }
          }
        } else if (scopeSpan.scope.name === "mcp-test-server") {
          serverSpan = scopeSpan.spans[0];
        }
      }
    }
    expect(rootSpan).toBeDefined();
    expect(rootSpan!.name).toBe("root");
    expect(serverSpan).toBeDefined();
    expect(serverSpan!.name).toBe("hello");
    expect(whoamiSpan).toBeDefined();
    expect(whoamiSpan!.name).toBe("whoami");
    expect(serverSpan!.traceId).toBe(rootSpan!.traceId);
    expect(serverSpan!.parentSpanId).toBe(rootSpan!.spanId);
    expect(whoamiSpan!.traceId).toBe(rootSpan!.traceId);
    expect(whoamiSpan!.parentSpanId).toBe(serverSpan!.spanId);
  });
  it("propagates context - streamable http", async () => {
    const child = spawn("tsx", [path.join(__dirname, "mcpserver.ts")], {
      env: {
        ...process.env,
        MCP_TRANSPORT: "streamableHttp",
        OTEL_EXPORTER_OTLP_ENDPOINT: otlpEndpoint,
      },
      stdio: ["ignore", "pipe", "inherit"],
    });
    const childClosedPromise = new Promise<void>((resolve) => {
      child.on("close", () => {
        resolve();
      });
    });
    const serverUrl = await new Promise<string>((resolve) => {
      child.stdout.on("data", (data: Buffer) => {
        const line = data.toString().trim();
        if (line.startsWith("Server running on ")) {
          const endpoint = line.substring("Server running on ".length);
          resolve(endpoint);
        }
      });
    });
    try {
      const transport = new StreamableHTTPClientTransport(new URL(serverUrl));
      const client = new MCPClientModule.Client({
        name: "MCP",
        version: "0.1.0",
      });
      // @ts-expect-error TS2589: MCP SDK's deeply nested Zod generics exceed TypeScript's recursion limit
      client.setRequestHandler(z.object({ method: z.literal("whoami") }), () => {
        return tracer.startActiveSpan("whoami", (span) => {
          try {
            return {
              name: "OpenInference",
            };
          } finally {
            span.end();
          }
        });
      });
      await client.connect(transport);
      try {
        await tracer.startActiveSpan("root", async (span) => {
          try {
            const { tools } = await client.listTools();
            expect(tools).toHaveLength(1);
            expect(tools[0].name).toBe("hello");

            const res = (await client.callTool({ name: "hello" })) as {
              content: Array<{ type: string; text: string }>;
            };
            expect(res.content).toHaveLength(1);
            expect(res.content[0].text).toBe("Hello OpenInference!");
          } finally {
            span.end();
          }
        });
      } finally {
        await client.close();
      }
    } finally {
      child.kill();
    }
    await tracerProvider.forceFlush();
    await childClosedPromise;
    expect(telemetry.resourceSpans).toHaveLength(2);
    let rootSpan: ExportedSpan | undefined;
    let serverSpan: ExportedSpan | undefined;
    let whoamiSpan: ExportedSpan | undefined;
    for (const resourceSpan of telemetry.resourceSpans) {
      expect(resourceSpan.scopeSpans).toHaveLength(1);
      for (const scopeSpan of resourceSpan.scopeSpans) {
        if (scopeSpan.scope.name === "mcp-test-client") {
          for (const span of scopeSpan.spans) {
            if (span.name === "root") {
              rootSpan = span;
            } else if (span.name === "whoami") {
              whoamiSpan = span;
            }
          }
        } else if (scopeSpan.scope.name === "mcp-test-server") {
          serverSpan = scopeSpan.spans[0];
        }
      }
    }
    expect(rootSpan).toBeDefined();
    expect(rootSpan!.name).toBe("root");
    expect(serverSpan).toBeDefined();
    expect(serverSpan!.name).toBe("hello");
    expect(whoamiSpan).toBeDefined();
    expect(whoamiSpan!.name).toBe("whoami");
    expect(serverSpan!.traceId).toBe(rootSpan!.traceId);
    expect(serverSpan!.parentSpanId).toBe(rootSpan!.spanId);
    expect(whoamiSpan!.traceId).toBe(rootSpan!.traceId);
    expect(whoamiSpan!.parentSpanId).toBe(serverSpan!.spanId);
  });
});
