pi-ext

import { execFile } from "node:child_process";
import {
	createBashToolDefinition,
	type ExtensionAPI,
	type ExtensionContext,
	type ThemeColor,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { findNu, freshOperations, NuSession } from "./runtime.ts";
import { closeDebug, dbg } from "./src/debug.ts";
import {
	parseBooleanSetting,
	registerSettingFlag,
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";

// Precedence: --nushell, PI_NUSHELL_ENABLE, trusted project, user, false.
const ENABLE_SETTING: SettingDeclaration<boolean> = {
	key: "nushell.enable",
	parse: parseBooleanSetting,
	default: false,
	flag: {
		name: "nushell",
		type: "boolean",
		description:
			"Enable Nushell tools at startup instead of bash and powershell",
	},
	env: "PI_NUSHELL_ENABLE",
};

const NU_COLORS: Record<number, ThemeColor> = {
	32: "syntaxString",
	33: "syntaxOperator",
	34: "syntaxPunctuation",
	35: "syntaxVariable",
	36: "syntaxFunction",
	90: "syntaxVariable",
	96: "syntaxKeyword",
};

// Translate Nu's fixed ANSI palette into the active Pi theme at render time.
function themeNuHighlight(
	ansi: string,
	theme: Parameters<
		NonNullable<ReturnType<typeof createBashToolDefinition>["renderCall"]>
	>[1],
): string {
	let output = "";
	let start = 0;
	let color: ThemeColor = "toolOutput";
	for (const match of ansi.matchAll(/\x1b\[([0-9;]*)m/g)) {
		output += theme.fg(color, ansi.slice(start, match.index));
		const codes = match[1]?.split(";").map(Number) ?? [];
		const foreground = codes.findLast(
			(code) => (code >= 30 && code <= 37) || (code >= 90 && code <= 97),
		);
		color =
			foreground === 35 && codes.includes(1)
				? "syntaxPunctuation"
				: (NU_COLORS[foreground ?? -1] ?? "toolOutput");
		start = match.index + match[0].length;
	}
	return output + theme.fg(color, ansi.slice(start));
}

// Nu parses source without evaluating it; only the command heading is highlighted.
const highlightCache = new WeakMap<object, { command: string; text: string }>();
function renderNuCall(
	args: { command: string; timeout?: number },
	theme: Parameters<
		NonNullable<ReturnType<typeof createBashToolDefinition>["renderCall"]>
	>[1],
	context: Parameters<
		NonNullable<ReturnType<typeof createBashToolDefinition>["renderCall"]>
	>[2],
	label: string,
) {
	const state = context.state;
	if (context.executionStarted && state.startedAt === undefined) {
		state.startedAt = Date.now();
		state.endedAt = undefined;
	}
	const command = typeof args?.command === "string" ? args.command : "";
	let cached = highlightCache.get(state);
	if (context.argsComplete && cached?.command !== command) {
		const entry = { command, text: command };
		cached = entry;
		highlightCache.set(state, entry);
		const binary = command && command.length <= 16_000 ? findNu() : undefined;
		if (binary) {
			execFile(
				binary,
				[
					"--no-config-file",
					"--no-history",
					"-c",
					"$env.PI_NU_HIGHLIGHT_INPUT | nu-highlight",
				],
				{
					env: { ...process.env, PI_NU_HIGHLIGHT_INPUT: command },
					encoding: "utf8",
					timeout: 1000,
					maxBuffer: 128 * 1024,
					windowsHide: true,
				},
				(error, stdout) => {
					if (error || highlightCache.get(state) !== entry) return;
					entry.text = stdout.trimEnd();
					context.invalidate();
				},
			);
		}
	}
	const highlighted = cached?.command === command ? cached.text : command;
	const text =
		context.lastComponent instanceof Text
			? context.lastComponent
			: new Text("", 0, 0);
	text.setText(
		`${theme.fg("toolTitle", theme.bold(label))} ${themeNuHighlight(highlighted || "...", theme)}${args?.timeout ? theme.fg("muted", ` (timeout ${args.timeout}s)`) : ""}`,
	);
	return text;
}

export default function nushellExtension(pi: ExtensionAPI) {
	const session = new NuSession();
	let enabled = false;
	let displaced: string[] = [];

	registerSettingFlag(pi, ENABLE_SETTING);

	function activate(ctx: ExtensionContext, quiet = false): void {
		if (enabled) return;
		if (!findNu()) {
			ctx.ui.notify(
				"Nushell executable not found on PATH; tools unchanged",
				"warning",
			);
			dbg?.("tools.enable.error", { kind: "nu_not_found" });
			return;
		}
		const available = new Set(pi.getAllTools().map((tool) => tool.name));
		const missing = ["nu", "nu_session"].filter((name) => !available.has(name));
		if (missing.length) {
			ctx.ui.notify(
				`Nushell tools excluded by tool allowlist: ${missing.join(", ")}; tools unchanged`,
				"warning",
			);
			dbg?.("tools.enable.error", {
				kind: "tools_excluded",
				count: missing.length,
			});
			return;
		}
		const active = pi.getActiveTools();
		displaced = active.filter(
			(tool) => tool === "bash" || tool === "powershell",
		);
		pi.setActiveTools([
			...active.filter((tool) => !displaced.includes(tool)),
			"nu",
			"nu_session",
		]);
		enabled = true;
		dbg?.("tools.enabled", { count: displaced.length });
		if (!quiet)
			ctx.ui.notify(
				"Nushell tools on: nu (fresh), nu_session (persistent)",
				"info",
			);
	}

	const freshTool = createBashToolDefinition(process.cwd(), {
		operations: {
			exec: (command, cwd, options) => {
				const binary = findNu();
				if (!binary) throw new Error("Nushell executable not found on PATH");
				return freshOperations(binary).exec(command, cwd, options);
			},
		},
	});
	pi.registerTool({
		...freshTool,
		name: "nu",
		label: " nu",
		description:
			"Execute Nushell commands in a fresh process in Pi's cwd. Syntax is Nushell, not Bash or PowerShell. Each call has independent variables, functions, environment, and cwd. Streams output; timeout in seconds is optional.",
		promptSnippet:
			"Execute independent Nushell commands with nu; use nu_session only for state across calls",
		promptGuidelines: [
			"nu already runs Nushell. Pass source directly; never wrap it in nu -c or ^nu -c. Prefer native commands like ls, pwd, open, mkdir, save; use ^name only for necessary external executables. Nu flags differ from GNU flags: check help <command> before guessing options. Batch short independent checks when readable.",
			"Use Nushell syntax: structured pipelines, $env.NAME, (command), o> file or save. Use try/catch for expected command failures.",
		],
		renderCall: (args, theme, context) =>
			renderNuCall(args, theme, context, " nu>"),
	});

	pi.registerTool({
		name: "nu_session",
		label: " nu session",
		description:
			"Evaluate Nushell in a persistent session. Variables, functions, environment, and cwd survive subsequent nu_session calls; nu calls stay independent. Prefer nu unless state must persist. MCP captures expression results as NUON; use ^command | complete and inspect exit_code for external status. Do not rely on print: MCP uses stderr for diagnostics.",
		promptSnippet:
			"Evaluate persistent Nushell state with nu_session when later calls depend on earlier state",
		promptGuidelines: [
			"Use nu by default; choose nu_session only for state shared across calls. Both tools already run Nu, so do not wrap commands in nu -c or ^nu -c. Nu session is sequential; /nushell reset clears its state. Cancelled commands can leave detached external processes running.",
		],
		executionMode: "sequential",
		parameters: Type.Object({
			command: Type.String({
				description: "Nushell source to evaluate in the persistent session",
			}),
		}),
		async execute(_id, { command }, signal, _onUpdate, ctx) {
			const binary = findNu();
			if (!binary) throw new Error("Nushell executable not found on PATH");
			const text = await session.evaluate(binary, ctx.cwd, command, signal);
			return { content: [{ type: "text", text }], details: undefined };
		},
		renderCall: (args, theme, context) =>
			renderNuCall(args, theme, context, " nu session>"),
		renderResult: freshTool.renderResult,
	});

	pi.registerCommand("nushell", {
		description:
			"Toggle Nushell tools and replace active bash/powershell; /nushell reset clears persistent state",
		getArgumentCompletions: (prefix) => {
			const matches = ["on", "off", "reset"].filter((value) =>
				value.startsWith(prefix),
			);
			return matches.length
				? matches.map((value) => ({ value, label: value }))
				: null;
		},
		async handler(arg, ctx) {
			const action = arg.trim() || (enabled ? "off" : "on");
			if (action === "reset") {
				await session.reset("explicit");
				ctx.ui.notify("Nushell session reset", "info");
				return;
			}
			if (action === "on") {
				activate(ctx);
				return;
			}
			if (action === "off") {
				if (!enabled) return;
				await session.reset("toggle_off");
				const active = pi
					.getActiveTools()
					.filter((tool) => tool !== "nu" && tool !== "nu_session");
				pi.setActiveTools([
					...active,
					...displaced.filter((tool) => !active.includes(tool)),
				]);
				dbg?.("tools.disabled", { count: displaced.length });
				displaced = [];
				enabled = false;
				ctx.ui.notify("Nushell tools off; persistent state cleared", "info");
				return;
			}
			ctx.ui.notify("Usage: /nushell [on|off|reset]", "warning");
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		await session.reset("session_start");
		dbg?.("session.start");
		enabled = false;
		displaced = [];
		pi.setActiveTools(
			pi
				.getActiveTools()
				.filter((name) => name !== "nu" && name !== "nu_session"),
		);
		const setting = resolveSetting(pi, ctx, ENABLE_SETTING);
		if (!setting.ok) {
			ctx.ui.notify(`${setting.error}; tools unchanged`, "warning");
			dbg?.("settings.invalid", { source: setting.source });
		} else if (setting.value) activate(ctx, true);
	});
	pi.on("session_tree", async () => {
		await session.reset("session_tree");
	});
	pi.on("session_shutdown", async () => {
		await session.reset("session_shutdown");
		dbg?.("session.shutdown");
		closeDebug();
	});
}