pi-ext

import { type ChildProcessWithoutNullStreams, spawn } from "node:child_process";
import { randomBytes } from "node:crypto";
import { constants } from "node:fs";
import { access, mkdtemp, rm, stat, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
	type BashOperations,
	type ExtensionAPI,
	truncateTail,
} from "@earendil-works/pi-coding-agent";
import type {
	AutocompleteItem,
	AutocompleteProvider,
} from "@earendil-works/pi-tui";
import { closeDebug, dbg, span } from "./src/debug.ts";
import {
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";

interface ProcessResult {
	stdout: string;
	stderr: string;
	output: string;
	code: number | undefined;
	signal: NodeJS.Signals | null;
}

interface BangCommand {
	command: string;
	prefixLength: number;
}

interface NuExecutionResult {
	output: string;
	exitCode: number | undefined;
	cancelled: boolean;
	truncated: boolean;
	fullOutputPath?: string;
}

const DEFAULT_NU_COMMAND = "nu";
const COMMAND_SETTING: SettingDeclaration<string> = {
	key: "nu-bang.command",
	env: "PI_NU_BANG_COMMAND",
	default: DEFAULT_NU_COMMAND,
	parse: (raw) => (typeof raw === "string" && raw.trim() ? raw : undefined),
};
// Non-login Nushell avoids sourcing user login startup for every bang command.
// That keeps execution deterministic and removes the dominant per-test/process cost.
const NU_BASE_ARGS = ["--no-history", "--error-style", "fancy"] as const;
const NU_TABLE_MODE = "none";
const TEMP_PREFIX = "pi-nu-bang-";
const TEMP_FILE_NAME = "command.nu";
const FULL_OUTPUT_PREFIX = "pi-nu-bang-output-";
const EXIT_STDIO_GRACE_MS = 100;

const tempDirs = new Set<string>();

function nuBaseArgs(): string[] {
	return [...NU_BASE_ARGS];
}

function nuExecArgs(command: string): string[] {
	return [...nuBaseArgs(), "--table-mode", NU_TABLE_MODE, "-c", command];
}

function fullOutputPath(): string {
	return join(
		tmpdir(),
		`${FULL_OUTPUT_PREFIX}${randomBytes(8).toString("hex")}.log`,
	);
}

async function persistFullOutput(output: string): Promise<string> {
	const path = fullOutputPath();
	await writeFile(path, output, { encoding: "utf8", mode: 0o600 });
	return path;
}

function escapeControlChars(text: string): string {
	return text.replace(
		/[\x00-\x1F\x7F-\x9F]/g,
		(char) => `\\u${char.charCodeAt(0).toString(16).padStart(4, "0")}`,
	);
}

function stripTerminalControls(output: string): string {
	return output
		.replace(/\x1B\][\s\S]*?(?:\x07|\x1B\\|\u009C)/g, "")
		.replace(/\u009D[\s\S]*?(?:\x07|\x1B\\|\u009C)/g, "")
		.replace(/\x1B[PX^_][\s\S]*?(?:\x1B\\|\u009C)/g, "")
		.replace(/[\u0090\u0098\u009E\u009F][\s\S]*?(?:\x1B\\|\u009C)/g, "")
		.replace(/\x1B[PX^_][^\r\n]*(?=\r?\n|$)/g, "")
		.replace(/\x1B\][^\r\n]*(?=\r?\n|$)/g, "")
		.replace(/[\u0090\u0098\u009D\u009E\u009F][^\r\n]*(?=\r?\n|$)/g, "")
		.replace(/\x1B\[[0-?]*[ -/]*[@-~]/g, "")
		.replace(/\u009B[0-?]*[ -/]*[@-~]/g, "")
		.replace(/\x1B[@-Z\\-_]/g, "")
		.replace(/[\x00-\x08\x0B-\x0D\x0E-\x1F\x7F-\x9F]/g, "");
}

function truncateOutput(output: string): {
	output: string;
	truncated: boolean;
} {
	const truncated = truncateTail(output);
	if (!truncated.truncated)
		return { output: truncated.content, truncated: false };

	const reason =
		truncated.truncatedBy === "lines"
			? `${truncated.outputLines}/${truncated.totalLines} lines`
			: `${truncated.outputBytes}/${truncated.totalBytes} bytes`;
	return {
		output: `${truncated.content}\n… (truncated, showing tail: ${reason})`,
		truncated: true,
	};
}

function nodeErrorCode(error: unknown): string | undefined {
	return error instanceof Error &&
		"code" in error &&
		typeof error.code === "string"
		? error.code
		: undefined;
}

async function workingDirectoryError(cwd: string): Promise<string | undefined> {
	const safeCwd = escapeControlChars(cwd);
	try {
		const info = await stat(cwd);
		if (!info.isDirectory())
			return `Working directory is not a directory: ${safeCwd}`;
		try {
			await access(cwd, constants.X_OK);
			return undefined;
		} catch (error) {
			const code = nodeErrorCode(error);
			if (code === "EACCES") {
				return `Working directory is not accessible: ${safeCwd}`;
			}
			if (code)
				return `Working directory cannot be checked: ${safeCwd} (${escapeControlChars(code)})`;
			throw error;
		}
	} catch (error) {
		const code = nodeErrorCode(error);
		if (code === "ENOENT") {
			return `Working directory does not exist: ${safeCwd}`;
		}
		if (code)
			return `Working directory cannot be checked: ${safeCwd} (${escapeControlChars(code)})`;
		throw error;
	}
}

function waitForForegroundProcess(
	child: ChildProcessWithoutNullStreams,
): Promise<{ code: number | null; signal: NodeJS.Signals | null }> {
	return new Promise((resolve, reject) => {
		let settled = false;
		let exited = false;
		let exitCode: number | null = null;
		let exitSignal: NodeJS.Signals | null = null;
		let postExitTimer: ReturnType<typeof setTimeout> | undefined;
		let stdoutEnded = false;
		let stderrEnded = false;

		function cleanup(): void {
			if (postExitTimer) clearTimeout(postExitTimer);
			child.removeListener("error", onError);
			child.removeListener("exit", onExit);
			child.removeListener("close", onClose);
			child.stdout.removeListener("end", onStdoutEnd);
			child.stderr.removeListener("end", onStderrEnd);
		}

		function finalize(
			code: number | null,
			processSignal: NodeJS.Signals | null,
		): void {
			if (settled) return;
			settled = true;
			cleanup();
			child.stdout.destroy();
			child.stderr.destroy();
			resolve({ code, signal: processSignal });
		}

		function maybeFinalizeAfterExit(): void {
			if (exited && stdoutEnded && stderrEnded) finalize(exitCode, exitSignal);
		}

		function onStdoutEnd(): void {
			stdoutEnded = true;
			maybeFinalizeAfterExit();
		}

		function onStderrEnd(): void {
			stderrEnded = true;
			maybeFinalizeAfterExit();
		}

		function onError(error: Error): void {
			if (settled) return;
			settled = true;
			cleanup();
			reject(error);
		}

		function onExit(
			code: number | null,
			processSignal: NodeJS.Signals | null,
		): void {
			exited = true;
			exitCode = code;
			exitSignal = processSignal;
			maybeFinalizeAfterExit();
			if (!settled)
				postExitTimer ??= setTimeout(
					() => finalize(code, processSignal),
					EXIT_STDIO_GRACE_MS,
				);
		}

		function onClose(
			code: number | null,
			processSignal: NodeJS.Signals | null,
		): void {
			finalize(code, processSignal);
		}

		child.stdout.once("end", onStdoutEnd);
		child.stderr.once("end", onStderrEnd);
		child.once("error", onError);
		child.once("exit", onExit);
		child.once("close", onClose);
	});
}

function runProcess(
	nuCommand: string,
	args: string[],
	cwd: string,
	signal?: AbortSignal,
): Promise<ProcessResult> {
	return new Promise((resolve, reject) => {
		let stdout = "";
		let stderr = "";
		let output = "";

		runStreamingProcess(nuCommand, args, cwd, {
			onStdout(chunk) {
				const text = chunk.toString("utf8");
				stdout += text;
				output += text;
			},
			onStderr(chunk) {
				const text = chunk.toString("utf8");
				stderr += text;
				output += text;
			},
			signal,
		})
			.then((result) =>
				resolve({
					stdout,
					stderr,
					output,
					code: result.exitCode ?? undefined,
					signal: result.signal,
				}),
			)
			.catch(reject);
	});
}

function runStreamingProcess(
	nuCommand: string,
	args: string[],
	cwd: string,
	options: {
		onStdout?: (chunk: Buffer) => void;
		onStderr?: (chunk: Buffer) => void;
		signal?: AbortSignal;
		timeout?: number;
		env?: NodeJS.ProcessEnv;
	},
): Promise<{ exitCode: number | null; signal: NodeJS.Signals | null }> {
	return new Promise((resolve, reject) => {
		let aborted = options.signal?.aborted ?? false;
		let timedOut = false;
		let forceKillTimer: ReturnType<typeof setTimeout> | undefined;
		let timeoutTimer: ReturnType<typeof setTimeout> | undefined;
		const detached = process.platform !== "win32";
		const child = spawn(nuCommand, args, {
			cwd,
			detached,
			env: options.env ?? process.env,
			stdio: ["ignore", "pipe", "pipe"],
		});

		function killProcessGroup(kind: NodeJS.Signals): void {
			if (child.pid === undefined) return;
			try {
				if (detached) process.kill(-child.pid, kind);
				else child.kill(kind);
			} catch {
				try {
					child.kill(kind);
				} catch {}
			}
		}

		function abortProcess(): void {
			aborted = true;
			killProcessGroup("SIGTERM");
			forceKillTimer ??= setTimeout(() => killProcessGroup("SIGKILL"), 250);
		}

		function cleanup(): void {
			if (forceKillTimer) clearTimeout(forceKillTimer);
			if (timeoutTimer) clearTimeout(timeoutTimer);
			options.signal?.removeEventListener("abort", abortProcess);
		}

		if (options.timeout !== undefined && options.timeout > 0) {
			timeoutTimer = setTimeout(() => {
				timedOut = true;
				killProcessGroup("SIGTERM");
				forceKillTimer ??= setTimeout(() => killProcessGroup("SIGKILL"), 250);
			}, options.timeout * 1000);
		}

		child.stdout.on("data", (chunk: Buffer) => options.onStdout?.(chunk));
		child.stderr.on("data", (chunk: Buffer) => options.onStderr?.(chunk));

		waitForForegroundProcess(child as unknown as ChildProcessWithoutNullStreams)
			.then((result) => {
				cleanup();
				if (timedOut) {
					reject(new Error(`timeout:${options.timeout}`));
					return;
				}
				resolve({
					exitCode: result.code,
					signal: aborted ? "SIGTERM" : result.signal,
				});
			})
			.catch((error: unknown) => {
				cleanup();
				if (aborted) {
					resolve({ exitCode: null, signal: "SIGTERM" });
					return;
				}
				reject(error);
			});

		if (options.signal) {
			options.signal.addEventListener("abort", abortProcess, { once: true });
			if (options.signal.aborted) abortProcess();
		}
	});
}

async function withTempNuFile<T>(
	code: string,
	fn: (file: string) => Promise<T>,
): Promise<T> {
	const dir = await mkdtemp(join(tmpdir(), TEMP_PREFIX));
	tempDirs.add(dir);
	try {
		const file = join(dir, TEMP_FILE_NAME);
		await writeFile(file, code, "utf8");
		return await fn(file);
	} finally {
		tempDirs.delete(dir);
		await rm(dir, { recursive: true, force: true }).catch(() => {});
	}
}

/** Return Nu source and prefix width for Pi's ! / !! input, or undefined for normal input. */
export function stripBangPrefix(text: string): BangCommand | undefined {
	if (text.startsWith("!!")) return { command: text.slice(2), prefixLength: 2 };
	if (text.startsWith("!")) return { command: text.slice(1), prefixLength: 1 };
	return undefined;
}

export function cursorOffset(
	lines: string[],
	cursorLine: number,
	cursorCol: number,
	prefixLength: number,
): number | undefined {
	if (cursorLine === 0 && cursorCol <= prefixLength) return undefined;
	let offset = cursorCol;
	for (let i = 0; i < cursorLine; i++) offset += (lines[i] ?? "").length + 1;
	return Math.max(0, offset - prefixLength);
}

function completionToken(before: string): string {
	let tokenStart = 0;
	let quote: string | undefined;
	let escaped = false;

	for (let i = 0; i < before.length; i++) {
		const char = before[i] ?? "";
		if (quote) {
			if (escaped) {
				escaped = false;
				continue;
			}
			if (char === "\\" && quote !== "`") {
				escaped = true;
				continue;
			}
			if (char === quote) quote = undefined;
			continue;
		}

		if (escaped) {
			escaped = false;
			continue;
		}
		if (char === "\\") {
			escaped = true;
			continue;
		}
		if (char === '"' || char === "'" || char === "`") {
			quote = char;
			continue;
		}
		if (/[\s|;(){}[\]]/.test(char)) tokenStart = i + 1;
	}

	return before.slice(tokenStart);
}

export function completionPrefix(code: string, offset: number): string {
	const token = completionToken(code.slice(0, offset));
	const dot = token.lastIndexOf(".");
	const isVariableCellPath = /^\$[^\s|;(){}[\]"'`]*$/.test(token);
	return isVariableCellPath && dot >= 0 ? token.slice(dot + 1) : token;
}

function unescapeBackslashToken(value: string): string {
	return value.replace(/\\(.)/g, "$1");
}

function escapedPathCompletionInput(
	code: string,
	offset: number,
	prefix: string,
): { code: string; offset: number } | undefined {
	if (!/\\./.test(prefix)) return undefined;
	const start = offset - prefix.length;
	if (start < 0) return undefined;

	const quotedPrefix = "`" + unescapeBackslashToken(prefix);
	return {
		code: `${code.slice(0, start)}${quotedPrefix}${code.slice(offset)}`,
		offset: start + quotedPrefix.length,
	};
}

function unquoteCompletionValue(value: string): string {
	const quote = value[0];
	return (quote === "`" || quote === "'" || quote === '"') &&
		value.endsWith(quote)
		? value.slice(1, -1)
		: value;
}

function filterEscapedPathCompletions(
	values: string[],
	prefix: string,
): string[] {
	const unescapedPrefix = unescapeBackslashToken(prefix);
	return values.filter((value) =>
		unescapeBackslashToken(unquoteCompletionValue(value)).startsWith(
			unescapedPrefix,
		),
	);
}

function stringArray(value: unknown): string[] {
	return Array.isArray(value)
		? value.filter((item): item is string => typeof item === "string")
		: [];
}

function parseCompletionJson(text: string): string[] | undefined {
	try {
		const parsed = JSON.parse(text) as { completions?: unknown };
		return stringArray(parsed.completions);
	} catch {
		return undefined;
	}
}

function readQuotedItem(
	body: string,
	index: number,
): { value: string; next: number } {
	let i = index;
	while (body[i] === '"') i++;

	let value = "";
	while (i < body.length) {
		const char = body[i];
		if (char === "\\" && i + 1 < body.length) {
			value += body[i + 1];
			i += 2;
			continue;
		}
		if (char === '"') {
			let afterQuotes = i;
			while (body[afterQuotes] === '"') afterQuotes++;
			let next = afterQuotes;
			while (/\s/.test(body[next] ?? "")) next++;
			if (next >= body.length || body[next] === ",")
				return { value, next: next + 1 };
			value += body.slice(i, afterQuotes);
			i = afterQuotes;
			continue;
		}
		value += char;
		i++;
	}
	return { value: value.trim(), next: i };
}

function parseMalformedCompletionList(text: string): string[] {
	const start = text.indexOf("[");
	const end = text.lastIndexOf("]");
	if (start < 0 || end <= start) return [];

	const body = text.slice(start + 1, end);
	const completions: string[] = [];
	let i = 0;
	while (i < body.length) {
		while (/[,\s]/.test(body[i] ?? "")) i++;
		if (i >= body.length) break;

		if (body[i] === '"') {
			const item = readQuotedItem(body, i);
			if (item.value) completions.push(item.value);
			i = item.next;
			continue;
		}

		const next = body.indexOf(",", i);
		const endOfItem = next < 0 ? body.length : next;
		const value = body.slice(i, endOfItem).trim();
		if (value) completions.push(value);
		i = endOfItem + 1;
	}
	return completions;
}

export function parseCompletionsOutput(stdout: string): string[] {
	const text = stdout.trim();
	if (!text) return [];
	return (parseCompletionJson(text) ?? parseMalformedCompletionList(text)).map(
		escapeControlChars,
	);
}

export async function nuCompletions(
	code: string,
	offset: number,
	cwd: string,
	signal?: AbortSignal,
	nuCommand = DEFAULT_NU_COMMAND,
): Promise<string[]> {
	const cwdProblem = await workingDirectoryError(cwd);
	if (cwdProblem) throw new Error(cwdProblem);

	return withTempNuFile(code, async (file) => {
		const result = await runProcess(
			nuCommand,
			[...nuBaseArgs(), "-I", cwd, "--ide-complete", String(offset), file],
			cwd,
			signal,
		);
		return parseCompletionsOutput(result.stdout);
	});
}

export function createNuBashOperations(
	fixedCommand?: string,
	nuCommand = DEFAULT_NU_COMMAND,
): BashOperations {
	return {
		async exec(command, cwd, options) {
			const finish = span?.("nu.execute", { kind: "process" });
			try {
				const cwdProblem = await workingDirectoryError(cwd);
				if (cwdProblem) throw new Error(cwdProblem);
				const result = await runStreamingProcess(
					nuCommand,
					nuExecArgs(fixedCommand ?? command),
					cwd,
					{
						onStdout: options.onData,
						onStderr: options.onData,
						signal: options.signal,
						timeout: options.timeout,
						env: options.env,
					},
				);
				finish?.("finish", {
					kind: "process",
					outcome:
						result.signal !== null
							? "cancelled"
							: result.exitCode === 0
								? "ok"
								: "nonzero",
				});
				return { exitCode: result.exitCode };
			} catch (error) {
				finish?.("error", { kind: "process", outcome: "error" });
				throw error;
			}
		},
	};
}

export async function executeNu(
	command: string,
	cwd: string,
	signal?: AbortSignal,
): Promise<NuExecutionResult> {
	try {
		const cwdProblem = await workingDirectoryError(cwd);
		if (cwdProblem) throw new Error(cwdProblem);

		const result = await runProcess(
			DEFAULT_NU_COMMAND,
			nuExecArgs(command),
			cwd,
			signal,
		);
		const output = stripTerminalControls(result.output);
		const truncated = truncateOutput(output);
		return {
			output: truncated.output,
			exitCode: result.code,
			cancelled: result.signal !== null,
			truncated: truncated.truncated,
			fullOutputPath: truncated.truncated
				? await persistFullOutput(output)
				: undefined,
		};
	} catch (error) {
		const message = error instanceof Error ? error.message : String(error);
		return {
			output: `Nu execution failed: ${escapeControlChars(message)}`,
			exitCode: 1,
			cancelled: false,
			truncated: false,
		};
	}
}

function fullText(lines: string[]): string {
	return lines.join("\n");
}

function completionItems(values: string[]): AutocompleteItem[] {
	return values.map((value) => ({ value, label: value }));
}

function matchingClosingQuote(prefix: string): string | undefined {
	const quote = prefix[0];
	return quote === '"' || quote === "'" || quote === "`" ? quote : undefined;
}

function replacePrefix(
	lines: string[],
	cursorLine: number,
	cursorCol: number,
	prefix: string,
	value: string,
) {
	const next = [...lines];
	const line = next[cursorLine] ?? "";
	const before = line.slice(0, cursorCol);
	const after = line.slice(cursorCol);
	const closingQuote = matchingClosingQuote(prefix);
	const adjustedAfter =
		closingQuote && after.startsWith(closingQuote) ? after.slice(1) : after;
	next[cursorLine] =
		before.slice(0, before.length - prefix.length) + value + adjustedAfter;
	return {
		lines: next,
		cursorLine,
		cursorCol: cursorCol - prefix.length + value.length,
	};
}

export function wrapNuAutocompleteProvider(
	current: AutocompleteProvider,
	cwd: string,
	nuCommand = DEFAULT_NU_COMMAND,
): AutocompleteProvider {
	const fallbackItems = new WeakSet<AutocompleteItem>();

	return {
		async getSuggestions(lines, cursorLine, cursorCol, options) {
			const bang = stripBangPrefix(fullText(lines));
			if (!bang)
				return current.getSuggestions(lines, cursorLine, cursorCol, options);

			const offset = cursorOffset(
				lines,
				cursorLine,
				cursorCol,
				bang.prefixLength,
			);
			if (offset === undefined) return null;

			try {
				const prefix = completionPrefix(bang.command, offset);
				const retry = escapedPathCompletionInput(bang.command, offset, prefix);
				let completions: string[];
				if (retry) {
					// Escaped path prefixes need a quoted Nu retry.
					// Run the retry first so the common path pays one Nu spawn, not two.
					const retryCompletions = await nuCompletions(
						retry.code,
						retry.offset,
						cwd,
						options.signal,
						nuCommand,
					);
					const source = retryCompletions.length
						? retryCompletions
						: await nuCompletions(
								bang.command,
								offset,
								cwd,
								options.signal,
								nuCommand,
							);
					completions = filterEscapedPathCompletions(source, prefix);
				} else {
					completions = await nuCompletions(
						bang.command,
						offset,
						cwd,
						options.signal,
						nuCommand,
					);
				}
				const items = completionItems(completions);
				return items.length === 0 ? null : { prefix, items };
			} catch {
				if (options.signal.aborted) return null;
				const fallback = await current.getSuggestions(
					lines,
					cursorLine,
					cursorCol,
					options,
				);
				for (const item of fallback?.items ?? []) fallbackItems.add(item);
				return fallback;
			}
		},
		applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
			if (!stripBangPrefix(fullText(lines)) || fallbackItems.has(item))
				return current.applyCompletion(
					lines,
					cursorLine,
					cursorCol,
					item,
					prefix,
				);
			return replacePrefix(lines, cursorLine, cursorCol, prefix, item.value);
		},
		shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
			if (stripBangPrefix(fullText(lines))) return true;
			return (
				current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ??
				true
			);
		},
	};
}

async function cleanupTempDirs(): Promise<void> {
	for (const dir of [...tempDirs]) {
		await rm(dir, { recursive: true, force: true }).catch(() => {});
		tempDirs.delete(dir);
	}
}

export default function (pi: ExtensionAPI) {
	let nuCommand = DEFAULT_NU_COMMAND;

	pi.on("session_shutdown", async () => {
		dbg?.("session.shutdown");
		try {
			await cleanupTempDirs();
		} finally {
			closeDebug();
		}
	});

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		const setting = resolveSetting(pi, ctx, COMMAND_SETTING);
		if (!setting.ok)
			ctx.ui.notify(`${setting.error}; using ${DEFAULT_NU_COMMAND}`, "warning");
		nuCommand = setting.ok ? setting.value : DEFAULT_NU_COMMAND;
		ctx.ui.addAutocompleteProvider?.((current) =>
			wrapNuAutocompleteProvider(current, ctx.cwd, nuCommand),
		);
	});

	pi.on("user_bash", (event) => {
		return { operations: createNuBashOperations(event.command, nuCommand) };
	});
}