pi-ext

import { mkdir, readFile, rm, writeFile } from "node:fs/promises";
import { join } from "node:path";

export interface BidiResponse {
	id: number;
	method?: string;
	type?: string;
	result?: unknown;
	error?: string;
	message?: string;
	skipped?: boolean;
}

export interface Frame {
	id: number;
	method: string;
	params?: Record<string, unknown>;
}

export interface BatchCtx {
	pending: Map<number, (r: BidiResponse) => void>;
	ledger: {
		subscriptions: string[];
		intercepts: string[];
		preloadScripts: string[];
	};
	spoolDir: string;
	spoolBytes?: number;
	spoolBudget?: { bytes: number };
}

const MAX_INLINE = 256 * 1024;
const MAX_PAYLOAD = 1024 * 1024;
export const MAX_SPOOL = 32 * 1024 * 1024;
const RESPONSE_BUDGET = MAX_PAYLOAD - 128 * 1024;

function payloadBudgetError(frame: Frame): BidiResponse {
	return {
		id: frame.id,
		method: frame.method,
		type: "error",
		error: "pi payload budget exceeded",
	};
}

function responseCost(text: string, index: number): number {
	return Buffer.byteLength(text) + (index ? 1 : 0);
}

function reserveSpool(b: BatchCtx, bytes: number): void {
	b.spoolBytes = (b.spoolBytes ?? 0) + bytes;
	if (b.spoolBudget) b.spoolBudget.bytes += bytes;
}

function releaseSpool(b: BatchCtx, bytes: number): void {
	b.spoolBytes = (b.spoolBytes ?? 0) - bytes;
	if (b.spoolBudget) b.spoolBudget.bytes -= bytes;
}

export function ledgerUpdate(
	b: Pick<BatchCtx, "ledger">,
	method: string,
	params: Record<string, unknown>,
	res: BidiResponse,
): void {
	const ok = res.type === "success";
	const result = (res.result ?? {}) as Record<string, unknown>;
	if (method === "session.subscribe" && ok) {
		for (const e of (params.events as string[]) ?? [])
			if (!b.ledger.subscriptions.includes(e)) b.ledger.subscriptions.push(e);
	} else if (method === "session.unsubscribe") {
		const gone = new Set((params.events as string[]) ?? []);
		b.ledger.subscriptions = b.ledger.subscriptions.filter((x) => !gone.has(x));
	} else if (
		method === "network.addIntercept" &&
		ok &&
		typeof result.intercept === "string"
	) {
		b.ledger.intercepts.push(result.intercept);
	} else if (method === "network.removeIntercept") {
		const ids = new Set(
			((params.intercepts as string[]) ?? []).concat(
				(params.intercept as string[]) ?? [],
			),
		);
		b.ledger.intercepts = b.ledger.intercepts.filter((x) => !ids.has(x));
	} else if (
		method === "script.addPreloadScript" &&
		ok &&
		typeof result.id === "string"
	) {
		b.ledger.preloadScripts.push(result.id);
	} else if (
		method === "script.removePreloadScript" &&
		typeof params.id === "string"
	) {
		b.ledger.preloadScripts = b.ledger.preloadScripts.filter(
			(x) => x !== params.id,
		);
	}
}

// Serialize each response exactly once: the size check and the payload text
// share one stringify. Oversized results are written to disk and replaced by
// a small marker before serialization. Returns responses plus their JSON
// texts so callers can assemble the payload without re-serializing.
export async function runBatch(
	ws: Pick<WebSocket, "send">,
	b: BatchCtx,
	frames: Frame[],
	byId: Map<number, BidiResponse>,
	timeoutMs: number,
	onResponse?: (response: BidiResponse, completed: number) => void,
	signal?: AbortSignal,
): Promise<{ responses: BidiResponse[]; texts: string[] }> {
	const responses: BidiResponse[] = [];
	const texts: string[] = [];
	const spoolWrites: Array<{
		id: number;
		file: string;
		bytes: number;
		done: Promise<Error | null>;
	}> = [];
	let aborted = false;
	let inlineBytes = 0;
	const record = (frame: Frame, response: BidiResponse): BidiResponse => {
		const index = responses.length;
		const reserve = frames
			.slice(index + 1)
			.reduce(
				(total, remaining, remainingIndex) =>
					total +
					responseCost(
						JSON.stringify(payloadBudgetError(remaining)),
						index + remainingIndex + 1,
					),
				0,
			);
		let out = response;
		let text = JSON.stringify(out);
		if (inlineBytes + responseCost(text, index) + reserve > RESPONSE_BUDGET) {
			out = payloadBudgetError(frame);
			text = JSON.stringify(out);
		}
		inlineBytes += responseCost(text, index);
		responses.push(out);
		texts.push(text);
		onResponse?.(out, responses.length);
		return out;
	};
	for (const frame of frames) {
		if (aborted) {
			record(frame, {
				id: frame.id,
				method: frame.method,
				skipped: true,
			});
			continue;
		}
		let wire: unknown;
		try {
			const paramsText = JSON.stringify(frame.params ?? {});
			if (Buffer.byteLength(paramsText) > MAX_PAYLOAD)
				throw new Error("BiDi params exceed 1 MiB budget");
			const loaded = new Map<number, BidiResponse>();
			for (const match of paramsText.matchAll(
				/(?<!\{)\{\{(\d+)\.result(?:\.\w+)+\}\}(?!\})/g,
			)) {
				const id = Number(match[1]);
				const response = byId.get(id);
				const file = (response?.result as { spooled?: string } | undefined)
					?.spooled;
				if (response && file && !loaded.has(id)) {
					const write = spoolWrites.find((entry) => entry.id === id);
					const failure = await write?.done;
					if (failure)
						throw new Error(
							`spool write failed for frame ${id}: ${failure.message}`,
						);
					loaded.set(id, response);
					byId.set(id, {
						...response,
						result: JSON.parse(await readFile(file, "utf8")),
					});
				}
			}
			try {
				wire = {
					id: frame.id,
					method: frame.method,
					params: interpolate(frame.params ?? {}, byId),
				};
			} finally {
				for (const [id, response] of loaded) byId.set(id, response);
			}
		} catch (e) {
			const r: BidiResponse = {
				id: frame.id,
				method: frame.method,
				type: "error",
				error: "pi interpolation",
				message: e instanceof Error ? e.message : String(e),
			};
			record(frame, r);
			aborted = true;
			continue;
		}
		const res = await new Promise<BidiResponse>((resolve) => {
			let settled = false;
			const finish = (r: BidiResponse) => {
				if (settled) return;
				settled = true;
				clearTimeout(t);
				signal?.removeEventListener("abort", onAbort);
				b.pending.delete(frame.id);
				resolve(r);
			};
			const onAbort = () => finish({ id: frame.id, type: "aborted" });
			const t = setTimeout(
				() =>
					finish({
						id: frame.id,
						type: "timeout",
						message: `no response within ${timeoutMs}ms; side effects may have applied`,
					}),
				timeoutMs,
			);
			signal?.addEventListener("abort", onAbort, { once: true });
			if (signal?.aborted) return onAbort();
			b.pending.set(frame.id, finish);
			try {
				const request = JSON.stringify(wire);
				if (Buffer.byteLength(request) > MAX_PAYLOAD)
					throw new Error("BiDi request exceeds 1 MiB budget");
				ws.send(request);
			} catch (e) {
				finish({ id: frame.id, type: "send-error", message: String(e) });
			}
		});
		if (
			["timeout", "connection-closed", "aborted", "send-error"].includes(
				res.type ?? "",
			)
		)
			aborted = true;
		ledgerUpdate(
			b,
			frame.method,
			(frame.params ?? {}) as Record<string, unknown>,
			res,
		);
		let out: BidiResponse = { ...res, id: frame.id, method: frame.method };
		let text: string;
		// Serialize the result first: it sizes the spool decision and, when
		// oversized, is the only big serialization this frame pays (the old
		// size-probe on the full frame plus a second result pass cost 2x).
		// Small results serialize twice (probe + wrapper); measured free.
		if (res.type === "success" && out.result !== undefined) {
			const resultText = JSON.stringify(out.result);
			const resultBytes = Buffer.byteLength(resultText);
			// wrapper slack keeps borderline frames under MAX_INLINE once marked
			if (
				resultBytes > MAX_INLINE - 1024 ||
				inlineBytes + resultBytes > RESPONSE_BUDGET
			) {
				if (
					(b.spoolBytes ?? 0) + resultBytes > MAX_SPOOL ||
					(b.spoolBudget?.bytes ?? 0) + resultBytes > MAX_SPOOL
				) {
					out.type = "error";
					out.result = undefined;
					out.error = "pi spool budget exhausted";
					text = JSON.stringify(out);
				} else {
					const file = join(
						b.spoolDir,
						`${frame.method.replace(/[^\w.-]/g, "_")}-${frame.id}.json`,
					);
					// exact pre-marker full-frame length without a second big
					// serialization: wrap a placeholder result, substitute its span
					const probe = JSON.stringify({ ...out, result: 0 });
					const fullLength = Buffer.byteLength(probe) - 1 + resultBytes;
					out.result = {
						spooled: file,
						size: fullLength,
						note: "full result written to file; read selectively",
					} as unknown;
					text = JSON.stringify(out);
					// Write off the critical path: frames keep flowing while the file
					// flushes; the drain below completes before we return.
					reserveSpool(b, resultBytes);
					spoolWrites.push({
						id: frame.id,
						file,
						bytes: resultBytes,
						done: writeFile(file, resultText).then(
							() => null,
							(error: unknown) =>
								error instanceof Error ? error : new Error(String(error)),
						),
					});
				}
			} else {
				text = JSON.stringify(out);
			}
		} else {
			text = JSON.stringify(out);
		}
		if (
			Buffer.byteLength(text) > MAX_INLINE ||
			inlineBytes + Buffer.byteLength(text) > RESPONSE_BUDGET
		) {
			out = payloadBudgetError(frame);
			text = JSON.stringify(out);
		}
		byId.set(frame.id, record(frame, out));
	}
	// Every write has a rejection handler from creation; drain before exposing
	// markers, then replace failed markers in both payload and reference cache.
	for (const write of spoolWrites) {
		const error = await write.done;
		if (!error) continue;
		try {
			await rm(write.file, { force: true });
			releaseSpool(b, write.bytes);
		} catch {
			// Keep the reservation if partial output cannot be removed.
		}
		const index = responses.findIndex((r) => r.id === write.id);
		const failed: BidiResponse = {
			id: write.id,
			method: responses[index].method,
			type: "error",
			error: "pi spool write failed",
		};
		const failedText = JSON.stringify(failed);
		inlineBytes +=
			responseCost(failedText, index) - responseCost(texts[index], index);
		responses[index] = failed;
		texts[index] = failedText;
		byId.set(write.id, failed);
	}
	return { responses, texts };
}

// Assemble the tool payload from pre-serialized response texts so large
// results are never serialized twice. Key order matches the historical
// object literal: session, responses, events, note?, wait?.
export function assemblePayload(
	session: unknown,
	responseTexts: string[],
	events: unknown[],
	eventCap: number,
	wait?: unknown,
): string {
	const bytes = (text: string) => Buffer.byteLength(text);
	const bounded = (value: unknown, limit: number) => {
		const text = JSON.stringify(value);
		return bytes(text) <= limit ? text : '{"truncated":true}';
	};
	const responseBytes = responseTexts.reduce(
		(sum, text, index) => sum + responseCost(text, index),
		0,
	);
	if (responseBytes > RESPONSE_BUDGET)
		throw new Error("BiDi responses exceed 896 KiB payload budget");
	const sessionJson = JSON.stringify(session);
	const sessionText =
		bytes(sessionJson) <= 32 * 1024
			? sessionJson
			: bounded(
					{
						sessionId: (session as { sessionId?: unknown })?.sessionId,
						port: (session as { port?: unknown })?.port,
						capabilities: { truncated: true },
					},
					32 * 1024,
				);
	const waitText =
		wait === undefined ? "" : ',"wait":' + bounded(wait, 72 * 1024);
	const prefix = `{"session":${sessionText},"responses":[${responseTexts.join(",")}],"events":[`;
	const note = ',"note":"event buffer capped or payload budget reached"';
	let remaining = MAX_PAYLOAD - bytes(prefix) - bytes(`]${note}${waitText}}`);
	if (remaining < 0)
		throw new Error("BiDi response payload exceeds 1 MiB budget");
	const kept: string[] = [];
	for (let i = events.length - 1; i >= 0; i--) {
		const text = bounded(events[i], 8192);
		const cost = bytes(text) + (kept.length ? 1 : 0);
		if (cost > remaining) break;
		kept.push(text);
		remaining -= cost;
	}
	const capped = kept.length < events.length || events.length === eventCap;
	return `${prefix}${kept.reverse().join(",")}]${capped ? note : ""}${waitText}}`;
}

// Strings in frame params may reference earlier responses: "{{2.result.context}}".
// Resolved after the referenced frame completes, before this frame is sent.
// Strings without the marker skip the regex entirely.
export function interpolate(
	value: unknown,
	byId: Map<number, BidiResponse>,
): unknown {
	if (typeof value === "string") {
		if (!value.includes("{{")) return value;
		const ref =
			/\{\{\{\{(\d+)((?:\.\w+(?:\.\d+)*)+)\}\}\}\}|\{\{(\d+)((?:\.\w+(?:\.\d+)*)+)\}\}/g;
		const resolve = (idStr: string, pathStr: string): unknown => {
			const res = byId.get(Number(idStr));
			if (res?.type !== "success")
				throw new Error(
					`{{${idStr}${pathStr}}}: frame ${idStr} has no successful response`,
				);
			let node: unknown = res;
			for (const seg of pathStr.replace(/^\./, "").split(".")) {
				if (
					node === null ||
					node === undefined ||
					!Object.hasOwn(Object(node), seg)
				)
					throw new Error(`{{${idStr}${pathStr}}}: path ends at ${seg}`);
				node = (node as Record<string, unknown>)[seg];
			}
			return node;
		};
		const whole = /^\{\{(\d+)((?:\.\w+(?:\.\d+)*)+)\}\}$/.exec(value);
		if (whole) return resolve(whole[1], whole[2]);
		// A doubled pair of braces protects page text without rescanning it.
		return value.replace(ref, (_match, escapedId, escapedPath, id, path) => {
			if (escapedId) return `{{${escapedId}${escapedPath}}}`;
			const resolved = resolve(id, path);
			if (resolved === null || typeof resolved === "object")
				throw new Error(
					`{{${id}${path}}}: embedded reference must resolve to a primitive`,
				);
			return String(resolved);
		});
	}
	if (Array.isArray(value)) return value.map((v) => interpolate(v, byId));
	if (value && typeof value === "object") {
		return Object.fromEntries(
			Object.entries(value as Record<string, unknown>).map(([k, v]) => [
				k,
				interpolate(v, byId),
			]),
		);
	}
	return value;
}

// Gecko requires destinationFolder to be an existing directory and silently
// falls back to ~/Downloads otherwise. Resolve per-call param > configured dir
// > Downloads default, and create the directory up front.
export async function resolveScreencastDestinations(
	frames: { method: string; params?: Record<string, unknown> }[],
	configured?: () => string | undefined,
): Promise<void> {
	for (const frame of frames) {
		if (frame.method !== "browsingContext.startScreencast") continue;
		frame.params ??= {};
		const dir =
			(typeof frame.params.destinationFolder === "string" &&
			frame.params.destinationFolder.trim()
				? frame.params.destinationFolder
				: undefined) ?? configured?.();
		if (!dir) continue; // Gecko default: ~/Downloads
		await mkdir(dir, { recursive: true });
		frame.params.destinationFolder = dir;
	}
}

// Cross-call {{N.result.path}} seeds grow without bound over a session;
// keep only the most recently inserted `keep` frame results.
export function trimById(
	byId: Map<number, BidiResponse>,
	keep: number,
): Map<number, BidiResponse> {
	if (
		byId.size <= keep &&
		[...byId.values()].reduce(
			(n, response) => n + Buffer.byteLength(JSON.stringify(response)),
			0,
		) <=
			2 * 1024 * 1024
	)
		return byId;
	const retained = new Map<number, BidiResponse>();
	let bytes = 0;
	for (const [id, response] of [...byId].reverse()) {
		const size = Buffer.byteLength(JSON.stringify(response));
		if (retained.size >= keep || bytes + size > 2 * 1024 * 1024) continue;
		retained.set(id, response);
		bytes += size;
	}
	return new Map([...retained].reverse());
}