pi-ext

import { setTimeout as delay } from "node:timers/promises";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { codexSearch, isCodexAvailable } from "./codex.js";
import type { SearchResult } from "./constants.js";
import {
	BROWSER_HEADERS,
	DDG_HTML_URL,
	DEFAULT_TIMEOUT_MS,
} from "./constants.js";
import { markSearchTime, parseDdgResults, throttleWait } from "./ddg.js";
import { fetchSearchHtml } from "./fetch.js";
import { googleSearch } from "./google.js";
import { isMinimaxAvailable, minimaxSearch } from "./minimax.js";
import { parallelSearch } from "./parallel.js";
import { isZaiAvailable, zaiSearch } from "./zai.js";

export type SearchSource =
	| "codex"
	| "minimax"
	| "parallel"
	| "google"
	| "ddg"
	| "zai";

interface SearchProvider {
	source: SearchSource;
	label: "Codex" | "Z.AI" | "MiniMax" | "Parallel" | "Google WML" | "DDG";
	search(): Promise<SearchResult[]>;
}

interface SearchFailure {
	label: SearchProvider["label"];
	err: unknown;
	message: string;
}

const MAX_ERROR_TEXT = 240;

function displayError(err: unknown): string {
	const message = (err instanceof Error ? err.message : String(err))
		.replace(/\s+/g, " ")
		.trim();
	if (!message) return "unknown error";
	if (message.length <= MAX_ERROR_TEXT) return message;
	return `${message.slice(0, MAX_ERROR_TEXT - 1)}…`;
}

export interface WebSearchOptions {
	query: string;
	count: number;
	timeout: number;
	signal?: AbortSignal;
	ctx?: ExtensionContext;
}

export interface WebSearchOutcome {
	results: SearchResult[];
	source: SearchSource;
}

export async function searchWeb(
	options: WebSearchOptions,
): Promise<WebSearchOutcome> {
	options.signal?.throwIfAborted();
	const providers = await buildProviders(options);
	const failures: SearchFailure[] = [];

	for (const provider of providers) {
		options.signal?.throwIfAborted();
		try {
			return { results: await provider.search(), source: provider.source };
		} catch (err) {
			options.signal?.throwIfAborted();
			failures.push({
				label: provider.label,
				err,
				message: displayError(err),
			});
		}
	}

	const summary = failures
		.map((failure) => `${failure.label}: ${failure.message}`)
		.join(" | ");

	throw new AggregateError(
		failures.map((failure) => failure.err),
		`Search unavailable: ${summary}`,
	);
}

export async function buildProviders(
	options: WebSearchOptions,
): Promise<SearchProvider[]> {
	const providers: SearchProvider[] = [
		{
			source: "google",
			label: "Google WML",
			search: () =>
				googleSearch(
					options.query,
					options.signal,
					options.count,
					options.timeout,
				),
		},
		{
			source: "parallel",
			label: "Parallel",
			search: () =>
				parallelSearch(
					options.query,
					options.signal,
					options.count,
					options.timeout,
				),
		},
	];

	const ctx = options.ctx;
	if (ctx && (await isCodexAvailable(ctx))) {
		providers.push({
			source: "codex",
			label: "Codex",
			search: () =>
				codexSearch(options.query, ctx, options.signal, options.count),
		});
	}

	if (ctx && (await isMinimaxAvailable(ctx))) {
		providers.push({
			source: "minimax",
			label: "MiniMax",
			search: () => searchMinimax(options.query, ctx, options.signal),
		});
	}

	providers.push({
		source: "ddg",
		label: "DDG",
		search: () =>
			searchDdg(options.query, options.signal, options.timeout, options.ctx),
	});

	if (await isZaiAvailable(options.ctx)) {
		providers.push({
			source: "zai",
			label: "Z.AI",
			search: () =>
				zaiSearch(options.query, options.ctx, options.signal, {
					count: options.count,
					timeoutMs: options.timeout,
				}),
		});
	}

	return providers;
}

async function searchMinimax(
	query: string,
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<SearchResult[]> {
	const response = await minimaxSearch(query, ctx, signal);
	return (response.organic ?? []).map((r) => ({
		title: r.title,
		url: r.link,
		snippet: r.snippet,
		date: r.date,
	}));
}

async function searchDdg(
	query: string,
	signal?: AbortSignal,
	timeout?: number,
	ctx?: ExtensionContext,
): Promise<SearchResult[]> {
	const deadline = AbortSignal.timeout(timeout ?? DEFAULT_TIMEOUT_MS);
	const requestSignal = signal ? AbortSignal.any([signal, deadline]) : deadline;
	requestSignal.throwIfAborted();
	const waitMs = throttleWait();
	if (waitMs > 0) {
		ctx?.ui?.setStatus(
			"web",
			`⏳ DDG rate limit: ${(waitMs / 1000).toFixed(1)}s`,
		);
		try {
			await delay(waitMs, undefined, { signal: requestSignal });
		} finally {
			ctx?.ui?.setStatus("web", undefined);
		}
	}

	const html = await fetchSearchHtml(DDG_HTML_URL, requestSignal, timeout, {
		method: "POST",
		headers: {
			...BROWSER_HEADERS,
			"Content-Type": "application/x-www-form-urlencoded",
			Origin: new URL(DDG_HTML_URL).origin,
			Referer: new URL("/", DDG_HTML_URL).href,
		},
		body: new URLSearchParams({ q: query, b: "", kl: "wt-wt" }),
	});
	markSearchTime();

	const parsed = parseDdgResults(html);
	if (parsed.error) throw new Error(parsed.error);
	return parsed.results.map((r) => ({
		title: r.title,
		url: r.url,
		snippet: r.snippet,
	}));
}