/**
 * web-search-middleware.ts — native provider middleware for Anthropic web search injection.
 *
 * Purpose: inject the web_search tool into Anthropic API requests before they are sent,
 * bypassing the jiti extension dispatch hot-path that fires on every LLM API call.
 * Placing injection here means it runs as compiled TypeScript with zero extension-dispatch
 * overhead, while remaining available to the search-the-web extension for delegation and
 * test coverage.
 *
 * Consumer: sdk.ts onPayload callback (native hot-path) and the search-the-web extension's
 * before_provider_request handler (delegation for test coverage and fallback).
 */

/** All custom search tool names that should be disabled when native search is active. */
export const CUSTOM_SEARCH_TOOL_NAMES = [
	"search-the-web",
	"search_and_read",
	"google_search",
];

/**
 * Maximum number of native web searches allowed per session (agent unit).
 *
 * Purpose: the Anthropic API's max_uses resets per-request. Without a session-level
 * cap, pause_turn → resubmit cycles allow unbounded total searches (#1309).
 * 15 = 3 full turns of 5 searches each — generous for research, but bounded.
 */
export const MAX_NATIVE_SEARCHES_PER_SESSION = 15;

/** Thinking block types that require signature validation by the API. */
const THINKING_TYPES = new Set(["thinking", "redacted_thinking"]);

/**
 * Module-level preference resolver, set by the search-the-web extension to expose
 * PREFERENCES.md-based search_provider overrides. Falls back to env-var only when unset.
 *
 * Purpose: decouple the middleware from the extension's preference system while still
 * respecting user-configured search provider overrides at injection time.
 */
let _preferBraveResolver: (() => boolean) | undefined;

/**
 * Register a function that determines whether to prefer Brave/custom search over
 * native Anthropic web search. Called by the search-the-web extension on activation
 * to expose PREFERENCES.md-based search_provider overrides to the middleware.
 *
 * Consumer: search-the-web extension native-search.js, on registration.
 */
export function setPreferBraveResolver(fn: () => boolean): void {
	_preferBraveResolver = fn;
}

/**
 * Returns true when the user prefers Brave or another custom search provider over
 * native Anthropic web search. Delegates to the registered resolver if available,
 * otherwise falls back to the PREFER_BRAVE_SEARCH env var.
 */
function preferBraveSearch(): boolean {
	if (_preferBraveResolver) return _preferBraveResolver();
	return (
		process.env.PREFER_BRAVE_SEARCH === "1" ||
		process.env.PREFER_BRAVE_SEARCH === "true"
	);
}

/**
 * Strip thinking/redacted_thinking blocks from assistant messages in the conversation
 * history.
 *
 * Purpose: the Pi SDK's streaming parser drops server_tool_use and
 * web_search_tool_result content blocks (unknown types). When the conversation is
 * replayed, assistant messages are incomplete — the Anthropic API detects the
 * modification and rejects the request with "thinking blocks cannot be modified."
 * Removing thinking blocks from history avoids the rejection; the model generates
 * fresh thinking for the current turn.
 *
 * Consumer: WebSearchMiddleware.applyToPayload, called on every Anthropic request.
 */
export function stripThinkingFromHistory(messages: unknown[]): void {
	for (const msg of messages as Array<{ role?: string; content?: unknown }>) {
		if (msg.role !== "assistant") continue;
		const content = msg.content;
		if (!Array.isArray(content)) continue;
		msg.content = content.filter(
			(block: unknown) =>
				!THINKING_TYPES.has((block as { type?: string })?.type ?? ""),
		);
	}
}

/**
 * Stateful per-session web search middleware.
 *
 * Purpose: inject the web_search tool definition into Anthropic API requests and enforce
 * a session-level search budget so pause_turn → resubmit loops cannot accumulate
 * unbounded searches (#1309). A high-water mark prevents budget resets when context
 * compaction removes web_search_tool_result blocks from the visible history.
 *
 * Consumer: sdk.ts createAgentSession (singleton instance, native hot-path) and the
 * search-the-web extension's before_provider_request handler (delegation).
 */
export class WebSearchMiddleware {
	private sessionSearchCount = 0;

	/**
	 * Reset the session-level search counter to zero.
	 *
	 * Purpose: prevent budget from carrying over across sessions. Must be called on
	 * session_start and on each registerNativeSearchHooks invocation (test isolation).
	 *
	 * Consumer: session_start handler in native-search.js; registerNativeSearchHooks reset.
	 */
	resetSession(): void {
		this.sessionSearchCount = 0;
	}

	/**
	 * Apply Anthropic web search injection to an API request payload. Mutates the
	 * payload in-place and returns it when changes were made, or undefined when no
	 * changes were needed (matching the extension before_provider_request return
	 * convention so the extension runner does not overwrite the payload reference).
	 *
	 * Consumer: sdk.ts onPayload callback (return value ignored, mutation used) and
	 * search-the-web extension before_provider_request handler (return value forwarded).
	 */
	applyToPayload(
		payload: unknown,
		model?: { provider?: string; id?: string },
	): unknown {
		if (!payload) return undefined;

		const p = payload as Record<string, unknown>;

		// Determine whether this is an Anthropic request. The model argument is the
		// most reliable source (comes directly from the resolved Model). Without it,
		// fall back to the model name heuristic, excluding known Claude-compatible
		// providers that do NOT support the web_search tool (Copilot, MiniMax, Kimi).
		let isAnthropic: boolean;
		if (model?.provider) {
			isAnthropic = model.provider === "anthropic";
		} else {
			const name = String(p.model ?? "").toLowerCase();
			isAnthropic =
				name.startsWith("claude-") &&
				!name.includes("minimax") &&
				!name.includes("kimi") &&
				!name.includes("copilot");
		}

		if (!isAnthropic) return undefined;

		// Strip thinking blocks from history to avoid signature validation errors
		// caused by the SDK dropping server_tool_use/web_search_tool_result blocks.
		const messages = p.messages;
		if (Array.isArray(messages)) {
			stripThinkingFromHistory(messages as unknown[]);
		}

		// When preferring Brave/custom search, skip native injection entirely.
		if (preferBraveSearch()) return undefined;

		if (!Array.isArray(p.tools)) p.tools = [];
		let tools = p.tools as Array<Record<string, unknown>>;

		// Don't double-inject if already present (e.g., sdk.ts ran before extension).
		if (tools.some((t) => t.type === "web_search_20250305")) return undefined;

		// Remove custom search tool definitions — native web_search is server-side
		// and more reliable; keeping both confuses the model (#custom-search-conflict).
		tools = tools.filter(
			(t) => !CUSTOM_SEARCH_TOOL_NAMES.includes(t.name as string),
		);
		p.tools = tools;

		// ── Session-level search budget (#1309, compaction-safe) ──────────────────
		// Count web_search_tool_result blocks in history to determine how many native
		// searches have already been used this session. Use the monotonic high-water
		// mark: take the max of the history count and the running counter so budget
		// does not reset when compaction removes result blocks from the visible history.
		if (Array.isArray(messages)) {
			let historySearchCount = 0;
			for (const msg of messages as Array<{ content?: unknown }>) {
				const content = msg.content;
				if (!Array.isArray(content)) continue;
				for (const block of content) {
					if ((block as { type?: string })?.type === "web_search_tool_result") {
						historySearchCount++;
					}
				}
			}
			this.sessionSearchCount = Math.max(
				this.sessionSearchCount,
				historySearchCount,
			);
		}

		const remaining = Math.max(
			0,
			MAX_NATIVE_SEARCHES_PER_SESSION - this.sessionSearchCount,
		);

		if (remaining <= 0) {
			// Budget exhausted — don't inject the search tool at all.
			// Return payload to signal it was processed (custom tools were filtered).
			return payload;
		}

		tools.push({
			type: "web_search_20250305",
			name: "web_search",
			// Cap per-request searches to the lesser of 5 (per-turn cap) or the
			// remaining session budget (#1309).
			max_uses: Math.min(5, remaining),
		});

		return payload;
	}
}

/**
 * Module-level singleton instance shared by sdk.ts (native hot-path) and the
 * search-the-web extension (delegation + session reset).
 *
 * Purpose: a single counter ensures the session budget is accurate regardless of
 * which code path processes each request. The extension resets the counter on
 * session_start and on each registerNativeSearchHooks call (test isolation).
 *
 * Consumer: sdk.ts onPayload, search-the-web extension native-search.js.
 */
export const webSearchMiddleware = new WebSearchMiddleware();
