/**
 * tool-ranker — ranks tools for `tool_search` and codemode's `searchTools()` with Sai.
 *
 * Each candidate tool is one yes/no question ("would this tool help with the request?");
 * the model's probability is the score. This is the `toolSearch` route of
 * `saiAdvisory` (on by default). It asks the shared Sai model; set
 * `saiAdvisory.routes.toolSearch.model` to use another decisions model from `./classifier.ts`.
 * The ranker never throws and never blocks discovery: with no
 * session, the route off, no Token-In credential, or any request failing, it returns the local
 * BM25 ranking, and tools the model did not score are appended in BM25 order.
 *
 * Which tools are sent: up to TOOL_RANK_MAX_CANDIDATES in BM25 order, in one fitted request.
 * Unjudged tools keep their local fallback semantics.
 */
import { Bm25Ranker, getSettingsPath, registerToolRanker } from "@selesai/code";
import type {
	ExtensionAPI,
	ToolRankOptions,
	ToolRanker,
	ToolSearchDocument,
	ToolSearchMatch,
} from "@selesai/code";
import { FindRequests, type JudgeItem, type JudgeStage, judgeNouls } from "./ask-tool.ts";
import {
	askSaiAnswers,
	confidenceBucket,
	emitSaiTelemetry,
	type SaiAdvisoryConfig,
	saiConnection,
	type SaiRuntime,
	saiUnavailable,
	readSaiAdvisoryConfig,
} from "./decisions.ts";

// ponytail: semantic recall is limited to 16 BM25-shortlisted candidates and one fitted batch;
// widen only when measured recall justifies the additional cost.
export const TOOL_RANK_MAX_CANDIDATES = 16;
/** Initial classifier dispatches per ranking. */
export const TOOL_RANK_MAX_REQUESTS = 1;
/** At most one retry, and only for an eligible transient failure. */
export const TOOL_RANK_RETRIES = 1;
/** A tool is a match at this probability or above: matches become part of the model's context. */
export const TOOL_RANK_KEEP = 0.5;
/** Characters of a tool's summary shown to Sai. */
export const TOOL_RANK_SUMMARY_CHARS = 480;

const STAGE: Pick<JudgeStage, "stateKey" | "criteria"> = {
	stateKey: "tools",
	criteria: {
		true: "The tool's name and description show it can do what the request asks for, or is a direct step toward it.",
		false: "The tool does something unrelated to the request.",
	},
};

export interface SaiToolRankerOptions {
	/** Receives route telemetry (shape and outcome only; never the query or tool text). */
	events?: { emit(channel: string, data: unknown): void };
	/** Ranks when Sai cannot, and orders the candidates sent past the cap. Default: BM25. */
	fallback?: ToolRanker;
	/** Reads the `saiAdvisory` settings. Default: the agent settings file. */
	readConfig?: () => SaiAdvisoryConfig;
}

/** Stop waiting for credential or classifier work that ignores the operation signal. */
function abortable<T>(work: Promise<T>, signal: AbortSignal): Promise<T> {
	return new Promise<T>((resolve, reject) => {
		const onAbort = () => {
			signal.removeEventListener("abort", onAbort);
			reject(new Error("aborted"));
		};
		if (signal.aborted) onAbort();
		else signal.addEventListener("abort", onAbort, { once: true });
		// Handle a late rejection too, and remove the listener even when work settles first.
		work.then(
			(value) => { signal.removeEventListener("abort", onAbort); resolve(value); },
			(error) => { signal.removeEventListener("abort", onAbort); reject(error); },
		);
	});
}

export function createSaiToolRanker(options: SaiToolRankerOptions = {}): ToolRanker {
	const fallback = options.fallback ?? new Bm25Ranker();
	const readConfig = options.readConfig ?? (() => readSaiAdvisoryConfig(getSettingsPath()));

	return {
		async rank(
			query: string,
			documents: readonly ToolSearchDocument[],
			limit: number,
			rankOptions?: ToolRankOptions,
		): Promise<ToolSearchMatch[]> {
			const started = Date.now();
			const local = await fallback.rank(query, documents, documents.length, rankOptions);
			const ctx = rankOptions?.ctx;
			if (limit <= 0) return [];
			if (!ctx || documents.length === 0 || query.trim() === "" || rankOptions?.signal?.aborted) return local.slice(0, limit);

			const telemetry = (outcome: "sai" | "fallback", extra: Record<string, unknown> = {}) =>
				emitSaiTelemetry(options.events, "decision", {
					route: "toolSearch",
					outcome,
					candidates: documents.length,
					elapsedMs: Date.now() - started,
					...extra,
				});

			let timer: ReturnType<typeof setTimeout> | undefined;
			let signal: AbortSignal | undefined;
			try {
				// Same cast as ask_sai: the session registry is wider than SaiRuntime's header type.
				const runtime = ctx as unknown as SaiRuntime;
				const config = readConfig();
				const route = config.routes.toolSearch;
				if (!route.enabled) return local.slice(0, limit);
				const deadlineAt = started + route.timeoutMs;
				const deadline = new AbortController();
				signal = rankOptions?.signal ? AbortSignal.any([rankOptions.signal, deadline.signal]) : deadline.signal;
				if (Date.now() >= deadlineAt) deadline.abort();
				else timer = setTimeout(() => deadline.abort(), deadlineAt - Date.now());
				if (signal.aborted) throw new Error("aborted");
				const connection = saiConnection(config, route);
				const unreachable = await abortable(saiUnavailable(runtime, connection), signal);
				if (Date.now() >= deadlineAt) deadline.abort();
				if (signal.aborted) throw new Error("aborted");
				if (unreachable) {
					telemetry("fallback", { confidence: confidenceBucket(undefined), reason: unreachable });
					return local.slice(0, limit);
				}

				// BM25 hits first, then the rest in document order: the cap drops the least likely tools.
				const byName = new Map(documents.map((document) => [document.name, document]));
				const hits = new Set(local.map((match) => match.name));
				const order = [...local.map((match) => match.name), ...documents.map((d) => d.name).filter((n) => !hits.has(n))];
				const shortlist = order.slice(0, TOOL_RANK_MAX_CANDIDATES).map((name) => byName.get(name)!);
				const items: JudgeItem[] = shortlist.map((document) => ({
					key: document.name,
					text: (document.summary ?? document.text).slice(0, TOOL_RANK_SUMMARY_CHARS),
					question: `Would calling the tool "${document.name}" help carry out the request?`,
				}));

				const requests = new FindRequests(
					(payload, requestSignal) => askSaiAnswers(runtime, connection, { payload, maxBytes: route.payloadBytes, signal: requestSignal }),
					TOOL_RANK_MAX_REQUESTS + TOOL_RANK_RETRIES,
					{ signal, deadlineAt },
				);
				const outcome = await abortable(
					judgeNouls(requests, { request: query, ...STAGE }, items, route.payloadBytes, {
						send: TOOL_RANK_MAX_REQUESTS,
						retry: TOOL_RANK_MAX_REQUESTS + TOOL_RANK_RETRIES,
					}),
					signal,
				);
				if (Date.now() >= deadlineAt) deadline.abort();
				if (signal.aborted) throw new Error("aborted");

				const unscored = new Set(documents.map((document) => document.name));
				const position = new Map(order.map((name, index) => [name, index]));
				const scored: ToolSearchMatch[] = [];
				items.forEach((item, index) => {
					const score = outcome.scores[index];
					if (typeof score !== "number" || !Number.isFinite(score)) return;
					unscored.delete(item.key);
					if (score >= TOOL_RANK_KEEP) scored.push({ name: item.key, score });
				});
				if (unscored.size === documents.length) {
					telemetry("fallback", { confidence: confidenceBucket(undefined), reason: outcome.reasons[0] ?? "missing" });
					return local.slice(0, limit);
				}
				scored.sort((a, b) => b.score - a.score || (position.get(a.name) ?? 0) - (position.get(b.name) ?? 0));
				const extras = local.filter((match) => unscored.has(match.name)).map((match) => ({ name: match.name, score: 0 }));
				telemetry("sai", {
					judged: documents.length - unscored.size,
					matched: scored.length,
					confidence: confidenceBucket(scored[0]?.score),
					...(unscored.size > 0 ? { unscored: unscored.size } : {}),
				});
				return [...scored, ...extras].slice(0, limit);
			} catch {
				// Discovery must work without Sai: an abort, a transport error, or a bad answer ranks locally.
				telemetry("fallback", { confidence: confidenceBucket(undefined), reason: signal?.aborted ? "timeout" : "transport" });
				return local.slice(0, limit);
			} finally {
				if (timer !== undefined) clearTimeout(timer);
			}
		},
	};
}

export default function saiToolRankerExtension(pi: ExtensionAPI): void {
	const dispose = registerToolRanker(createSaiToolRanker({ events: pi.events }));
	pi.on("session_shutdown", () => {
		dispose();
	});
}