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

import fs from "node:fs";
import os from "node:os";
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 } from "graphql";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { makeNoopPrimeDeps } from "../../../__tests__/helpers/prime-deps.js";
import { type PrimeDeps } from "../../../lib/prime-schema.js";
import { registerSfGqlConnectTool } from "../sf-gql-connect.js";

const ORG = "test-tool-connect";
const ORG_URL = "https://test-tool-connect.my.salesforce.com";
const SCHEMA = buildSchema(`type Query { x: Int }`);

// Succeeds on the first download (prime), then fails every subsequent one — lets
// a single server prime a real cache and then fail a refresh over it.
function flakyPrimeDeps(): PrimeDeps {
	const noop = makeNoopPrimeDeps(ORG, ORG_URL, SCHEMA);
	let n = 0;
	return {
		getOrgAuth: noop.getOrgAuth,
		downloadSchema: async (auth) => {
			n++;
			if (n === 1) return noop.downloadSchema(auth);
			throw new Error("introspection failed");
		},
	};
}

function failingPrimeDeps(): PrimeDeps {
	const noop = makeNoopPrimeDeps(ORG, ORG_URL, SCHEMA);
	return {
		getOrgAuth: noop.getOrgAuth,
		downloadSchema: async () => {
			throw new Error("introspection failed");
		},
	};
}

interface ParsedConnect {
	org: string;
	instanceUrl: string;
	refreshed: boolean;
	cached: boolean;
	durationMs: number;
	warnings?: string[];
}

async function connect(
	primeDeps: PrimeDeps = makeNoopPrimeDeps(ORG, ORG_URL, SCHEMA),
): Promise<{ client: Client; server: McpServer }> {
	const server = new McpServer({ name: "graphiti-mcp", version: "test" });
	registerSfGqlConnectTool(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 };
}

function parse(result: Awaited<ReturnType<Client["callTool"]>>): ParsedConnect {
	const content = result.content as { type: string; text?: string }[];
	return JSON.parse(content[0]?.text ?? "{}") as ParsedConnect;
}

describe("mcp/tools/sf-gql-connect", () => {
	let tmpRoot: string;

	beforeEach(() => {
		tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), "graphiti-tool-connect-"));
		process.env.GRAPHITI_HOME = tmpRoot;
	});

	afterEach(() => {
		delete process.env.GRAPHITI_HOME;
		fs.rmSync(tmpRoot, { recursive: true, force: true });
	});

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

	it("primes on first connect and reports cached on the second (forceRefresh defaults to false)", async () => {
		const { client, server } = await connect();
		try {
			const first = parse(
				await client.callTool({ name: "sf_gql_connect", arguments: { org: ORG } }),
			);
			expect(first.cached).toBe(false);
			expect(first.refreshed).toBe(false);
			expect(first.instanceUrl).toBe(ORG_URL);

			const second = parse(
				await client.callTool({ name: "sf_gql_connect", arguments: { org: ORG } }),
			);
			expect(second.cached).toBe(true);
			expect(second.refreshed).toBe(false);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("forceRefresh:true re-downloads and reports refreshed:true", async () => {
		const { client, server } = await connect();
		try {
			await client.callTool({ name: "sf_gql_connect", arguments: { org: ORG } }); // prime
			const refreshed = parse(
				await client.callTool({
					name: "sf_gql_connect",
					arguments: { org: ORG, forceRefresh: true },
				}),
			);
			expect(refreshed.refreshed).toBe(true);
			expect(refreshed.cached).toBe(false);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("refresh failure over a surviving cache returns a soft staleness warning (not isError)", async () => {
		const { client, server } = await connect(flakyPrimeDeps());
		try {
			await client.callTool({ name: "sf_gql_connect", arguments: { org: ORG } }); // prime (ok)
			const result = await client.callTool({
				name: "sf_gql_connect",
				arguments: { org: ORG, forceRefresh: true }, // download now fails
			});
			expect(result.isError).toBeFalsy();
			const out = parse(result);
			expect(out.refreshed).toBe(false);
			expect(out.cached).toBe(true);
			expect(out.warnings?.[0]).toMatch(/keeping the previously cached schema/i);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("refresh failure with no prior cache surfaces an isError envelope", async () => {
		const { client, server } = await connect(failingPrimeDeps());
		try {
			const result = await client.callTool({
				name: "sf_gql_connect",
				arguments: { org: ORG, forceRefresh: true },
			});
			expect(result.isError).toBe(true);
		} finally {
			await client.close();
			await server.close();
		}
	});

	it("tools/call with missing org returns a validation error", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({ name: "sf_gql_connect", arguments: {} });
			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 an org alias containing shell metacharacters", async () => {
		const { client, server } = await connect();
		try {
			const result = await client.callTool({
				name: "sf_gql_connect",
				arguments: { org: "evil; rm -rf /" },
			});
			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();
		}
	});
});