/**
 * 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 PrimeDeps } from "../../../lib/prime-schema.js";
import { primeSchemaCache } from "../../../lib/walker.js";
import { registerSfGqlAggregateTool } from "../sf-gql-aggregate.js";

const ORG = "test-tool-aggregate";
const ORG_URL = "https://test-tool-aggregate.my.salesforce.com";
const SCHEMA = buildSchema(`
	type Query { uiapi: UIAPI! }
	type UIAPI { aggregate: RecordQueryAggregate! }
	type RecordQueryAggregate {
		Account(first: Int, after: String, where: Account_Filter, orderBy: Account_OrderBy, groupBy: Account_GroupBy): AccountAggregateConnection
	}
	input Account_Filter { Industry: PicklistOperators }
	input Account_GroupBy { Industry: GroupByClause }
	input Account_OrderBy { Industry: OrderByValue }
	input PicklistOperators { eq: String }
	input GroupByClause { group: Boolean }
	input OrderByValue { order: String }

	type AccountAggregateConnection {
		edges: [AccountAggregateEdge!]!
		pageInfo: PageInfo!
	}
	type AccountAggregateEdge { node: AccountResult!, cursor: String! }
	type AccountResult { aggregate: AccountAggregate }
	type AccountAggregate {
		Id: IDAggregate
		Industry: PicklistAggregate
	}
	type IDAggregate { value: ID, count: LongValue, countDistinct: LongValue }
	type PicklistAggregate { value: String, count: LongValue, countDistinct: LongValue }
	type LongValue { value: Float }
	type PageInfo { hasNextPage: Boolean!, endCursor: 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;
	},
};

async function connect(): Promise<{ client: Client; server: McpServer }> {
	const server = new McpServer({ name: "graphiti-mcp", version: "test" });
	registerSfGqlAggregateTool(server, { primeDeps });
	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-aggregate", () => {
	it("tools/list advertises sf_gql_aggregate with org/object/groupBy/aggregations properties", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_aggregate");
			expect(tool).toBeDefined();
			const props =
				(tool!.inputSchema as { properties?: Record<string, unknown> }).properties ?? {};
			expect(props.org).toBeDefined();
			expect(props.object).toBeDefined();
			expect(props.groupBy).toBeDefined();
			expect(props.aggregations).toBeDefined();
			// W-22899736: object + operationName advertise the GraphQL Name pattern in tools/list.
			expect((props.object as { pattern?: string }).pattern).toBe("^[A-Za-z_][A-Za-z0-9_]*$");
			expect((props.operationName as { pattern?: string }).pattern).toBe(
				"^[A-Za-z_][A-Za-z0-9_]*$",
			);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call sf_gql_aggregate returns ToolOutput JSON envelope", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
				},
			});
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.type).toBe("text");
			const parsed = JSON.parse(content[0]?.text ?? "{}") as {
				query: string;
				variables: unknown[];
				types: string;
				warnings: string[];
			};
			expect(parsed.query).toMatch(/\bAccountAggregate\b/);
			expect(parsed.query).toMatch(/countId\s*:\s*Id\s*\{\s*count\s*\{\s*value/s);
			expect(parsed.query).toMatch(/Industry\s*:\s*\{\s*group\s*:\s*true\s*\}/s);
			expect(Array.isArray(parsed.variables)).toBe(true);
			expect(typeof parsed.types).toBe("string");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call sf_gql_aggregate with no groupBy and no aggregations defaults to count(Id)", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: { org: ORG, object: "Account" },
			});
			expect(result.isError).toBeFalsy();
			const content = result.content as { type: string; text?: string }[];
			const parsed = JSON.parse(content[0]?.text ?? "{}");
			expect(parsed.query).toMatch(/countId\s*:\s*Id\s*\{\s*count\s*\{\s*value/s);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call sf_gql_aggregate with sum aggregation missing field returns validation error", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: [],
					aggregations: [{ function: "sum" }],
				},
			});
			expect(result.isError).toBe(true);
			const content = result.content as { type: string; text?: string }[];
			expect(content[0]?.text ?? "").toMatch(/field/i);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call accepts a whole-argument $filter placeholder", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
					filter: "$filter",
				},
			});
			expect(result.isError).toBeFalsy();
			const content = result.content as { type: string; text?: string }[];
			const parsed = JSON.parse(content[0]?.text ?? "{}") as {
				query: string;
				variables: { name: string; type: string }[];
			};
			expect(parsed.query).toMatch(/where\s*:\s*\$filter\b/);
			expect(parsed.variables.find((v) => v.name === "filter")?.type).toBe("Account_Filter");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call rejects a malformed whole-argument filter placeholder", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
					filter: "not-a-placeholder",
				},
			});
			expect(result.isError).toBe(true);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/list advertises orderBy/filter/first $var branch inline (no lost $ref branches)", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_aggregate");
			const props = (tool!.inputSchema as { properties: Record<string, unknown> }).properties;

			// The filter/orderBy/first unions must be self-contained: their advertised
			// JSON Schema must not rely on $ref to express the string branch, because
			// the MCP SDK's converter resolves $refs in ways that have dropped that
			// branch (regressing client-side validation of "$var" inputs). The
			// per-use-site varPlaceholder() factory in AGGREGATE_INPUT exists to keep
			// these branches inlined — this test guards against a DRY refactor to a
			// shared instance silently reintroducing $refs on these args.
			// (Note: aggregations[].alias legitimately emits a $ref elsewhere in the
			// schema — a shared field deduped by zod-to-json-schema, unrelated to and
			// harmless for the $var-promotion path — so we assert per-arg, not whole.)
			for (const key of ["orderBy", "filter", "first"]) {
				const schema = JSON.stringify(props[key]);
				expect(schema, `${key} must not express its union via $ref`).not.toContain("$ref");
				expect(schema, `${key} must inline a string $var branch`).toMatch(/"type":"string"/);
				expect(schema, `${key} must carry the $var regex pattern`).toContain("\\\\$"); // pattern ^\$...
			}
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("accepts a whole-argument $first string and $orderBy string", async () => {
		const { client, server } = await connect();
		try {
			const res = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
					first: "$first",
					orderBy: "$orderBy",
				},
			});
			expect(res.isError).toBeFalsy();
			const content = res.content as { type: string; text?: string }[];
			const out = JSON.parse(content[0]?.text ?? "{}") as {
				query: string;
				variables: { name: string; type: string }[];
			};
			expect(out.query).toMatch(/first\s*:\s*\$first\b/);
			expect(out.query).toMatch(/orderBy\s*:\s*\$orderBy\b/);
			expect(out.variables.find((v) => v.name === "first")).toBeDefined();
			expect(out.variables.find((v) => v.name === "orderBy")?.type).toBe("Account_OrderBy");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects an invalid $first placeholder", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: { org: ORG, object: "Account", first: "$1bad" },
			});
			expect(result.isError).toBe(true);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("accepts a JSON-stringified filter object (model serialization tolerance)", async () => {
		const { client, server } = await connect();
		try {
			const res = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
					filter: '{"Industry":{"eq":"$q"}}',
				},
			});
			const content = res.content as { type: string; text?: string }[];
			const out = JSON.parse(content[0]?.text ?? "{}") as {
				query: string;
				variables: { name: string }[];
			};
			// coerced to object, leaf $q promoted
			expect(out.variables.find((v) => v.name === "q")).toBeDefined();
			expect(out.query).toMatch(/Industry\s*:\s*\{\s*eq\s*:\s*\$q/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("accepts a stringified number for first", async () => {
		const { client, server } = await connect();
		try {
			const res = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: { org: ORG, object: "Account", first: "25" },
			});
			const content = res.content as { type: string; text?: string }[];
			const out = JSON.parse(content[0]?.text ?? "{}") as { query: string };
			expect(out.query).toMatch(/first\s*:\s*25\b/);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("accepts a double-quoted $var placeholder (model double-stringify)", async () => {
		const { client, server } = await connect();
		try {
			const res = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: ["Industry"],
					aggregations: [{ function: "count" }],
					filter: '"$filter"',
				},
			});
			const content = res.content as { type: string; text?: string }[];
			const out = JSON.parse(content[0]?.text ?? "{}") as {
				query: string;
				variables: { name: string }[];
			};
			expect(out.query).toMatch(/where\s*:\s*\$filter\b/);
			expect(out.variables.find((v) => v.name === "filter")).toBeDefined();
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects a non-JSON garbage string for filter", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: { org: ORG, object: "Account", filter: "not json{" },
			});
			expect(result.isError).toBe(true);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("rejects garbage string for first (abc)", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: { org: ORG, object: "Account", first: "abc" },
			});
			expect(result.isError).toBe(true);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call rejects invalid GraphQL Names in object/operationName", async () => {
		const { client, server } = await connect();
		try {
			for (const args of [
				{ org: ORG, object: "1Bad" },
				{ org: ORG, object: "Account", operationName: "1Bad" },
			]) {
				const result = await client.callTool({ name: "sf_gql_aggregate", arguments: args });
				expect(result.isError).toBe(true);
				const content = result.content as { type: string; text?: string }[];
				expect(content[0]?.text ?? "").toMatch(/GraphQL Name/);
			}
		} finally {
			await client.close();
			await server.close();
		}
	});
});

// W-23336443: the groupBy `function` enum is a PLAIN z.union member, so it is
// wrapped in enumStripControlChars (a plain enum's invalid_enum_value echoes the
// received value verbatim through the SDK's upstream validation). The aggregate
// `function` DISCRIMINATOR is deliberately NOT wrapped — a bad discriminator
// raises invalid_union_discriminator, which lists only the expected options and
// does NOT echo the input, so there is no reflection channel to close. Both
// claims are proven end-to-end through the real client.callTool channel below.
// Oracle: assert the specific injected code points absent, not a whole-class
// scan (the SDK pretty-prints issues with \p{Cc} newlines). See
// sf-gql-discover.spec.ts for the full rationale.
describe("mcp/tools/sf-gql-aggregate — control-char neutralization (W-23336443)", () => {
	const BIDI = "\u{202e}";
	const ZWSP = "\u{200b}";
	const DEL = "\x7f";

	it("a poisoned-but-valid groupBy function strips to the member and is ACCEPTED", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: [{ field: "Industry", function: `CALENDAR_MONTH${BIDI}` }],
					aggregations: [{ function: "count" }],
				},
			});
			expect(result.isError).toBeFalsy();
			const content = result.content as { type: string; text?: string }[];
			const parsed = JSON.parse(content[0]?.text ?? "{}") as { query: string };
			// Stripped to CALENDAR_MONTH, which renders as the date-bucket function.
			expect(parsed.query).toMatch(/function\s*:\s*CALENDAR_MONTH\b/s);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("a genuinely-invalid poisoned groupBy function rejects with NO raw injected code point", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					groupBy: [{ field: "Industry", function: `${BIDI}bo${ZWSP}gus${DEL}` }],
					aggregations: [{ function: "count" }],
				},
			});
			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);
			expect(text).toContain("bogus");
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("the un-wrapped aggregate function DISCRIMINATOR leaks no injected code point either", async () => {
		const { client, server } = await connect();
		try {
			// Negative control: a poisoned discriminator is NOT stripped (the
			// discriminator is intentionally left un-wrapped), yet the SDK's
			// invalid_union_discriminator issue lists only the expected options and
			// never echoes the received value — so nothing leaks regardless.
			const result = await client.callTool({
				name: "sf_gql_aggregate",
				arguments: {
					org: ORG,
					object: "Account",
					aggregations: [{ function: `count${BIDI}${ZWSP}${DEL}`, field: "Amount" }],
				},
			});
			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);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/list still advertises the groupBy function enum after wrapping", async () => {
		const { client, server } = await connect();
		try {
			const list = await client.listTools();
			const tool = list.tools.find((t) => t.name === "sf_gql_aggregate");
			// groupBy is an array whose items are a union; the object branch carries
			// the wrapped `function` enum. Assert the enum survives somewhere in the
			// advertised groupBy schema (converter shape varies: anyOf/items).
			const groupBy = (tool!.inputSchema as { properties: Record<string, unknown> }).properties
				.groupBy;
			const serialized = JSON.stringify(groupBy);
			for (const fn of ["CALENDAR_MONTH", "FISCAL_YEAR", "WEEK_IN_YEAR"]) {
				expect(serialized).toContain(fn);
			}
		} finally {
			await client.close();
			await server.close();
		}
	});
});