import { afterEach, describe, expect, test } from "vitest";
import { ZaiApiError, zaiSearch } from "../zai.js";

const originalFetch = globalThis.fetch;
const originalZaiKey = process.env.ZAI_API_KEY;
const originalZaiBase = process.env.ZAI_BASE_URL;
const originalZaiSearchBase = process.env.ZAI_SEARCH_BASE_URL;
const originalZaiSearchMcp = process.env.ZAI_SEARCH_MCP_URL;
const originalZaiSearchLocation = process.env.ZAI_SEARCH_LOCATION;
const originalZaiSearchContentSize = process.env.ZAI_SEARCH_CONTENT_SIZE;

afterEach(() => {
	globalThis.fetch = originalFetch;
	if (originalZaiKey === undefined) delete process.env.ZAI_API_KEY;
	else process.env.ZAI_API_KEY = originalZaiKey;
	if (originalZaiBase === undefined) delete process.env.ZAI_BASE_URL;
	else process.env.ZAI_BASE_URL = originalZaiBase;
	if (originalZaiSearchBase === undefined)
		delete process.env.ZAI_SEARCH_BASE_URL;
	else process.env.ZAI_SEARCH_BASE_URL = originalZaiSearchBase;
	if (originalZaiSearchMcp === undefined) delete process.env.ZAI_SEARCH_MCP_URL;
	else process.env.ZAI_SEARCH_MCP_URL = originalZaiSearchMcp;
	if (originalZaiSearchLocation === undefined)
		delete process.env.ZAI_SEARCH_LOCATION;
	else process.env.ZAI_SEARCH_LOCATION = originalZaiSearchLocation;
	if (originalZaiSearchContentSize === undefined)
		delete process.env.ZAI_SEARCH_CONTENT_SIZE;
	else process.env.ZAI_SEARCH_CONTENT_SIZE = originalZaiSearchContentSize;
	delete process.env.ZAI_SEARCH_FRESHNESS;
});

function sse(data: unknown): string {
	return `id:1\nevent:message\ndata:${JSON.stringify(data)}\n\n`;
}

describe("zaiSearch", () => {
	test("calls Coding Plan MCP search and maps results", async () => {
		process.env.ZAI_API_KEY = "test-key";

		const requestUrls: string[] = [];
		let callBody: Record<string, unknown> = {};
		globalThis.fetch = (async (input, init) => {
			requestUrls.push(String(input));
			const body = JSON.parse(String(init?.body));
			expect(init?.headers).toMatchObject({ Authorization: "Bearer test-key" });

			if (body.method === "initialize") {
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2024-11-05" },
					}),
					{
						status: 200,
						headers: {
							"Content-Type": "text/event-stream",
							"mcp-session-id": "session-1",
						},
					},
				);
			}
			if (body.method === "notifications/initialized") {
				expect(init?.headers).toMatchObject({ "mcp-session-id": "session-1" });
				return new Response("", { status: 200 });
			}

			callBody = body;
			expect(init?.headers).toMatchObject({ "mcp-session-id": "session-1" });
			const toolText = JSON.stringify(
				JSON.stringify([
					{
						title: "Example",
						content: "Summary",
						link: "https://example.com",
						publish_date: "2026-05-09",
					},
				]),
			);
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					result: { content: [{ type: "text", text: toolText }] },
				}),
				{ status: 200, headers: { "Content-Type": "text/event-stream" } },
			);
		}) as typeof fetch;

		const results = await zaiSearch("example", undefined, undefined, {
			count: 3,
			freshness: "week",
		});

		expect(requestUrls).toEqual([
			"https://api.z.ai/api/mcp/web_search_prime/mcp",
			"https://api.z.ai/api/mcp/web_search_prime/mcp",
			"https://api.z.ai/api/mcp/web_search_prime/mcp",
		]);
		expect(callBody).toMatchObject({
			jsonrpc: "2.0",
			method: "tools/call",
			params: {
				name: "web_search_prime",
				arguments: {
					search_query: "example",
					content_size: "medium",
					location: "us",
					search_recency_filter: "oneWeek",
				},
			},
		});
		expect(results).toEqual([
			{
				title: "Example",
				url: "https://example.com",
				snippet: "Summary",
				date: "2026-05-09",
			},
		]);
	});

	test("uses ZAI_SEARCH_FRESHNESS default", async () => {
		process.env.ZAI_API_KEY = "test-key";
		process.env.ZAI_SEARCH_FRESHNESS = "month";

		let callArgs: Record<string, unknown> = {};
		globalThis.fetch = (async (_input, init) => {
			const body = JSON.parse(String(init?.body));
			if (body.method === "initialize") {
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2025-03-26" },
					}),
					{
						status: 200,
						headers: { "mcp-session-id": "s" },
					},
				);
			}
			if (body.method === "notifications/initialized")
				return new Response("", { status: 200 });
			callArgs = body.params.arguments;
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					result: { content: [{ type: "text", text: "[]" }] },
				}),
				{ status: 200 },
			);
		}) as typeof fetch;

		await zaiSearch("example", undefined, undefined, { count: 3 });
		expect(callArgs).toMatchObject({ search_recency_filter: "oneMonth" });
	});

	test("normalizes base URLs to the MCP endpoint", async () => {
		process.env.ZAI_API_KEY = "test-key";
		process.env.ZAI_SEARCH_MCP_URL =
			"https://api.z.ai/api/mcp/web_search_prime/mcp";

		let requestUrl = "";
		globalThis.fetch = (async (input, init) => {
			requestUrl = String(input);
			const body = JSON.parse(String(init?.body));
			if (body.method === "initialize")
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2025-03-26" },
					}),
					{
						status: 200,
						headers: { "mcp-session-id": "s" },
					},
				);
			if (body.method === "notifications/initialized")
				return new Response("", { status: 200 });
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					result: { content: [{ type: "text", text: "[]" }] },
				}),
				{ status: 200 },
			);
		}) as typeof fetch;

		await zaiSearch("example", undefined);
		expect(requestUrl).toBe("https://api.z.ai/api/mcp/web_search_prime/mcp");
	});

	test("ZAI_SEARCH_BASE_URL overrides with MCP endpoint normalization", async () => {
		process.env.ZAI_API_KEY = "test-key";
		process.env.ZAI_SEARCH_BASE_URL = "https://api.z.ai/api/paas/v4";

		let requestUrl = "";
		globalThis.fetch = (async (input, init) => {
			requestUrl = String(input);
			const body = JSON.parse(String(init?.body));
			if (body.method === "initialize")
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2025-03-26" },
					}),
					{
						status: 200,
						headers: { "mcp-session-id": "s" },
					},
				);
			if (body.method === "notifications/initialized")
				return new Response("", { status: 200 });
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					result: { content: [{ type: "text", text: "[]" }] },
				}),
				{ status: 200 },
			);
		}) as typeof fetch;

		await zaiSearch("example", undefined);
		expect(requestUrl).toBe("https://api.z.ai/api/mcp/web_search_prime/mcp");
	});

	test("preserves JSON API error code, status, and raw message", async () => {
		process.env.ZAI_API_KEY = "test-key";
		globalThis.fetch = (async () =>
			new Response(
				JSON.stringify({ error: { code: "1002", message: "bad key" } }),
				{ status: 401, statusText: "Unauthorized" },
			)) as typeof fetch;

		try {
			await zaiSearch("example", undefined);
			throw new Error("expected failure");
		} catch (err) {
			expect(err).toBeInstanceOf(ZaiApiError);
			expect("kind" in (err as object)).toBe(false);
			expect((err as ZaiApiError).code).toBe("1002");
			expect((err as ZaiApiError).status).toBe(401);
			expect((err as ZaiApiError).rawMessage).toBe("bad key");
			expect((err as Error).message).toBe("Z.AI API error (1002): bad key");
		}
	});

	test("preserves HTTP plaintext errors", async () => {
		process.env.ZAI_API_KEY = "test-key";
		globalThis.fetch = (async () =>
			new Response("Not Found", {
				status: 404,
				statusText: "Not Found",
			})) as typeof fetch;

		try {
			await zaiSearch("example", undefined);
			throw new Error("expected failure");
		} catch (err) {
			expect(err).toBeInstanceOf(ZaiApiError);
			expect("kind" in (err as object)).toBe(false);
			expect((err as ZaiApiError).code).toBeUndefined();
			expect((err as ZaiApiError).status).toBe(404);
			expect((err as ZaiApiError).rawMessage).toBe("Not Found: Not Found");
			expect((err as Error).message).toBe(
				"Z.AI API error (HTTP 404): Not Found: Not Found",
			);
		}
	});

	test("preserves MCP JSON-RPC errors", async () => {
		process.env.ZAI_API_KEY = "test-key";
		globalThis.fetch = (async (_input, init) => {
			const body = JSON.parse(String(init?.body));
			if (body.method === "initialize")
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2025-03-26" },
					}),
					{
						status: 200,
						headers: { "mcp-session-id": "s" },
					},
				);
			if (body.method === "notifications/initialized")
				return new Response("", { status: 200 });
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					error: { code: -32602, message: "bad params" },
				}),
				{ status: 200 },
			);
		}) as typeof fetch;

		try {
			await zaiSearch("example", undefined);
			throw new Error("expected failure");
		} catch (err) {
			expect(err).toBeInstanceOf(ZaiApiError);
			expect("kind" in (err as object)).toBe(false);
			expect((err as ZaiApiError).code).toBe("-32602");
			expect((err as ZaiApiError).status).toBe(200);
			expect((err as ZaiApiError).rawMessage).toBe("bad params");
			expect((err as Error).message).toBe(
				"Z.AI API error (-32602): bad params",
			);
		}
	});

	test("preserves MCP tool errors without fake HTTP 200 label", async () => {
		process.env.ZAI_API_KEY = "test-key";
		globalThis.fetch = (async (_input, init) => {
			const body = JSON.parse(String(init?.body));
			if (body.method === "initialize")
				return new Response(
					sse({
						jsonrpc: "2.0",
						id: 1,
						result: { protocolVersion: "2025-03-26" },
					}),
					{
						status: 200,
						headers: { "mcp-session-id": "s" },
					},
				);
			if (body.method === "notifications/initialized")
				return new Response("", { status: 200 });
			return new Response(
				sse({
					jsonrpc: "2.0",
					id: 2,
					result: {
						content: [{ type: "text", text: "MCP denied" }],
						isError: true,
					},
				}),
				{ status: 200 },
			);
		}) as typeof fetch;

		try {
			await zaiSearch("example", undefined);
			throw new Error("expected failure");
		} catch (err) {
			expect(err).toBeInstanceOf(ZaiApiError);
			expect("kind" in (err as object)).toBe(false);
			expect((err as ZaiApiError).code).toBeUndefined();
			expect((err as ZaiApiError).status).toBe(200);
			expect((err as ZaiApiError).rawMessage).toBe("MCP denied");
			expect((err as Error).message).toBe("Z.AI API error: MCP denied");
		}
	});
});