import type { AgentToolResult, ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { StringEnum } from "@earendil-works/pi-ai/compat";
import { Type } from "typebox";
import { normalizeSearchProviderSelection, search, type SearchProviderSelection } from "../search/router.ts";
import type { RecencyFilter } from "../search/types.ts";
import { getConfigOrEmpty } from "../config/store.ts";
import { runWithProxy } from "../net/proxy.ts";

function normalizeQueries(query: unknown, queries: unknown): string[] {
	const values = Array.isArray(queries) ? queries : query === undefined ? [] : [query];
	return values.flatMap(value => typeof value === "string" && value.trim() ? [value.trim()] : []);
}

function normalizeRecency(value: unknown): RecencyFilter | undefined {
	return value === "day" || value === "week" || value === "month" || value === "year" ? value : undefined;
}

function requestedProvider(): SearchProviderSelection {
	const config = getConfigOrEmpty();
	return normalizeSearchProviderSelection(config.searchProvider ?? config.provider);
}

function providerSelectionLabel(provider: SearchProviderSelection): string {
	return Array.isArray(provider) ? provider.join(", ") : provider;
}

function formatSearchResult(query: string, response: Awaited<ReturnType<typeof search>>): string {
	let output = `## ${query}\n\n`;
	if (response.answer.trim()) output += `${response.answer.trim()}\n\n`;
	if (response.results.length > 0) {
		output += "### Sources\n\n";
		for (const [index, result] of response.results.entries()) {
			output += `${index + 1}. [${result.title || result.url}](${result.url})`;
			if (result.snippet) output += ` — ${result.snippet}`;
			output += "\n";
		}
	} else {
		output += `No structured citations returned by ${response.provider}.\n`;
	}
	return output.trim();
}

export function registerWebSearch(pi: ExtensionAPI): void {
	pi.registerTool({
		name: "web_search",
		label: "Web Search",
		description: "Search the web through configurable Providers and return answers with source citations.",
		promptSnippet: "Use for web research. Provider selection and fallback are configured in web-search.json.",
		parameters: Type.Object({
			query: Type.Optional(Type.String({ description: "Single search query" })),
			queries: Type.Optional(Type.Array(Type.String(), { description: "Multiple search queries" })),
			numResults: Type.Optional(Type.Integer({ minimum: 1, maximum: 20, description: "Results per query" })),
			recencyFilter: Type.Optional(StringEnum(["day", "week", "month", "year"])),
			domainFilter: Type.Optional(Type.Array(Type.String(), { description: "Domains to include; prefix with - to exclude" })),
			proxy: Type.Optional(Type.String({ description: "Optional HTTP(S) proxy URL" })),
		}),
		async execute(_callId, params, signal, onUpdate, ctx): Promise<AgentToolResult<Record<string, unknown>>> {
			return runWithProxy(typeof params.proxy === "string" ? params.proxy : undefined, async () => {
				const queries = normalizeQueries(params.query, params.queries);
				if (queries.length === 0) {
					return { content: [{ type: "text", text: "Error: provide query or queries." }], details: { error: "No query provided" } };
				}
				const provider = requestedProvider();
				const results: Array<{ query: string; provider?: string; answer?: string; sourceCount: number; error?: string }> = [];
				const sections: string[] = [];
				for (const [index, query] of queries.entries()) {
					if (signal?.aborted) throw new Error("Aborted");
					onUpdate?.({ content: [{ type: "text", text: `Searching ${index + 1}/${queries.length}: ${query}` }], details: { phase: "search", progress: index / queries.length } });
					try {
						const response = await search(query, {
							provider,
							numResults: params.numResults,
							recencyFilter: normalizeRecency(params.recencyFilter),
							domainFilter: params.domainFilter,
							signal,
							extensionContext: ctx,
						});
						sections.push(formatSearchResult(query, response));
						results.push({ query, provider: response.provider, answer: response.answer, sourceCount: response.results.length });
					} catch (error) {
						if (signal?.aborted || String(error).toLowerCase().includes("abort")) throw error;
						const message = error instanceof Error ? error.message : String(error);
						sections.push(`## ${query}\n\nError: ${message}`);
						results.push({ query, sourceCount: 0, error: message });
					}
				}
				return {
					content: [{ type: "text", text: sections.join("\n\n") }],
					details: {
						provider: providerSelectionLabel(provider),
						queryCount: queries.length,
						successfulQueries: results.filter(result => !result.error).length,
						totalResults: results.reduce((total, result) => total + result.sourceCount, 0),
						results,
					},
				};
			});
		},
		renderCall(args, theme) {
			const input = args as { query?: unknown; queries?: unknown };
			const queries = normalizeQueries(input.query, input.queries);
			if (queries.length <= 1) {
				return new Text(theme.fg("toolTitle", theme.bold("search ")) + theme.fg("accent", queries[0] ?? "(no query)"), 0, 0);
			}
			const lines = [theme.fg("toolTitle", theme.bold("search"))];
			for (const query of queries) lines.push(theme.fg("accent", `  ${query}`));
			return new Text(lines.join("\n"), 0, 0);
		},
		renderResult(result, { isPartial, expanded }, theme) {
			if (isPartial) return new Text(theme.fg("accent", "searching..."), 0, 0);
			const details = result.details as { successfulQueries?: number; queryCount?: number; provider?: string; error?: string };
			const provider = details.provider ? ` [${details.provider}]` : "";
			const status = details.error
				? theme.fg("error", `search failed${provider}`)
				: theme.fg("success", `${details.successfulQueries ?? 0}/${details.queryCount ?? 0} queries${provider}`);
			const text = result.content.find(part => part.type === "text")?.text ?? "";
			if (!text) return new Text(status, 0, 0);
			const limit = expanded ? text.length : 600;
			const preview = text.length > limit ? `${text.slice(0, limit)}...` : text;
			const hint = !expanded && text.length > limit ? theme.fg("muted", "\n\n... (ctrl+o to expand)") : "";
			return new Text(`${status}\n${theme.fg(details.error ? "error" : "toolOutput", preview)}${hint}`, 0, 0);
		},
	});
}