pi-ext

import { delimiter, join } from "node:path";
import {
	type BashSpawnContext,
	createBashTool,
	type ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";
import { findExecutable } from "./src/pi-ext-executable.ts";

const INTERCEPTED_COMMANDS_PATH = join(
	import.meta.dirname,
	"intercepted-commands",
);
const PIP_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?pip\s*(?:$|\s)/m;
const PIP3_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?pip3\s*(?:$|\s)/m;
const POETRY_COMMAND = /(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?poetry\s*(?:$|\s)/m;
const PYTHON_COMMAND = String.raw`(?:^|\n|[;|&]{1,2})\s*(?:\S+\/)?python(?:3(?:\.\d+)?)?\b[^\n;|&]*`;
const PYTHON_PIP_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*pip\\b|\\s-mpip\\b)`,
	"m",
);
const PYTHON_VENV_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*venv\\b|\\s-mvenv\\b)`,
	"m",
);
const PYTHON_COMPILE_COMMAND = new RegExp(
	`${PYTHON_COMMAND}(?:\\s-m\\s*py_compile\\b|\\s-mpy_compile\\b)`,
	"m",
);

async function uvInstalled(
	env = process.env,
	platform = process.platform,
): Promise<boolean> {
	return Boolean(
		findExecutable(
			platform === "win32" ? ["uv.exe", "uv.cmd", "uv.bat"] : ["uv"],
			{ path: env[pathEnvKey(env, platform)] ?? "", platform },
		),
	);
}

function pathEnvKey(
	env: NodeJS.ProcessEnv,
	platform = process.platform,
): string {
	if (platform !== "win32") return "PATH";
	return Object.keys(env).find((key) => key.toLowerCase() === "path") ?? "PATH";
}

function prependInterceptedCommands(
	ctx: BashSpawnContext,
	platform = process.platform,
): BashSpawnContext {
	if (platform === "win32") return ctx;
	const key = pathEnvKey(ctx.env);
	const path = [INTERCEPTED_COMMANDS_PATH, ctx.env[key]]
		.filter(Boolean)
		.join(delimiter);
	return { ...ctx, env: { ...ctx.env, [key]: path } };
}

function blocked(message: string, ...body: string[]): string {
	return [
		`Error: ${message}. Use uv instead:`,
		"",
		...body.map((line) => `  ${line}`),
		"",
	].join("\n");
}

const PIP_BLOCKED = blocked(
	"pip is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const PIP3_BLOCKED = blocked(
	"pip3 is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const POETRY_BLOCKED = blocked(
	"poetry is disabled",
	"To initialize a project: uv init",
	"To add a dependency: uv add PACKAGE",
	"To sync dependencies: uv sync",
	"To run commands: uv run COMMAND",
);
const PYTHON_PIP_BLOCKED = blocked(
	"'python -m pip' is disabled",
	"To install a package for a script: uv run --with PACKAGE python script.py",
	"To add a dependency to the project: uv add PACKAGE",
);
const PYTHON_VENV_BLOCKED = blocked(
	"'python -m venv' is disabled",
	"To create a virtual environment: uv venv",
);
const PYTHON_COMPILE_BLOCKED = [
	"Error: 'python -m py_compile' is disabled because it writes .pyc files to __pycache__.",
	"",
	"  To verify syntax without bytecode output: uv run python -m ast path/to/file.py >/dev/null",
	"",
].join("\n");

function getBlockedCommandMessage(command: string): string | null {
	if (PIP_COMMAND.test(command)) return PIP_BLOCKED;
	if (PIP3_COMMAND.test(command)) return PIP3_BLOCKED;
	if (POETRY_COMMAND.test(command)) return POETRY_BLOCKED;
	if (PYTHON_PIP_COMMAND.test(command)) return PYTHON_PIP_BLOCKED;
	if (PYTHON_VENV_COMMAND.test(command)) return PYTHON_VENV_BLOCKED;
	if (PYTHON_COMPILE_COMMAND.test(command)) return PYTHON_COMPILE_BLOCKED;
	return null;
}

export const __test = {
	getBlockedCommandMessage,
	prependInterceptedCommands,
	uvInstalled,
};

export default function uvExtension(pi: ExtensionAPI) {
	let generation = 0;

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start");
		const ticket = ++generation;
		const finish = span?.("availability.check");
		void uvInstalled().then(
			(installed) => {
				finish?.("finish", { installed });
				if (ticket === generation && !installed && ctx.hasUI)
					ctx.ui.notify("uv extension: uv not found on PATH", "warning");
			},
			(error) => {
				finish?.("error");
				throw error;
			},
		);
	});
	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
		generation++;
	});

	const bash = createBashTool(process.cwd(), {
		spawnHook: (ctx) => {
			const message = getBlockedCommandMessage(ctx.command);
			if (message) throw new Error(message);
			return prependInterceptedCommands(ctx);
		},
	});
	const promptMetadata = bash as typeof bash & {
		promptGuidelines?: readonly string[];
	};
	Object.assign(promptMetadata, {
		description: `${bash.description} This replacement routes PATH-based Python through uv and rejects pip, Poetry, python -m pip, python -m venv, and python -m py_compile.`,
		promptGuidelines: [
			...(promptMetadata.promptGuidelines ?? []).map((guideline) =>
				guideline ===
				"You can inspect PI_* environment variables for current model and session details."
					? "Use bash to inspect PI_* environment variables for current model and session details."
					: guideline,
			),
			"Use bash with uv instead of pip, pip3, Poetry, python -m pip, or python -m venv for Python package and environment operations; use uv run python -m ast PATH >/dev/null instead of python -m py_compile for syntax checks.",
		],
	});
	pi.registerTool(bash);
}