import assert from "node:assert/strict";
import test from "node:test";
import { startStreamableHttpServer } from "../src/streamable-http.js";

test("accepts a synchronous session cleanup callback during shutdown", async () => {
  const closed = [];
  const listener = await startStreamableHttpServer({
    serverName: "test-postgres-mcp",
    serverVersion: "1.0.0",
    tools: [],
    callTool: async () => ({}),
    apiKeys: ["agent=0123456789abcdef0123456789abcdef"],
    host: "127.0.0.1",
    port: 0,
    onSessionClose: ({ id }) => { closed.push(id); }
  });

  const response = await fetch(`http://127.0.0.1:${listener.port}/mcp`, {
    method: "POST",
    headers: {
      authorization: "Bearer 0123456789abcdef0123456789abcdef",
      "content-type": "application/json"
    },
    body: JSON.stringify({
      jsonrpc: "2.0",
      id: 1,
      method: "initialize",
      params: { protocolVersion: "2025-03-26", capabilities: {}, clientInfo: { name: "test", version: "1" } }
    })
  });
  const sessionId = response.headers.get("mcp-session-id");
  assert.equal(response.status, 200);
  assert.ok(sessionId);

  await listener.close();
  assert.deepEqual(closed, [sessionId]);
});

const token = "0123456789abcdef0123456789abcdef";

async function startPinned(options = {}) {
  return startStreamableHttpServer({
    serverName: "test-postgres-mcp",
    serverVersion: "1.0.0",
    tools: (context) => context.pinned["x-postgres-scope"] !== undefined ? [{ name: "internal_only" }] : [{ name: "external_only" }],
    callTool: async (name, _args, context) => ({ name, pinned: context.pinned }),
    apiKeys: [`agent=${token}`],
    host: "127.0.0.1",
    port: 0,
    pinnedHeaders: ["x-postgres-scope"],
    ...options
  });
}

function rpc(listener, body, headers = {}) {
  return fetch(`http://127.0.0.1:${listener.port}/mcp`, {
    method: "POST",
    headers: { authorization: `Bearer ${token}`, "content-type": "application/json", ...headers },
    body: JSON.stringify({ jsonrpc: "2.0", id: 1, ...body })
  });
}

const initialize = { method: "initialize", params: { protocolVersion: "2025-03-26", capabilities: {}, clientInfo: { name: "test", version: "1" } } };

test("pins the bucket header at session open and offers that session's tools", async () => {
  const initialized = [];
  const listener = await startPinned({ onSessionInitialize: (context) => { initialized.push(context.pinned); } });
  try {
    const opened = await rpc(listener, initialize, { "x-postgres-scope": "exec-1" });
    const sessionId = opened.headers.get("mcp-session-id");
    assert.equal(opened.status, 200);
    assert.deepEqual(initialized, [{ "x-postgres-scope": "exec-1" }]);

    const listed = await (await rpc(listener, { method: "tools/list" }, { "mcp-session-id": sessionId })).json();
    assert.deepEqual(listed.result.tools.map((tool) => tool.name), ["internal_only"]);

    const called = await (await rpc(listener, { method: "tools/call", params: { name: "internal_only" } }, { "mcp-session-id": sessionId })).json();
    assert.deepEqual(called.result.structuredContent.pinned, { "x-postgres-scope": "exec-1" });

    const refused = await (await rpc(listener, { method: "tools/call", params: { name: "external_only" } }, { "mcp-session-id": sessionId })).json();
    assert.match(refused.error.message, /Unknown tool/);
  } finally {
    await listener.close();
  }
});

test("refuses a bucket header that changes within a session", async () => {
  const listener = await startPinned();
  try {
    const opened = await rpc(listener, initialize, { "x-postgres-scope": "exec-1" });
    const sessionId = opened.headers.get("mcp-session-id");
    const changed = await rpc(listener, { method: "tools/list" }, { "mcp-session-id": sessionId, "x-postgres-scope": "exec-2" });
    assert.equal(changed.status, 400);
    assert.match((await changed.json()).error, /cannot change within a session/);
  } finally {
    await listener.close();
  }
});

test("offers the external tools to a session opened without the bucket header", async () => {
  const listener = await startPinned();
  try {
    const opened = await rpc(listener, initialize);
    const listed = await (await rpc(listener, { method: "tools/list" }, { "mcp-session-id": opened.headers.get("mcp-session-id") })).json();
    assert.deepEqual(listed.result.tools.map((tool) => tool.name), ["external_only"]);
  } finally {
    await listener.close();
  }
});

test("does not open a session whose initialization fails", async () => {
  const closed = [];
  const listener = await startPinned({
    onSessionInitialize: () => { throw Object.assign(new Error("not a valid bucket name"), { statusCode: 400 }); },
    onSessionClose: ({ id }) => { closed.push(id); }
  });
  try {
    const refused = await rpc(listener, initialize, { "x-postgres-scope": "Bad" });
    assert.equal(refused.status, 400);
    assert.match((await refused.json()).error, /Unable to initialize the session: not a valid bucket name/);
    assert.equal(refused.headers.get("mcp-session-id"), null);
    assert.equal(closed.length, 1);
  } finally {
    await listener.close();
  }
});