/**
 * Copyright (c) 2026, Salesforce, Inc.,
 * All rights reserved.
 * For full license text, see the LICENSE.txt file
 */

import path from "node:path";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { buildSchema, introspectionFromSchema } from "graphql";
import { describe, expect, it } from "vitest";
import { atomicWriteJson } from "../../../lib/fs-utils.js";
import {
	schemaCacheKeyForInstanceUrl,
	schemaDir,
	type SchemaMetadata,
} from "../../../lib/introspect.js";
import { type ObjectInfoResult } from "../../../lib/object-info.js";
import { type PrimeDeps } from "../../../lib/prime-schema.js";
import { primeSchemaCache } from "../../../lib/walker.js";
import { registerSfGqlDiscoverTool } from "../sf-gql-discover.js";

const ORG = "test-tool-discover";
const ORG_URL = "https://test-tool-discover.my.salesforce.com";
const SCHEMA = buildSchema(`
	type Query { uiapi: UIAPI! }
	type UIAPI { query: RecordQuery! }
	type RecordQuery {
		Account(first: Int, after: String): AccountConnection!
	}
	type AccountConnection { edges: [AccountEdge!]! }
	type AccountEdge { node: Account! }
	type Account { Id: ID!, Name: StringValue }
	type StringValue { value: String }
`);
primeSchemaCache(ORG, SCHEMA);
primeSchemaCache(ORG_URL, SCHEMA);

const primeDeps: PrimeDeps = {
	getOrgAuth: async () => ({
		alias: ORG,
		username: "u",
		instanceUrl: ORG_URL,
		accessToken: "t",
		orgId: "00D",
	}),
	downloadSchema: async (auth) => {
		const cacheKey = schemaCacheKeyForInstanceUrl(auth.instanceUrl);
		const filePath = path.join(schemaDir(), `${cacheKey}.json`);
		atomicWriteJson(filePath, { data: introspectionFromSchema(SCHEMA) });
		const meta: SchemaMetadata = {
			cacheKey,
			instanceUrl: auth.instanceUrl,
			typeCount: 0,
			downloadedAt: new Date().toISOString(),
			filePath,
		};
		return meta;
	},
};

const ACCOUNT_INFO: ObjectInfoResult = {
	apiName: "Account",
	label: "Account",
	labelPlural: "Accounts",
	createable: true,
	deletable: true,
	updateable: true,
	queryable: true,
	searchable: true,
	custom: false,
	keyPrefix: "001",
	nameFields: ["Name"],
	defaultRecordTypeId: null,
	fields: [
		{
			apiName: "Id",
			label: "Account ID",
			dataType: "ID",
			required: false,
			createable: false,
			updateable: false,
			calculated: false,
			custom: false,
			filterable: true,
			sortable: true,
			nameField: false,
			reference: false,
			relationshipName: null,
			compound: false,
			compoundFieldName: null,
			defaultedOnCreate: false,
			extraTypeInfo: null,
			inlineHelpText: null,
			precision: 0,
			scale: 0,
			referenceToInfos: [],
			controllerName: null,
			controllingFields: [],
		},
	],
	childRelationships: [],
	recordTypeInfos: [],
	picklists: [],
	fetchedAt: new Date().toISOString(),
};

async function connect(): Promise<{ client: Client; server: McpServer }> {
	const server = new McpServer({ name: "graphiti-mcp", version: "test" });
	registerSfGqlDiscoverTool(server, {
		primeDeps,
		getOrgAuth: async () => ({
			alias: ORG,
			username: "u",
			instanceUrl: ORG_URL,
			accessToken: "t",
			orgId: "00D",
		}),
		getObjectInfo: async (_auth, _alias, name) => {
			if (name === "Account") return ACCOUNT_INFO;
			throw new Error(`No fixture for ${name}`);
		},
	});
	const [c, s] = InMemoryTransport.createLinkedPair();
	const client = new Client({ name: "test", version: "0.0.0" });
	await Promise.all([server.connect(s), client.connect(c)]);
	return { client, server };
}

describe("mcp/tools/sf-gql-discover", () => {
	it("tools/list advertises sf_gql_discover with org/mode properties", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_discover");
			expect(tool).toBeDefined();
			const props =
				(tool!.inputSchema as { properties?: Record<string, unknown> }).properties ?? {};
			expect(props.org).toBeDefined();
			expect(props.mode).toBeDefined();
		} finally {
			await client.close();
			await server.close();
		}
	});

	// W-23636056: the `search` Cc/Cf guard must be enforced with `.refine()`, not
	// `.regex()`. `.regex()` serializes into the advertised JSON-Schema as a
	// `pattern` whose `\p{...}` property escape is legal in ECMAScript/RE2 but
	// ILLEGAL in Python's `re` — so a Python-side JSON-Schema consumer (the eval
	// harness validating tool-call args with `jsonschema`) throws `bad escape \p`
	// on EVERY discover call. This asserts on the schema the SDK actually
	// advertises over tools/list (the exact conversion the host serves), so a
	// silent revert to `.regex()` reopens the crash and fails here. The runtime
	// rejection rule itself is covered by input-schemas.spec.ts (unaffected by
	// this .regex()->.refine() change, since both enforce DISCOVER_SEARCH_RE).
	it("advertises `search` with NO `pattern` keyword and no `\\p{` escape (cross-runtime portable)", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_discover");
			const inputSchema = tool!.inputSchema as {
				properties: Record<string, { type?: string; maxLength?: number; pattern?: string }>;
			};
			const search = inputSchema.properties.search;
			// search is still advertised (type + length cap), just without a pattern.
			expect(search).toBeDefined();
			expect(search.type).toBe("string");
			expect(search.maxLength).toBe(100);
			// The crux: no `pattern` (so no regex reaches a Python jsonschema consumer)...
			expect(search.pattern).toBeUndefined();
			// ...and the illegal-in-Python `\p{...}` escape appears nowhere in the
			// whole emitted schema (guards against it leaking via any other field too).
			expect(JSON.stringify(inputSchema)).not.toContain("\\p{");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call list_objects returns the schema's queryable SObjects", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "list_objects" },
			});
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.type).toBe("text");
			const parsed = JSON.parse(content[0]?.text ?? "{}") as {
				mode: string;
				objects: { name: string }[];
			};
			expect(parsed.mode).toBe("list_objects");
			expect(parsed.objects.map((o) => o.name)).toContain("Account");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call describe_object returns ObjectDescription", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "describe_object", object: "Account" },
			});
			const content = result.content as { type: string; text?: string }[];
			const parsed = JSON.parse(content[0]?.text ?? "{}") as {
				mode: string;
				object: { name: string; fields: { name: string }[] };
			};
			expect(parsed.mode).toBe("describe_object");
			expect(parsed.object.name).toBe("Account");
			expect(parsed.object.fields.map((f) => f.name)).toContain("Id");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call with missing required arg returns validation error", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/mode/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call describe_object without object returns error", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "describe_object" },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/object/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects org alias containing shell metacharacters", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: "evil; rm -rf /", mode: "list_objects" },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/org/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects object name containing path traversal", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "describe_object", object: "../etc/passwd" },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/object/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("accepts org alias at the 80-char limit", async () => {
		const { client, server } = await connect();
		try {
			const longOrg = "a".repeat(80);
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: longOrg, mode: "list_objects" },
			});
			// Org isn't primed so the call fails downstream, but it must
			// pass zod validation — surface as a non-validation error.
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").not.toMatch(/org must match/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects org alias one over the 80-char limit", async () => {
		const { client, server } = await connect();
		try {
			const tooLongOrg = "a".repeat(81);
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: tooLongOrg, mode: "list_objects" },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/org/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects search containing control characters", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "list_objects", search: "foo\x00bar" },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/search/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects search exceeding 100 characters", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: "list_objects", search: "a".repeat(101) },
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/search/);
		} finally {
			await client.close();
			await server.close();
		}
	});
});

// W-23336443: the `mode` enum is wrapped in enumStripControlChars. A z.enum
// rejection is reflected VERBATIM by the MCP SDK's input validation, which runs
// UPSTREAM of runTool — so the tool adapter's neutralizeControlChars never sees
// it. These tests exercise the REAL SDK channel (client.callTool round-trip),
// which is the actual vulnerable path, not a bare schema.safeParse.
//
// Oracle discipline: the SDK pretty-prints zod issues with JSON.stringify(_, 2),
// which inserts real newlines (themselves \p{Cc}) into the envelope, so a
// whole-class /[\p{Cc}\p{Cf}]/ scan false-positives. We assert instead that the
// SPECIFIC injected code points (U+202E bidi, U+200B ZWSP, U+007F DEL — the trio
// the SDK's JSON.stringify leaves raw, since it escapes only C0) are absent.
describe("mcp/tools/sf-gql-discover — control-char neutralization (W-23336443)", () => {
	// One representative per surviving class. C0 (e.g. \n) is escaped by the SDK's
	// JSON.stringify already; these three are the ones that survived raw pre-fix.
	const BIDI = "\u{202e}";
	const ZWSP = "\u{200b}";
	const DEL = "\x7f";

	it("a poisoned-but-otherwise-valid mode strips to the member and is ACCEPTED", async () => {
		const { client, server } = await connect();
		try {
			// "describe_object" + trailing bidi override. Pre-fix this rejected and
			// reflected the raw U+202E; post-fix it strips to the valid member and
			// proceeds (the Account fixture describe succeeds).
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: `describe_object${BIDI}`, object: "Account" },
			});
			expect(result.isError).toBeFalsy();
			const content = result.content as { type: string; text?: string }[];
			const parsed = JSON.parse(content[0]?.text ?? "{}") as { mode: string };
			expect(parsed.mode).toBe("describe_object");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("a genuinely-invalid poisoned mode rejects with NO raw injected code point", async () => {
		const { client, server } = await connect();
		try {
			// "bogus" wrapped in bidi + zero-width + DEL. Strips to "bogus", which is
			// not a member, so the SDK rejects — but the reflected value is stripped.
			const result = await client.callTool({
				name: "sf_gql_discover",
				arguments: { org: ORG, mode: `${BIDI}bo${ZWSP}gus${DEL}` },
			});
			expect(result.isError).toBe(true);
			const text = (result.content as { text?: string }[])[0]?.text ?? "";
			expect(text).not.toContain(BIDI);
			expect(text).not.toContain(ZWSP);
			expect(text).not.toContain(DEL);
			// It still reflects the stripped token so the LLM can self-correct.
			expect(text).toContain("bogus");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/list still advertises the mode enum + description after wrapping", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_discover");
			const mode = (tool!.inputSchema as { properties: Record<string, unknown> }).properties
				.mode as { enum?: string[]; description?: string };
			expect(mode.enum).toEqual(["list_objects", "describe_object", "describe_field"]);
			expect(mode.description).toMatch(/Discovery mode/);
			// mode stays REQUIRED through the preprocess wrapper.
			const required = (tool!.inputSchema as { required?: string[] }).required ?? [];
			expect(required).toContain("mode");
		} finally {
			await client.close();
			await server.close();
		}
	});
});