import { test } from "node:test";
import assert from "node:assert/strict";
import http from "node:http";
import { AddressInfo } from "node:net";

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";

import { buildGateway, SEP, type GatewayOptions } from "../src/gateway.js";
import type { EnabledAgent } from "../src/config.js";

/**
 * Fake upstream A2A agent: speaks the JSON-RPC subset the gateway uses
 * (initialize, tools/list, tools/call). One handler per server instance.
 */
function startFakeAgent(opts: {
  tools: Array<{ name: string; description: string; inputSchema: any; outputSchema?: any }>;
  onCall: (name: string, args: any) => any;
  authHeader?: string;
}): Promise<{ url: string; close: () => Promise<void>; requests: any[] }> {
  const requests: any[] = [];
  return new Promise((resolve) => {
    const server = http.createServer((req, res) => {
      if (req.method !== "POST" || req.url !== "/mcp") {
        res.statusCode = 404;
        res.end();
        return;
      }
      const auth = req.headers["authorization"];
      const chunks: Buffer[] = [];
      req.on("data", (c) => chunks.push(c));
      req.on("end", () => {
        const body = JSON.parse(Buffer.concat(chunks).toString("utf8"));
        requests.push({ auth, body });
        let result: any;
        if (body.method === "tools/list") {
          result = { tools: opts.tools };
        } else if (body.method === "tools/call") {
          try {
            result = opts.onCall(body.params.name, body.params.arguments);
          } catch (err) {
            res.setHeader("content-type", "application/json");
            res.end(
              JSON.stringify({
                jsonrpc: "2.0",
                id: body.id,
                error: { code: -32603, message: (err as Error).message },
              }),
            );
            return;
          }
        } else {
          result = {};
        }
        res.setHeader("content-type", "application/json");
        res.end(JSON.stringify({ jsonrpc: "2.0", id: body.id, result }));
      });
    });
    server.listen(0, "127.0.0.1", () => {
      const { port } = server.address() as AddressInfo;
      resolve({
        url: `http://127.0.0.1:${port}`,
        requests,
        close: () => new Promise((r) => server.close(() => r())),
      });
    });
  });
}

function startSseAgent(opts: {
  onCall: (body: any, res: http.ServerResponse) => void;
}): Promise<{ url: string; close: () => Promise<void>; requests: any[] }> {
  const requests: any[] = [];
  return new Promise((resolve) => {
    const server = http.createServer((req, res) => {
      if (req.method !== "POST" || req.url !== "/mcp") {
        res.statusCode = 404;
        res.end();
        return;
      }
      const chunks: Buffer[] = [];
      req.on("data", (c) => chunks.push(c));
      req.on("end", () => {
        const body = JSON.parse(Buffer.concat(chunks).toString("utf8"));
        requests.push({ body });
        if (body.method !== "tools/call") {
          res.setHeader("content-type", "application/json");
          res.end(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: {} }));
          return;
        }
        res.writeHead(200, { "content-type": "text/event-stream" });
        opts.onCall(body, res);
      });
    });
    server.listen(0, "127.0.0.1", () => {
      const { port } = server.address() as AddressInfo;
      resolve({
        url: `http://127.0.0.1:${port}`,
        requests,
        close: () => new Promise((r) => server.close(() => r())),
      });
    });
  });
}

/** Drive an in-process gateway Server via the SDK's request handlers. */
async function callGateway(
  agents: EnabledAgent[],
  token: string | null,
  method: "tools/list" | "tools/call",
  params: any,
  opts: {
    gateway?: Partial<Omit<GatewayOptions, "agents" | "token" | "server">>;
    extra?: any;
  } = {},
) {
  const mcpServer = new Server(
    { name: "test", version: "0.0.0" },
    { capabilities: { tools: { listChanged: false } } },
  );
  await buildGateway({ agents, token, server: mcpServer, ...opts.gateway });
  // Reach into the server's registered handlers — the SDK exposes them via
  // a private map, so we route through the public schema-keyed lookup we
  // know is wired up.
  const schema = method === "tools/list" ? ListToolsRequestSchema : CallToolRequestSchema;
  const handler = (mcpServer as any)._requestHandlers.get(schema.shape.method.value);
  if (!handler) throw new Error("handler not registered");
  return handler({ method, params }, opts.extra ?? ({} as any));
}

test("tools/list aggregates upstream tools with agent prefix", async () => {
  const upstream = await startFakeAgent({
    tools: [
      {
        name: "greet",
        description: "Greet someone",
        inputSchema: { type: "object", properties: { who: { type: "string" } } },
      },
    ],
    onCall: () => ({ content: [{ type: "text", text: "hi" }], isError: false }),
  });
  try {
    const out = await callGateway(
      [{ name: "greeter", url: upstream.url, addedAt: "now" }],
      null,
      "tools/list",
      {},
    );
    assert.equal(out.tools.length, 1);
    assert.equal(out.tools[0].name, `greeter${SEP}greet`);
    assert.match(out.tools[0].description, /\[greeter\]/);
  } finally {
    await upstream.close();
  }
});

test("tools/list survives a single upstream failure", async () => {
  const good = await startFakeAgent({
    tools: [
      {
        name: "ping",
        description: "p",
        inputSchema: { type: "object", properties: {} },
      },
    ],
    onCall: () => ({ content: [], isError: false }),
  });
  try {
    const out = await callGateway(
      [
        { name: "broken", url: "http://127.0.0.1:1", addedAt: "now" },
        { name: "good", url: good.url, addedAt: "now" },
      ],
      null,
      "tools/list",
      {},
    );
    const names = out.tools.map((t: any) => t.name);
    assert.deepEqual(names, [`good${SEP}ping`]);
  } finally {
    await good.close();
  }
});

test("tools/call routes by agent prefix and forwards the bearer token", async () => {
  const upstream = await startFakeAgent({
    tools: [],
    onCall: (name, args) => ({
      content: [{ type: "text", text: `called ${name} with ${args.who}` }],
      structuredContent: { result: `hello ${args.who}` },
      isError: false,
    }),
  });
  try {
    const result = await callGateway(
      [{ name: "greeter", url: upstream.url, addedAt: "now" }],
      "tkn-123",
      "tools/call",
      { name: `greeter${SEP}greet`, arguments: { who: "world" } },
    );
    assert.equal(result.isError, false);
    assert.equal(result.structuredContent.result, "hello world");
    const req = upstream.requests.find((r) => r.body.method === "tools/call");
    assert.equal(req.auth, "Bearer tkn-123");
    assert.equal(req.body.params.name, "greet"); // prefix stripped
  } finally {
    await upstream.close();
  }
});

test("tools/call with unknown agent returns soft error", async () => {
  const result = await callGateway([], null, "tools/call", {
    name: `ghost${SEP}skill`,
    arguments: {},
  });
  assert.equal(result.isError, true);
  assert.match(result.content[0].text, /unknown agent: ghost/);
});

test("tools/call without prefix returns soft error", async () => {
  const result = await callGateway([], null, "tools/call", {
    name: "no_prefix",
    arguments: {},
  });
  assert.equal(result.isError, true);
  assert.match(result.content[0].text, /missing agent prefix/);
});

test("tools/call surfaces upstream JSON-RPC error as soft error", async () => {
  const upstream = await startFakeAgent({
    tools: [],
    onCall: () => {
      throw new Error("kaboom");
    },
  });
  try {
    const result = await callGateway(
      [{ name: "boom", url: upstream.url, addedAt: "now" }],
      null,
      "tools/call",
      { name: `boom${SEP}explode`, arguments: {} },
    );
    assert.equal(result.isError, true);
    assert.match(result.content[0].text, /kaboom/);
  } finally {
    await upstream.close();
  }
});

test("tools/call forwards upstream SSE progress notifications", async () => {
  const upstream = await startSseAgent({
    onCall: (body, res) => {
      res.write(
        `data: ${JSON.stringify({
          jsonrpc: "2.0",
          method: "notifications/progress",
          params: { progressToken: "p1", message: "halfway" },
        })}\n\n`,
      );
      res.end(
        `data: ${JSON.stringify({
          jsonrpc: "2.0",
          id: body.id,
          result: {
            content: [{ type: "text", text: "done" }],
            isError: false,
          },
        })}\n\n`,
      );
    },
  });
  const notifications: any[] = [];
  try {
    const result = await callGateway(
      [{ name: "streamer", url: upstream.url, addedAt: "now" }],
      null,
      "tools/call",
      { name: `streamer${SEP}build`, arguments: {} },
      {
        extra: {
          requestId: "req-1",
          sendNotification: async (notification: any) => {
            notifications.push(notification);
          },
        },
      },
    );
    assert.equal(result.isError, false);
    assert.deepEqual(notifications, [
      {
        method: "notifications/progress",
        params: { progressToken: "p1", message: "halfway" },
      },
    ]);
  } finally {
    await upstream.close();
  }
});

test("tools/call reports upstream SSE idle timeout", async () => {
  const upstream = await startSseAgent({
    onCall: (_body, res) => {
      res.write(": connected\n\n");
    },
  });
  try {
    const result = await callGateway(
      [{ name: "slow", url: upstream.url, addedAt: "now" }],
      null,
      "tools/call",
      { name: `slow${SEP}build`, arguments: {} },
      { gateway: { upstreamStreamIdleTimeoutMs: 25 } },
    );
    assert.equal(result.isError, true);
    assert.match(result.content[0].text, /stream idle timeout/);
  } finally {
    await upstream.close();
  }
});