/**
 * Google Search Extension
 *
 * Provides a `google_search` tool that performs web searches via Gemini's
 * Google Search grounding feature through @google/gemini-cli-core.
 *
 * The tool sends queries to Gemini Flash with `googleSearch: {}` enabled.
 * Gemini internally performs Google searches, synthesizes an answer, and
 * returns it with source URLs from grounding metadata.
 */
import { Type } from "@sinclair/typebox";
import {
	DEFAULT_MAX_BYTES,
	DEFAULT_MAX_LINES,
	formatSize,
	truncateHead,
} from "@singularity-forge/coding-agent";
import { Text } from "@singularity-forge/tui";
import {
	getBraveApiKey,
	getTavilyApiKey,
	resolveSearchProvider,
} from "./provider.js";

/**
 * Build a Code Assist server through @google/gemini-cli-core.
 *
 * The OAuth fallback uses cli-core directly instead of carrying SF's own
 * Cloud Code Assist wire client.
 */
async function buildCodeAssistServer(accessToken, projectId) {
	const [{ CodeAssistServer }, { OAuth2Client }] = await Promise.all([
		import("@google/gemini-cli-core"),
		import("google-auth-library"),
	]);
	const authClient = new OAuth2Client();
	authClient.setCredentials({ access_token: accessToken });
	return new CodeAssistServer(authClient, projectId, { headers: {} });
}
/**
 * Perform a search using OAuth credentials via @google/gemini-cli-core.
 */
async function searchWithOAuth(query, accessToken, projectId, signal) {
	const model = process.env.GEMINI_SEARCH_MODEL || "gemini-2.5-flash";
	const server = await buildCodeAssistServer(accessToken, projectId);
	const promptId = `sf-google-search-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;
	const chunks = server.generateContentStream(
		{
			model,
			contents: [{ role: "user", parts: [{ text: query }] }],
			config: {
				tools: [{ googleSearch: {} }],
			},
		},
		promptId,
		"USER",
	);
	let answer = "";
	let grounding;
	for await (const chunk of chunks) {
		if (signal?.aborted) {
			throw new Error("Request was aborted");
		}
		const candidate = chunk?.candidates?.[0];
		const text = candidate?.content?.parts?.find((p) => p.text)?.text;
		if (text) {
			answer += text;
		}
		if (candidate?.groundingMetadata) {
			grounding = candidate.groundingMetadata;
		}
	}
	const sources = [];
	const seenTitles = new Set();
	if (grounding?.groundingChunks) {
		for (const chunk of grounding.groundingChunks) {
			if (chunk.web) {
				const title = chunk.web.title ?? "Untitled";
				if (seenTitles.has(title)) continue;
				seenTitles.add(title);
				const domain = chunk.web.domain ?? title;
				sources.push({
					title,
					uri: chunk.web.uri ?? "",
					domain,
				});
			}
		}
	}
	const searchQueries = grounding?.webSearchQueries ?? [];
	return { answer, sources, searchQueries, cached: false };
}

// ── Fallback search via search-the-web providers ────────────────────────────
async function executeTavilyFallback(query, signal) {
	const response = await fetch("https://api.tavily.com/search", {
		method: "POST",
		headers: {
			"Content-Type": "application/json",
			Authorization: `Bearer ${getTavilyApiKey()}`,
		},
		body: JSON.stringify({
			query,
			max_results: 10,
			search_depth: "basic",
			include_answer: true,
		}),
		signal,
	});
	if (!response.ok) {
		throw new Error(`Tavily fallback failed: ${response.status}`);
	}
	const data = await response.json();
	const answer = data.answer || "";
	const sources = (data.results || [])
		.filter((r) => typeof r.url === "string" && r.url.length > 0)
		.map((r) => ({
			title: r.title || "(untitled)",
			uri: r.url,
			domain: r.url ? new URL(r.url).hostname : "",
		}));
	return { answer, sources, searchQueries: [query], cached: false };
}

async function executeBraveFallback(query, signal) {
	const url = new URL("https://api.search.brave.com/res/v1/web/search");
	url.searchParams.append("q", query);
	url.searchParams.append("count", "10");
	url.searchParams.append("extra_snippets", "true");
	url.searchParams.append("text_decorations", "false");
	const response = await fetch(url.toString(), {
		method: "GET",
		headers: {
			Accept: "application/json",
			"Accept-Encoding": "gzip",
			"X-Subscription-Token": getBraveApiKey(),
		},
		signal,
	});
	if (!response.ok) {
		throw new Error(`Brave fallback failed: ${response.status}`);
	}
	const data = await response.json();
	const rawResults = data.web?.results ?? [];
	const answerParts = [];
	if (data.summary && Array.isArray(data.summary)) {
		answerParts.push(
			data.summary
				.filter((s) => s.type === "token" || s.type === "text")
				.map((s) => s.data)
				.join(""),
		);
	}
	const answer = answerParts.join("\n");
	const sources = rawResults
		.filter((r) => typeof r.url === "string" && r.url.length > 0)
		.map((r) => ({
			title: r.title || "(untitled)",
			uri: r.url,
			domain: r.url ? new URL(r.url).hostname : "",
		}));
	return { answer, sources, searchQueries: [query], cached: false };
}

async function executeFallbackSearch(query, signal) {
	const provider = resolveSearchProvider();
	if (!provider || provider === "combosearch") {
		throw new Error("No fallback search provider available");
	}
	if (provider === "tavily") {
		return executeTavilyFallback(query, signal);
	}
	if (provider === "brave") {
		return executeBraveFallback(query, signal);
	}
	// For other providers, try Tavily if available, else Brave
	if (getTavilyApiKey()) {
		return executeTavilyFallback(query, signal);
	}
	if (getBraveApiKey()) {
		return executeBraveFallback(query, signal);
	}
	throw new Error("No fallback search provider available");
}

function cacheKey(query) {
	return query.toLowerCase().trim();
}
// ── Tool registration ────────────────────────────────────────────────────────
export function googleSearchExtension(pi) {
	return registerGoogleSearchTool(pi);
}
export function registerGoogleSearchTool(pi) {
	pi.registerTool({
		name: "google_search",
		label: "Google Search",
		description:
			"Search the web using Google Search via Gemini. " +
			"Returns an AI-synthesized answer grounded in Google Search results, plus source URLs. " +
			"Use this when you need current information from the web: recent events, documentation, " +
			"product details, technical references, news, etc. " +
			"Requires Google Code Assist OAuth via the google-gemini-cli provider. Alternative to Brave-based search tools.",
		promptSnippet:
			"Search the web via Google Search to get current information with sources",
		promptGuidelines: [
			"Use google_search when you need up-to-date web information that isn't in your training data.",
			"Be specific with queries for better results, e.g. 'Next.js 15 app router migration guide' not just 'Next.js'.",
			"The tool returns both an answer and source URLs. Cite sources when sharing results with the user.",
			"Results are cached per-session, so repeated identical queries are free.",
			"You can still use fetch_page to read a specific URL if needed after getting results from google_search.",
		],
		parameters: Type.Object({
			query: Type.String({
				description:
					"The search query, e.g. 'latest Node.js LTS version' or 'how to configure Tailwind v4'",
			}),
			maxSources: Type.Optional(
				Type.Number({
					description:
						"Maximum number of source URLs to include (default 5, max 10).",
					minimum: 1,
					maximum: 10,
				}),
			),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			const startTime = Date.now();
			const maxSources = Math.min(Math.max(params.maxSources ?? 5, 1), 10);
			let oauthToken;
			let projectId;
			const oauthRaw =
				await ctx.modelRegistry.getApiKeyForProvider("google-gemini-cli");
			if (oauthRaw) {
				try {
					const parsed = JSON.parse(oauthRaw);
					oauthToken = parsed.token;
					projectId = parsed.projectId;
				} catch {
					// Fall through to error
				}
			}
			if (!oauthToken || !projectId) {
				// No Gemini credentials — try fallback through search-the-web providers
				try {
					const fallbackResult = await executeFallbackSearch(
						params.query,
						signal,
					);
					resultCache.set(cacheKey(params.query), fallbackResult);
					const rawOutput = formatOutput(fallbackResult, maxSources);
					const truncation = truncateHead(rawOutput, {
						maxLines: DEFAULT_MAX_LINES,
						maxBytes: DEFAULT_MAX_BYTES,
					});
					let finalText = truncation.content;
					if (truncation.truncated) {
						finalText +=
							`\n\n[Truncated: showing ${truncation.outputLines}/${truncation.totalLines} lines` +
							` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`;
					}
					return {
						content: [{ type: "text", text: finalText }],
						details: {
							query: params.query,
							sourceCount: fallbackResult.sources.length,
							cached: false,
							durationMs: Date.now() - startTime,
							fallbackProvider: resolveSearchProvider(),
						},
					};
				} catch {
					return {
						content: [
							{
								type: "text",
								text: "Error: No Google Code Assist OAuth found for google_search. Configure the google-gemini-cli provider, or use a non-Google fallback search provider.",
							},
						],
						isError: true,
						details: {
							query: params.query,
							sourceCount: 0,
							cached: false,
							durationMs: Date.now() - startTime,
							error: "auth_error: No credentials set",
						},
					};
				}
			}
			// Check cache
			const key = cacheKey(params.query);
			if (resultCache.has(key)) {
				const cached = resultCache.get(key);
				const output = formatOutput(cached, maxSources);
				return {
					content: [{ type: "text", text: output }],
					details: {
						query: params.query,
						sourceCount: cached.sources.length,
						cached: true,
						durationMs: Date.now() - startTime,
					},
				};
			}
			// Call Gemini with Google Search grounding
			let result;
			try {
				result = await searchWithOAuth(
					params.query,
					oauthToken,
					projectId,
					signal,
				);
			} catch (err) {
				const msg = err instanceof Error ? err.message : String(err);
				let errorType = "api_error";
				if (msg.includes("401") || msg.includes("UNAUTHENTICATED")) {
					errorType = "auth_error";
				} else if (
					msg.includes("429") ||
					msg.includes("RESOURCE_EXHAUSTED") ||
					msg.includes("quota")
				) {
					errorType = "rate_limit";
				}
				return {
					content: [
						{
							type: "text",
							text: `Google Search failed (${errorType}): ${msg}`,
						},
					],
					isError: true,
					details: {
						query: params.query,
						sourceCount: 0,
						cached: false,
						durationMs: Date.now() - startTime,
						error: `${errorType}: ${msg}`,
					},
				};
			}
			// Cache the result
			resultCache.set(key, result);
			// Format and truncate output
			const rawOutput = formatOutput(result, maxSources);
			const truncation = truncateHead(rawOutput, {
				maxLines: DEFAULT_MAX_LINES,
				maxBytes: DEFAULT_MAX_BYTES,
			});
			let finalText = truncation.content;
			if (truncation.truncated) {
				finalText +=
					`\n\n[Truncated: showing ${truncation.outputLines}/${truncation.totalLines} lines` +
					` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`;
			}
			return {
				content: [{ type: "text", text: finalText }],
				details: {
					query: params.query,
					sourceCount: result.sources.length,
					cached: false,
					durationMs: Date.now() - startTime,
				},
			};
		},
		renderCall(args, theme) {
			let text = theme.fg("toolTitle", theme.bold("google_search "));
			text += theme.fg("accent", `"${args.query}"`);
			return new Text(text, 0, 0);
		},
		renderResult(result, { isPartial, expanded }, theme) {
			const d = result.details;
			if (isPartial)
				return new Text(theme.fg("warning", "Searching Google..."), 0, 0);
			if (result.isError || d?.error) {
				return new Text(
					theme.fg("error", `Error: ${d?.error ?? "unknown"}`),
					0,
					0,
				);
			}
			let text = theme.fg("success", `${d?.sourceCount ?? 0} sources`);
			text += theme.fg("dim", ` (${d?.durationMs ?? 0}ms)`);
			if (d?.cached) text += theme.fg("dim", " · cached");
			if (expanded) {
				const content = result.content[0];
				if (content?.type === "text") {
					const preview = content.text.split("\n").slice(0, 8).join("\n");
					text += "\n\n" + theme.fg("dim", preview);
					if (content.text.split("\n").length > 8) {
						text += "\n" + theme.fg("muted", "...");
					}
				}
			}
			return new Text(text, 0, 0);
		},
	});
	// ── Session cleanup ─────────────────────────────────────────────────────
	pi.on("session_shutdown", async () => {
		resultCache.clear();
	});
	// ── Startup notification ─────────────────────────────────────────────────
	// Defensive: Only check for OAuth if modelRegistry and authStorage exist (test mocks may omit)
	pi.on("session_start", async (_event, ctx) => {
		if (ctx?.modelRegistry?.authStorage?.hasAuth) {
			const hasOAuth =
				await ctx.modelRegistry.authStorage.hasAuth("google-gemini-cli");
			if (!hasOAuth && ctx?.ui?.notify) {
				ctx.ui.notify(
					"Google Search: No google-gemini-cli OAuth set. Configure Google Code Assist OAuth to use google_search.",
					"warning",
				);
			}
		}
	});
}
const resultCache = new Map();
/** Reset the google_search result cache (called by search-the-web extension on session_start). */
export function resetGoogleSearchCache() {
	resultCache.clear();
}
// ── Output formatting ────────────────────────────────────────────────────────
function formatOutput(result, maxSources) {
	const lines = [];
	// Answer
	if (result.answer) {
		lines.push(result.answer);
	} else {
		lines.push("(No answer text returned from search)");
	}
	// Sources
	if (result.sources.length > 0) {
		lines.push("");
		lines.push("Sources:");
		const sourcesToShow = result.sources.slice(0, maxSources);
		for (let i = 0; i < sourcesToShow.length; i++) {
			const s = sourcesToShow[i];
			lines.push(`[${i + 1}] ${s.title} - ${s.domain}`);
			lines.push(`    ${s.uri}`);
		}
		if (result.sources.length > maxSources) {
			lines.push(
				`(${result.sources.length - maxSources} more sources omitted)`,
			);
		}
	} else {
		lines.push("");
		lines.push("(No source URLs found in grounding metadata)");
	}
	// Search queries
	if (result.searchQueries.length > 0) {
		lines.push("");
		lines.push(
			`Searches performed: ${result.searchQueries.map((q) => `"${q}"`).join(", ")}`,
		);
	}
	return lines.join("\n");
}
