import { assert, assertEquals, assertRejects } from "@std/assert";
import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import { testExtension } from "@berry/api/testing.ts";
import { tempDir, until } from "@berry/util";
import mcp from "../../extensions/mcp/mod.ts";
import { BOOTSTRAP, Connection, RemoteEnvironment } from "./connection.ts";

Deno.test("MCP servers use the worker's environment and remote force-stop leaves the worker usable", async () => {
  await using dir = await tempDir({ prefix: "remote-mcp-" });
  await Deno.writeTextFile(
    dir.path + "/server.ts",
    await Deno.readTextFile(new URL("../../extensions/mcp/test-server.ts", import.meta.url)),
  );
  await Deno.writeTextFile(
    dir.path + "/mcp.json",
    JSON.stringify({
      mcpServers: {
        probe: {
          command: Deno.execPath(),
          args: ["run", "-A", "server.ts", "stubborn"],
          env: { TOKEN: "${BERRY_MACHINE}" },
        },
      },
    }),
  );
  const connection = await Connection.open("mcp-test", [
    "env",
    "BERRY_MACHINE=mcp-test",
    Deno.execPath(),
    "eval",
    "--quiet",
    BOOTSTRAP,
  ]);
  const env = new RemoteEnvironment("mcp-test", connection, dir.path);
  const ext = testExtension(mcp, { env });
  try {
    await ext.configure({ configFile: "mcp.json" });
    const result = await ext.call("mcp_call", { server: "probe", tool: "echo" }) as {
      structuredContent: { token: string; cwd: string };
    };
    assertEquals(result.structuredContent.token, "mcp-test");
    assertEquals(result.structuredContent.cwd, dir.path);
    await ext.dispose();
    assertEquals((await env.readTextFile("mcp.json", context)).ok, true);

    const child = await env.spawn(["cat"], {}, context);
    await child.stdout.cancel();
    const writer = child.stdin.getWriter();
    await writer.write(new TextEncoder().encode("discarded output"));
    await writer.close();
    assertEquals(await child.exited, 0);
    assertEquals((await env.fileInfo("server.ts", context)).ok, true);

    const abandoned = await env.spawn(["cat"], {}, context);
    const pending = assertRejects(() => new Response(abandoned.stdout).text(), Error, "lost");
    await connection.close();
    await pending;
  } finally {
    await ext.dispose();
    await connection.close();
  }
});

Deno.test("HTTP MCP calls reuse a live worker but the first call after worker replacement opens a fresh session", async () => {
  await using dir = await tempDir({ prefix: "remote-http-mcp-" });
  let initialized = 0;
  let called = 0;
  let session = "";
  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.method !== "initialize" && request.headers.get("Mcp-Session-Id") !== session) {
      return new Response(null, { status: 404 });
    }
    const result = message.method === "initialize"
      ? {
        protocolVersion: message.params.protocolVersion,
        capabilities: { tools: {} },
        serverInfo: { name: "http", version: "1" },
      }
      : { content: [], structuredContent: { session, called: ++called } };
    if (message.method === "initialize") session = String(++initialized);
    return Response.json({ jsonrpc: "2.0", id: message.id, result }, {
      headers: { "Mcp-Session-Id": session },
    });
  });
  await Deno.writeTextFile(
    dir.path + "/mcp.json",
    JSON.stringify({
      mcpServers: { http: { url: `http://127.0.0.1:${server.addr.port}/mcp` } },
    }),
  );
  const open = () => Connection.open("http", [Deno.execPath(), "eval", "--quiet", BOOTSTRAP]);
  const first = await open();
  let second: Connection | undefined;
  const ext = testExtension(mcp, { env: new RemoteEnvironment("http", first, dir.path) });
  const call = (env?: RemoteEnvironment) =>
    ext.call("mcp_call", { server: "http", tool: "echo" }, { env });
  try {
    await ext.configure({ configFile: "mcp.json" });
    await call();
    await call(new RemoteEnvironment("http", first, dir.path));
    assertEquals(initialized, 1);
    await first.close();
    assert(first.lifetime.aborted);
    second = await open();
    const result = await call(new RemoteEnvironment("http", second, dir.path)) as {
      structuredContent: { session: string; called: number };
    };
    assertEquals(result.structuredContent, { session: "2", called: 3 });
    assertEquals(initialized, 2);
    assertEquals(called, 3);
  } finally {
    await ext.dispose();
    await first.close();
    await second?.close();
    await server.shutdown();
  }
});

Deno.test({
  name: "a stalled remote worker cannot prevent fetch cancellation or HTTP MCP disposal",
  ignore: Deno.build.os !== "linux",
  async fn() {
    await using dir = await tempDir({ prefix: "stalled-mcp-" });
    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 });
        const result = message.method === "initialize"
          ? {
            protocolVersion: message.params.protocolVersion,
            capabilities: { tools: {} },
            serverInfo: { name: "http", version: "1" },
          }
          : { content: [] };
        return Response.json({ jsonrpc: "2.0", id: message.id, result }, {
          headers: { "Mcp-Session-Id": "session" },
        });
      },
    );
    const url = `http://127.0.0.1:${server.addr.port}/mcp`;
    await Deno.writeTextFile(
      dir.path + "/mcp.json",
      JSON.stringify({ mcpServers: { http: { url } } }),
    );
    const connection = await Connection.open("stalled", [
      "sh",
      "-c",
      'echo $$ > "$1"; shift; exec "$@"',
      "sh",
      dir.path + "/pid",
      Deno.execPath(),
      "eval",
      "--quiet",
      BOOTSTRAP,
    ]);
    const pid = Number(await Deno.readTextFile(dir.path + "/pid"));
    const env = new RemoteEnvironment("stalled", connection, dir.path);
    const ext = testExtension(mcp, { env });
    try {
      await ext.configure({ configFile: "mcp.json" });
      await ext.call("mcp_call", { server: "http", tool: "echo" });
      Deno.kill(pid, "SIGSTOP");
      const abort = new AbortController();
      const reason = new Error("request canceled");
      let rejected: unknown;
      const fetching = env.fetch(url, { signal: abort.signal }).catch((error: unknown) =>
        rejected = error
      );
      abort.abort(reason);
      await until(() => rejected !== undefined, {
        timeout: 1000,
        message: "fetch abort waited for a worker reply",
      });
      await fetching;
      assertEquals(rejected, reason);
      let disposed = false;
      const disposing = ext.dispose().then(() => disposed = true);
      await until(() => disposed, {
        timeout: 3000,
        message: "HTTP MCP disposal waited for a stalled worker",
      });
      await disposing;
      assert(!connection.lifetime.aborted);
    } finally {
      Deno.kill(pid, "SIGCONT");
      await ext.dispose();
      await connection.close();
      await server.shutdown();
    }
  },
});