import {
  assert,
  assertEquals,
  assertNotEquals,
  assertRejects,
  assertStringIncludes,
} from "@std/assert";
import { rootedEnvironment, testExtension } from "@berry/api/testing.ts";
import type { Environment, JsonValue, Process } from "@berry/api";
import { tempDir, until } from "@berry/util";
import mcp from "./mod.ts";

type Result = {
  content: { type: string; text: string }[];
  isError: boolean;
  structuredContent: {
    pid: number;
    count: number;
    cwd: string;
    token: string;
    cancelled: number;
    arguments: Record<string, JsonValue>;
    roots: { uri: string }[];
  };
};

async function within<T>(promise: Promise<T>, ms = 5000): Promise<T> {
  let timer: ReturnType<typeof setTimeout> | undefined;
  try {
    return await Promise.race([
      promise,
      new Promise<never>((_, reject) =>
        timer = setTimeout(() => reject(new Error("operation did not finish")), ms)
      ),
    ]);
  } finally {
    clearTimeout(timer);
  }
}

async function fixture() {
  const dir = await tempDir({ prefix: "mcp-test-" });
  await Deno.mkdir(dir.path + "/project");
  await Deno.mkdir(dir.path + "/other");
  await Deno.writeTextFile(
    dir.path + "/server.ts",
    await Deno.readTextFile(new URL("./test-server.ts", import.meta.url)),
  );
  const env = rootedEnvironment(dir.path, {
    cwd: "/project",
    variables: { MCP_TOKEN: "remote-secret" },
  });
  const spawn = env.spawn;
  const children: Process[] = [];
  env.spawn = async (...args) => {
    const child = await spawn(...args);
    children.push(child);
    return child;
  };
  const ext = testExtension(mcp, { env });
  await ext.configure({ configFile: "/servers.json" });
  const server = {
    command: Deno.execPath(),
    args: ["run", "-A", "../server.ts"],
    env: { TOKEN: "${MCP_TOKEN}" },
  };
  const configure = (servers: Record<string, unknown>) =>
    Deno.writeTextFile(dir.path + "/servers.json", JSON.stringify({ mcpServers: servers }));
  await configure({ probe: server });
  return {
    ext,
    env,
    children,
    configure,
    server,
    root: dir.path,
    async [Symbol.asyncDispose]() {
      await ext.dispose();
      await dir[Symbol.asyncDispose]();
    },
  };
}

Deno.test("MCP discovery is lazy, shares a connection, and calls preserve schemas, results, and progress", async () => {
  await using f = await fixture();
  await f.ext.prompt();
  assertEquals(f.children.length, 0);
  const [listed, described] = await Promise.all([
    f.ext.call("mcp_list", {}),
    f.ext.call("mcp_describe", { server: "probe", tool: "echo" }),
  ]);
  assertEquals(f.children.length, 1);
  assertEquals((listed as { tools: { name: string }[] }[])[0].tools[0].name, "echo");
  const description = described as { instructions: string; tools: Record<string, JsonValue>[] };
  assertStringIncludes(description.instructions, "echo");
  assertEquals(description.tools[0].annotations, { readOnlyHint: true });
  assert(description.tools[0].inputSchema);
  assert(description.tools[0].outputSchema);
  const progress: JsonValue[] = [];
  const result = await f.ext.call("mcp_call", {
    server: "probe",
    tool: "echo",
    arguments: { fail: true, query: "猫" },
  }, { progress: (value) => progress.push(value) }) as Result;
  assertEquals(result.content, [{ type: "text", text: "café, 猫" }]);
  assertEquals(result.isError, true);
  assertEquals(result.structuredContent.arguments, { fail: true, query: "猫" });
  assertEquals(result.structuredContent.cwd, f.root + "/project");
  assertEquals(result.structuredContent.token, "remote-secret");
  assertStringIncludes(result.structuredContent.roots[0].uri, "/project");
  assertEquals((progress[0] as { progress: number }).progress, 1);
  assertEquals(f.ext.state, { configFile: "/servers.json" });
});

Deno.test("unavailable MCP servers say why, without leaking credentials or hiding the others", async () => {
  await using f = await fixture();
  await f.configure({
    probe: f.server,
    broken: { command: "berry-mcp-no-such-program" },
    disabled: { enabled: false, command: "${MISSING_TOKEN}" },
    malformed: { command: "deno", args: "not-an-array" },
    exits: {
      command: Deno.execPath(),
      args: ["eval", "console.error('missing API_KEY'); Deno.exit(23)"],
    },
    unparsable: { url: "http://[${MCP_TOKEN}]/mcp" },
    userinfo: { url: "http://user:${MCP_TOKEN}@localhost/mcp" },
  });
  const listed = await within(f.ext.call("mcp_list", {})) as {
    server: string;
    status: string;
    error?: string;
  }[];
  assertEquals(listed.map(({ server, status }) => [server, status]), [
    ["probe", "connected"],
    ["broken", "failed"],
    ["disabled", "disabled"],
    ["malformed", "failed"],
    ["exits", "failed"],
    ["unparsable", "failed"],
    ["userinfo", "failed"],
  ]);
  assertEquals(f.children.length, 2);
  const error = (name: string) => listed.find(({ server }) => server === name)!.error!;
  assertStringIncludes(error("exits"), "missing API_KEY");
  assertStringIncludes(error("exits"), "23");
  assertStringIncludes(error("unparsable"), "URL");
  assert(!JSON.stringify(listed).includes("remote-secret"));
  await assertRejects(
    () => f.ext.call("mcp_call", { server: "disabled", tool: "echo" }),
    Error,
    "disabled",
  );
  await assertRejects(
    () => f.ext.call("mcp_describe", { server: "probe", tool: "missing" }),
    Error,
    "No tool",
  );
  await assertRejects(
    () => f.ext.call("mcp_call", { server: "missing", tool: "echo" }),
    Error,
    "No MCP server",
  );
});

Deno.test("clients follow the machine, directory, and current server configuration", async () => {
  await using f = await fixture();
  const call = (env?: Environment) =>
    f.ext.call("mcp_call", { server: "probe", tool: "echo" }, { env }) as Promise<Result>;
  const first = (await call()).structuredContent;
  assertEquals((await call()).structuredContent.pid, first.pid);
  f.env.cwd = "/other";
  const moved = (await call()).structuredContent;
  assertNotEquals(moved.pid, first.pid);
  assertEquals(moved.cwd, f.root + "/other");
  const another = { ...f.env, id: "another-machine", variables: { MCP_TOKEN: "another-secret" } };
  const remote = (await call(another)).structuredContent;
  assertNotEquals(remote.pid, moved.pid);
  assertEquals(remote.token, "another-secret");
  await f.configure({ probe: { ...f.server, env: { TOKEN: "changed" } } });
  const changed = (await call()).structuredContent;
  assertNotEquals(changed.pid, moved.pid);
  assertEquals(changed.token, "changed");
});

Deno.test("cancellation reaches the server without discarding the shared connection", async () => {
  await using f = await fixture();
  const abort = new AbortController();
  const waiting = f.ext.call("mcp_call", { server: "probe", tool: "wait" }, {
    signal: abort.signal,
  });
  const rejected = assertRejects(() => waiting);
  await until(() => Deno.stat(f.root + "/project/waiting").then(() => true, () => false));
  abort.abort();
  await rejected;
  const next = await f.ext.call("mcp_call", { server: "probe", tool: "echo" }) as Result;
  assertEquals(next.structuredContent.cancelled, 1);
  assertEquals(f.children.length, 1);
});

Deno.test("a crashed server reconnects for the next call, never by replaying the failed call", async () => {
  await using f = await fixture();
  await assertRejects(() => f.ext.call("mcp_call", { server: "probe", tool: "crash" }));
  assertEquals(f.children.length, 1);
  const next = await Promise.all([
    f.ext.call("mcp_call", { server: "probe", tool: "echo" }),
    f.ext.call("mcp_call", { server: "probe", tool: "echo" }),
  ]) as Result[];
  assertEquals(next.map((result) => result.structuredContent.count).sort(), [1, 2]);
  assertEquals(f.children.length, 2);
});

Deno.test("reload disposal force-stops a server that ignores stdin closure and SIGTERM", async () => {
  await using f = await fixture();
  await f.configure({ probe: { ...f.server, args: [...f.server.args, "stubborn"] } });
  await f.ext.call("mcp_list", {});
  await f.ext.dispose();
  assertEquals(await f.children[0].exited, null);
});

Deno.test({
  name: "reload disposal stops descendants even when their wrapper has already exited",
  ignore: Deno.build.os === "windows",
  async fn() {
    await using f = await fixture();
    await f.configure({ probe: { ...f.server, args: [...f.server.args, "descendant"] } });
    await f.ext.call("mcp_list", {});
    const pid = Number(await Deno.readTextFile(f.root + "/project/descendant"));
    await f.ext.dispose();
    await until(() => {
      try {
        Deno.kill(pid, 0);
        return false;
      } catch {
        return true;
      }
    });
  },
});

Deno.test("disposal during startup does not leave a late-spawned process running", async () => {
  await using f = await fixture();
  const spawn = f.env.spawn;
  const entered = Promise.withResolvers<void>();
  const release = Promise.withResolvers<void>();
  f.env.spawn = async (...args) => {
    const child = await spawn(...args);
    entered.resolve();
    await release.promise;
    return child;
  };
  const listing = f.ext.call("mcp_list", {});
  await entered.promise;
  try {
    await within(f.ext.dispose());
    await within(listing);
  } finally {
    release.resolve();
    assertEquals(await f.children[0].exited, null);
  }
});

Deno.test("Streamable HTTP uses the environment's network and preserves MCP errors", async () => {
  await using f = await fixture();
  const authorization: (string | null)[] = [];
  const methods: string[] = [];
  const server = Deno.serve({ hostname: "127.0.0.1", port: 0, onListen() {} }, async (request) => {
    authorization.push(request.headers.get("authorization"));
    methods.push(request.method);
    if (request.method === "GET") return new Response(null, { status: 405 });
    if (request.method === "DELETE") return new Response(null, { status: 204 });
    const message = await request.json();
    if (message.id === undefined) return new Response(null, { status: 202 });
    const result = message.method === "initialize"
      ? {
        protocolVersion: message.params.protocolVersion,
        capabilities: { tools: {} },
        serverInfo: { name: "http", version: "1" },
      }
      : message.method === "tools/list"
      ? { tools: [{ name: "echo", inputSchema: { type: "object" } }] }
      : { content: [], structuredContent: message.params.arguments.answer, isError: true };
    return Response.json({ jsonrpc: "2.0", id: message.id, result }, {
      headers: { "Mcp-Session-Id": "test-session" },
    });
  });
  try {
    let fetched = 0;
    const fetchHere = f.env.fetch;
    f.env.fetch = (input, init) => {
      assertEquals(String(input), "http://only-on-the-environment.invalid/mcp");
      fetched++;
      return fetchHere(`http://127.0.0.1:${server.addr.port}/mcp`, init);
    };
    await f.configure({
      remote: {
        url: "http://only-on-the-environment.invalid/mcp",
        headers: { Authorization: "Bearer ${MCP_TOKEN}" },
      },
    });
    const result = await f.ext.call("mcp_call", {
      server: "remote",
      tool: "echo",
      arguments: { answer: [1, "two", null] },
    }) as { structuredContent: JsonValue; isError: boolean };
    assertEquals(result.structuredContent, [1, "two", null]);
    assertEquals(result.isError, true);
    assert(fetched > 0);
    assert(authorization.every((header) => header === "Bearer remote-secret"));
    await f.ext.dispose();
    assert(methods.includes("DELETE"));
  } finally {
    await server.shutdown();
  }
});

Deno.test("an HTTP network failure retires the session without aborting other calls already using it", async () => {
  await using f = await fixture();
  const entered = Promise.withResolvers<void>();
  const release = Promise.withResolvers<void>();
  let initialized = 0;
  const server = Deno.serve({ hostname: "127.0.0.1", port: 0, onListen() {} }, async (request) => {
    if (request.method === "GET") return new Response(null, { status: 405 });
    if (request.method === "DELETE") return new Response(null, { status: 204 });
    const message = await request.json();
    if (message.id === undefined) return new Response(null, { status: 202 });
    if (message.params?.name === "slow") {
      entered.resolve();
      await release.promise;
    }
    const result = message.method === "initialize"
      ? {
        protocolVersion: message.params.protocolVersion,
        capabilities: { tools: {} },
        serverInfo: { name: "http", version: String(++initialized) },
      }
      : { content: [{ type: "text", text: message.params.name }] };
    return Response.json({ jsonrpc: "2.0", id: message.id, result }, {
      headers: { "Mcp-Session-Id": "session" },
    });
  });
  const fetchHere = f.env.fetch;
  f.env.fetch = (input, init) => {
    const message = typeof init?.body === "string" ? JSON.parse(init.body) : undefined;
    if (message?.params?.name === "failed") return Promise.reject(new TypeError("network lost"));
    return fetchHere(input, init);
  };
  const call = (tool: string) => f.ext.call("mcp_call", { server: "http", tool });
  try {
    await f.configure({ http: { url: `http://127.0.0.1:${server.addr.port}/mcp` } });
    const slow = call("slow").then((value) => ({ value }), (error: unknown) => ({ error }));
    await within(entered.promise);
    await assertRejects(() => call("failed"), TypeError, "network lost");
    await call("fresh");
    assertEquals(initialized, 2);
    release.resolve();
    const result = await within(slow);
    assert("value" in result);
    assertEquals((result.value as { content: JsonValue }).content, [{
      type: "text",
      text: "slow",
    }]);
  } finally {
    release.resolve();
    await f.ext.dispose();
    await server.shutdown();
  }
});