pi-ext

import { AsyncLocalStorage } from "node:async_hooks";
import { execFile } from "node:child_process";
import {
	appendFile,
	chmod,
	copyFile,
	mkdir,
	mkdtemp,
	rm,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import path, { delimiter } from "node:path";
import { promisify } from "node:util";
import {
	type Context,
	fauxAssistantMessage,
	fauxProvider,
} from "@earendil-works/pi-ai";
import { stripTerminalSequences } from "@earendil-works/pi-tui";
import { createTestSession } from "../../../test/harness";
import { cleanGitEnvironment } from "../git.js";
import type { Hunk, Plan } from "../types.js";

const nativeScope = new AsyncLocalStorage<{
	signal: AbortSignal;
	pending: Set<Promise<unknown>>;
	cancelling: Promise<unknown>[];
	cleanup: AsyncDisposableStack;
}>();

// Used by Vitest aroundEach; no runner import, so Playwright can use this file too.
export async function nativeTest(
	run: () => Promise<void>,
	context: { signal: AbortSignal },
) {
	const controller = new AbortController();
	const scope = {
		signal: AbortSignal.any([context.signal, controller.signal]),
		pending: new Set<Promise<unknown>>(),
		cancelling: [] as Promise<unknown>[],
		cleanup: new AsyncDisposableStack(),
	};
	await nativeScope.run(scope, async () => {
		try {
			await run();
		} finally {
			controller.abort();
			while (scope.pending.size) await Promise.allSettled([...scope.pending]);
			await scope.cleanup.disposeAsync();
		}
	});
}

export function testSignal(signal?: AbortSignal): AbortSignal {
	const parent = nativeScope.getStore()?.signal;
	return parent && signal
		? AbortSignal.any([parent, signal])
		: (parent ?? signal ?? new AbortController().signal);
}

export { own as nativeOperation };

function own<T>(work: () => Promise<T>): Promise<T> {
	const scope = nativeScope.getStore();
	let pending: Promise<T>;
	try {
		pending = work();
	} catch (error) {
		pending = Promise.reject(error);
	}
	scope?.pending.add(pending);
	void pending.then(
		() => scope?.pending.delete(pending),
		() => scope?.pending.delete(pending),
	);
	return pending;
}

export function cleanupNative(cleanup: () => Promise<void>): void {
	nativeScope.getStore()?.cleanup.defer(cleanup);
}

const exec = promisify(execFile);
export function git(cwd: string, ...args: string[]) {
	return own(async () => {
		const signal = testSignal();
		signal.throwIfAborted();
		const command = exec("git", args, {
			cwd,
			signal,
			env: cleanGitEnvironment(),
		});
		const closed = new Promise<void>((resolve) =>
			command.child.once("close", () => resolve()),
		);
		try {
			return (await command).stdout.trim();
		} finally {
			await closed;
		}
	});
}

export function withFakeGh(
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	return own(() => fakeGh(program, check));
}

async function fakeGh(
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	testSignal().throwIfAborted();
	const root = await mkdtemp(path.join(tmpdir(), "strata-gh-"));
	try {
		const bin = path.join(root, "bin");
		await mkdir(bin);
		const executable = path.join(
			bin,
			process.platform === "win32" ? "gh.exe" : "gh",
		);
		await copyFile(process.execPath, executable);
		await chmod(executable, 0o755);
		const preload = path.join(root, "fake-gh.cjs");
		await writeFile(
			preload,
			`if (/^gh(?:\\.exe)?$/i.test(require("node:path").basename(process.execPath)) && process.env.STRATA_FAKE_GH_PRELOAD === __filename) {\nconst original = process.argv.slice(1); original[0] = require("node:path").basename(original[0]); process.argv = [process.execPath, "fake-gh", ...(process.execArgv.includes("-c") ? ["-c"] : []), ...original]; const args = process.argv.slice(2);\n${program}\nprocess.exit(process.exitCode ?? 0);\n}`,
		);
		const previousPath = process.env.PATH;
		const previousOptions = process.env.NODE_OPTIONS;
		const previousPreload = process.env.STRATA_FAKE_GH_PRELOAD;
		process.env.PATH = `${bin}${delimiter}${previousPath ?? ""}`;
		process.env.STRATA_FAKE_GH_PRELOAD = preload;
		process.env.NODE_OPTIONS = [
			previousOptions,
			`--require=${JSON.stringify(preload)}`,
		]
			.filter(Boolean)
			.join(" ");
		try {
			await check();
		} finally {
			await Promise.allSettled(nativeScope.getStore()?.cancelling ?? []);
			if (previousPath === undefined) delete process.env.PATH;
			else process.env.PATH = previousPath;
			if (previousOptions === undefined) delete process.env.NODE_OPTIONS;
			else process.env.NODE_OPTIONS = previousOptions;
			if (previousPreload === undefined)
				delete process.env.STRATA_FAKE_GH_PRELOAD;
			else process.env.STRATA_FAKE_GH_PRELOAD = previousPreload;
		}
	} finally {
		await rm(root, { recursive: true, force: true });
	}
}

export function withFakeExecutable(
	name: string,
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	return own(() => fakeExecutable(name, program, check));
}

async function fakeExecutable(
	name: string,
	program: string,
	check: () => Promise<void>,
): Promise<void> {
	testSignal().throwIfAborted();
	const root = await mkdtemp(path.join(tmpdir(), "strata-executable-"));
	try {
		const bin = path.join(root, "bin");
		await mkdir(bin);
		const executable = path.join(
			bin,
			process.platform === "win32" ? `${name}.exe` : name,
		);
		await copyFile(process.execPath, executable);
		await chmod(executable, 0o755);
		const preload = path.join(root, "fake-executable.cjs");
		await writeFile(
			preload,
			`const executable = require("node:path").basename(process.execPath).toLowerCase();
if (executable === ${JSON.stringify(name.toLowerCase())} || executable === ${JSON.stringify(`${name.toLowerCase()}.exe`)}) {
const original = process.argv.slice(1); original[0] = require("node:path").basename(original[0]); process.argv = [process.execPath, ${JSON.stringify(name)}, ...(process.execArgv.includes("-c") ? ["-c"] : []), ...original]; const args = process.argv.slice(2);
${program}
process.exit(process.exitCode ?? 0);
}`,
		);
		const previousPath = process.env.PATH;
		const previousOptions = process.env.NODE_OPTIONS;
		process.env.PATH = `${bin}${delimiter}${previousPath ?? ""}`;
		process.env.NODE_OPTIONS = [
			previousOptions,
			`--require=${JSON.stringify(preload)}`,
		]
			.filter(Boolean)
			.join(" ");
		try {
			await check();
		} finally {
			await Promise.allSettled(nativeScope.getStore()?.cancelling ?? []);
			if (previousPath === undefined) delete process.env.PATH;
			else process.env.PATH = previousPath;
			if (previousOptions === undefined) delete process.env.NODE_OPTIONS;
			else process.env.NODE_OPTIONS = previousOptions;
		}
	} finally {
		await rm(root, { recursive: true, force: true });
	}
}

export function fixturePlan(hunks: Array<Pick<Hunk, "id" | "path">>): Plan {
	return {
		summary: "Bound retries before using them.",
		cohorts: [
			{
				title: "Payment retries",
				layers: hunks.map((hunk, index) => ({
					id: `layer-${index + 1}`,
					title: `Review ${hunk.path}`,
					summary: `Layer ${index + 1} explains the original captured change.`,
					hunks: [{ id: hunk.id, summary: `Explanation for ${hunk.path}.` }],
					flow: ["Contract", "Consumer", "Tests"],
				})),
			},
		],
	};
}

export function openFixture(mode: "tui" | "rpc" | "json" = "tui") {
	return own(() => createFixture(mode));
}

async function createFixture(mode: "tui" | "rpc" | "json") {
	const signal = testSignal();
	signal.throwIfAborted();
	const root = await mkdtemp(path.join(tmpdir(), "strata-pi-"));
	let t: Awaited<ReturnType<typeof createTestSession>> | undefined;
	let closing: Promise<void> | undefined;
	const close = () =>
		(closing ??= (async () => {
			try {
				// Harness dispose does not emit shutdown or await extension work.
				await t?.session.extensionRunner.emit({
					type: "session_shutdown",
					reason: "quit",
				});
			} finally {
				try {
					t?.dispose();
				} finally {
					await rm(root, { recursive: true, force: true });
				}
			}
		})());
	cleanupNative(close);
	try {
		await git(root, "init", "--quiet");
		await appendFile(
			path.join(root, ".git", "config"),
			"[user]\n\tname = Strata Test\n\temail = strata@example.invalid\n",
		);
		await writeFile(
			path.join(root, "contract.ts"),
			"export const attempts = 1;\n",
		);
		await writeFile(
			path.join(root, "consumer.ts"),
			"export const retry = false;\n",
		);
		await git(root, "add", ".");
		await git(root, "commit", "--quiet", "-m", "base");
		await writeFile(
			path.join(root, "contract.ts"),
			"export const attempts = 3;\n",
		);
		await writeFile(
			path.join(root, "consumer.ts"),
			"export const retry = true;\n",
		);
		t = await createTestSession({
			cwd: root,
			env: { STRATA_NO_OPEN: "1" },
			extensions: [path.resolve(import.meta.dirname, "../index.ts")],
		});
		const session = t;
		signal.throwIfAborted();
		const runner = t.session.extensionRunner;
		const emit = runner.emit.bind(runner);
		runner.emit = (...args: Parameters<typeof emit>) =>
			own(() => emit(...args));
		const base = runner.createContext();
		const faux = fauxProvider({
			provider: "strata-integration",
			models: [{ id: "review", maxTokens: 16_384 }],
		});
		let calls = 0;
		const prompts: string[] = [];
		const respond = (context: Context) => {
			calls++;
			const message = context.messages.findLast((item) => item.role === "user");
			const text =
				typeof message?.content === "string"
					? message.content
					: (message?.content ?? [])
							.filter((part) => part.type === "text")
							.map((part) => part.text)
							.join("");
			prompts.push(text);
			if (text.startsWith("Organize")) {
				const hunks = [
					...text.matchAll(
						/<file path="(contract\.ts|consumer\.ts)">\s*<hunk id="([^"]+)"/g,
					),
				].map(([, filePath, id]) => ({ id, path: filePath }));
				return fauxAssistantMessage(JSON.stringify(fixturePlan(hunks)));
			}
			return fauxAssistantMessage("Three attempts include the first call.");
		};
		faux.setResponses(Array.from({ length: 20 }, () => respond));
		base.modelRegistry.registerProvider(faux.provider);
		const model = base.modelRegistry.find("strata-integration", "review");
		if (!model) throw new Error("Fixture model missing");
		await t.session.setModel(model);
		const notifications: Array<{ text: string; severity: string }> = [];
		const widgets: Array<{
			key: string;
			content: unknown;
			options: unknown;
		}> = [];
		const ctx = Object.create(base);
		Object.defineProperties(ctx, {
			model: { value: model },
			hasUI: { value: mode !== "json" },
			waitForIdle: { value: async () => {} },
			mode: { value: mode },
			ui: {
				value: {
					...base.ui,
					notify: (text: string, severity = "info") =>
						notifications.push({ text, severity }),
					setWidget: (key: string, content: unknown, options: unknown) => {
						widgets.push({ key, content, options });
					},
				},
			},
		});
		const command = runner.extensions[0].commands.get("strata");
		const scope = nativeScope.getStore();
		const cancel = () => {
			const stopped = own(() => command.handler("cancel", ctx));
			scope?.cancelling.push(stopped);
		};
		signal.throwIfAborted();
		signal.addEventListener("abort", cancel, { once: true });
		cleanupNative(async () => {
			signal.removeEventListener("abort", cancel);
		});
		const run = (args = "") =>
			own(async () => {
				signal.throwIfAborted();
				await command.handler(args, ctx);
				signal.throwIfAborted();
			});
		return {
			t,
			root,
			ctx,
			runner,
			command,
			faux,
			notifications,
			widgets,
			prompts,
			run,
			widgetLines(width = 120): string[] {
				const content = widgets.findLast(
					(item) => item.key === "strata",
				)?.content;
				if (content === undefined) return [];
				if (Array.isArray(content)) return content as string[];
				if (typeof content !== "function") return [];
				const theme = {
					bold: (text: string) => text,
					fg: (_color: string, text: string) => text,
				};
				return content({}, theme).render(width);
			},
			widgetText(width = 120): string {
				return this.widgetLines(width).map(stripTerminalSequences).join("\n");
			},
			get calls() {
				return calls;
			},
			get url(): string {
				const line = this.widgetLines(160)[1] ?? "";
				const match = /\x1b]8;;([^\x1b]+)\x1b\\/.exec(line);
				if (match?.[1]) return match[1];
				const notice = notifications.findLast((item) =>
					item.text.startsWith("Strata review: "),
				);
				if (notice) return notice.text.slice("Strata review: ".length);
				const message = session.events.messages.findLast(
					(item) =>
						item.role === "custom" &&
						item.customType === "strata-status" &&
						typeof item.content === "string" &&
						item.content.startsWith("Strata review: "),
				);
				if (message?.role === "custom" && typeof message.content === "string")
					return message.content.slice("Strata review: ".length);
				throw new Error(`No review URL: ${line}`);
			},
			close,
		};
	} catch (error) {
		await close();
		throw error;
	}
}