import assert from "node:assert/strict";
import { afterEach, test } from "vitest";
import { googleSearchExtension } from "../resources/extensions/search-the-web/tool-google-search.js";

function createMockPI() {
	const handlers: any[] = [];
	const registeredTool: any = null;

	return {
		handlers,
		registeredTool,
		on(event: string, handler: any) {
			handlers.push({ event, handler });
		},
		registerTool(tool: any) {
			this.registeredTool = tool;
		},
		async fire(event: string, eventData: any, ctx: any) {
			for (const h of handlers) {
				if (h.event === event) {
					await h.handler(eventData, ctx);
				}
			}
		},
	};
}

function mockModelRegistry(oauthJson?: string) {
	return {
		authStorage: {
			hasAuth: async (_id: string) => !!oauthJson,
		},
		getApiKeyForProvider: async (_provider: string) => oauthJson,
	};
}

function restoreEnv(name: string, value: string | undefined) {
	if (value === undefined) delete process.env[name];
	else process.env[name] = value;
}

test("google_search_when_no_code_assist_oauth_warns_and_falls_back_or_errors", async () => {
	const originalKey = process.env.GEMINI_API_KEY;
	const originalAlias = process.env.GOOGLE_GENERATIVE_AI_API_KEY;
	const originalTavily = process.env.TAVILY_API_KEY;
	const originalBrave = process.env.BRAVE_API_KEY;
	const originalMinimax = process.env.MINIMAX_API_KEY;
	delete process.env.GEMINI_API_KEY;
	delete process.env.GOOGLE_GENERATIVE_AI_API_KEY;
	process.env.TAVILY_API_KEY = "";
	delete process.env.BRAVE_API_KEY;
	delete process.env.MINIMAX_API_KEY;

	afterEach(() => {
		restoreEnv("GEMINI_API_KEY", originalKey);
		restoreEnv("GOOGLE_GENERATIVE_AI_API_KEY", originalAlias);
		restoreEnv("TAVILY_API_KEY", originalTavily);
		restoreEnv("BRAVE_API_KEY", originalBrave);
		restoreEnv("MINIMAX_API_KEY", originalMinimax);
	});

	const pi = createMockPI();
	googleSearchExtension(pi as any);

	const notifications: any[] = [];
	const mockCtx = {
		ui: {
			notify(msg: string, level: string) {
				notifications.push({ msg, level });
			},
		},
		modelRegistry: mockModelRegistry(undefined),
	};

	await pi.fire("session_start", {}, mockCtx);
	assert.equal(notifications.length, 1);
	assert.ok(notifications[0].msg.includes("No google-gemini-cli OAuth set"));

	const registeredTool = (pi as any).registeredTool;
	const result = await registeredTool.execute(
		"call-1",
		{ query: "test" },
		new AbortController().signal,
		() => {},
		mockCtx,
	);
	assert.equal(result.isError, true);
	assert.ok(result.content[0].text.includes("No Google Code Assist OAuth"));
});

test("google_search_ignores_genai_api_key_env_and_requires_code_assist_oauth", async () => {
	const originalKey = process.env.GEMINI_API_KEY;
	const originalAlias = process.env.GOOGLE_GENERATIVE_AI_API_KEY;
	process.env.GEMINI_API_KEY = "mock-api-key";
	process.env.GOOGLE_GENERATIVE_AI_API_KEY = "mock-alias-key";

	afterEach(() => {
		restoreEnv("GEMINI_API_KEY", originalKey);
		restoreEnv("GOOGLE_GENERATIVE_AI_API_KEY", originalAlias);
	});

	const pi = createMockPI();
	googleSearchExtension(pi as any);

	const notifications: any[] = [];
	const mockCtx = {
		ui: {
			notify(msg: string, level: string) {
				notifications.push({ msg, level });
			},
		},
		modelRegistry: mockModelRegistry(undefined),
	};

	await pi.fire("session_start", {}, mockCtx);
	assert.equal(notifications.length, 1);
	assert.ok(notifications[0].msg.includes("google-gemini-cli OAuth"));
});
