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

const TOKEN = "0123456789abcdef0123456789abcdef";
const HEADER = "x-preview-key";

async function server(onCall) {
  const listener = await startStreamableHttpServer({
    serverName: "test", serverVersion: "1.0.0",
    tools: [{ name: "probe", description: "probe", inputSchema: { type: "object", additionalProperties: false } }],
    callTool: onCall,
    apiKeys: [`subject=${TOKEN}`],
    host: "127.0.0.1",
    port: 0,
    pinnedHeaders: [HEADER]
  });
  return listener;
}

async function post(listener, body, headers = {}) {
  const response = await fetch(`http://127.0.0.1:${listener.port}/mcp`, {
    method: "POST",
    headers: { authorization: `Bearer ${TOKEN}`, "content-type": "application/json", accept: "application/json", ...headers },
    body: JSON.stringify(body)
  });
  const text = await response.text();
  return { status: response.status, sessionId: response.headers.get("mcp-session-id"), body: text ? JSON.parse(text) : null };
}

const INITIALIZE = { jsonrpc: "2.0", id: 1, method: "initialize", params: { protocolVersion: "2025-06-18", capabilities: {}, clientInfo: { name: "t", version: "0" } } };
const CALL = { jsonrpc: "2.0", id: 2, method: "tools/call", params: { name: "probe", arguments: {} } };

test("a pinned header reaches the tool call as session scope", async () => {
  let seen = null;
  const listener = await server(async (name, args, context) => { seen = context.pinned; return { ok: true }; });
  try {
    const opened = await post(listener, INITIALIZE, { [HEADER]: "exec-a" });
    await post(listener, CALL, { "mcp-session-id": opened.sessionId });

    assert.deepEqual(seen, { [HEADER]: "exec-a" });
  } finally { await listener.close(); }
});

test("the scope cannot be changed part-way through a session", async () => {
  // Otherwise one open session would be enough to reach every execution in turn, which is exactly
  // what taking the key out of the tool arguments was meant to prevent.
  const listener = await server(async () => ({ ok: true }));
  try {
    const opened = await post(listener, INITIALIZE, { [HEADER]: "exec-a" });
    const switched = await post(listener, CALL, { "mcp-session-id": opened.sessionId, [HEADER]: "exec-b" });

    assert.equal(switched.status, 400);
    assert.match(switched.body.error, /cannot change within a session/);
  } finally { await listener.close(); }
});

test("repeating the same scope on a later request is accepted", async () => {
  const listener = await server(async () => ({ ok: true }));
  try {
    const opened = await post(listener, INITIALIZE, { [HEADER]: "exec-a" });
    const again = await post(listener, CALL, { "mcp-session-id": opened.sessionId, [HEADER]: "exec-a" });

    assert.equal(again.status, 200);
  } finally { await listener.close(); }
});

test("a session opened without the header carries no scope at all", async () => {
  let seen = "unset";
  const listener = await server(async (name, args, context) => { seen = context.pinned; return { ok: true }; });
  try {
    const opened = await post(listener, INITIALIZE);
    await post(listener, CALL, { "mcp-session-id": opened.sessionId });

    assert.deepEqual(seen, {}, "an absent header must not become an empty-string scope");
  } finally { await listener.close(); }
});

test("an oversized scope is refused at session open", async () => {
  // Control characters never get this far - the HTTP client and Node's parser both refuse them -
  // so the guard that earns its place here is the length one.
  const listener = await server(async () => ({ ok: true }));
  try {
    const opened = await post(listener, INITIALIZE, { [HEADER]: "x".repeat(257) });

    assert.equal(opened.status, 400);
    assert.match(opened.body.error, /Invalid x-preview-key header/);
  } finally { await listener.close(); }
});