pi-ext

import { spawn } from "node:child_process";
import { statSync } from "node:fs";
import { extname, isAbsolute, relative, resolve, sep } from "node:path";
import type {
	ExtensionAPI,
	ExtensionContext,
	SessionEntry,
} from "@earendil-works/pi-coding-agent";
import {
	Container,
	fuzzyFilter,
	Input,
	SelectList,
} from "@earendil-works/pi-tui";
import { closeDebug, dbg, span } from "./src/debug.ts";
import { findExecutable, findExecutableFile } from "./src/pi-ext-executable.ts";

const MAX_CANDIDATES = 20;
const SOURCE_ORDER = ["user", "assistant", "edit", "write"] as const;
const SPLIT_PERCENT = "69";

type CandidateSource = (typeof SOURCE_ORDER)[number];

export interface FileCandidate {
	path: string;
	displayPath: string;
	source: CandidateSource;
}

function envValue(
	env: NodeJS.ProcessEnv,
	key: string,
	platform: NodeJS.Platform,
): string | undefined {
	if (platform !== "win32") return env[key];
	const actual = Object.keys(env).find(
		(candidate) => candidate.toLowerCase() === key.toLowerCase(),
	);
	return actual ? env[actual] : undefined;
}

function executableNames(command: string, platform: NodeJS.Platform): string[] {
	if (platform !== "win32" || extname(command)) return [command];
	return [command, `${command}.exe`, `${command}.com`];
}

function executableOnPath(
	command: string,
	env: NodeJS.ProcessEnv = process.env,
	platform: NodeJS.Platform = process.platform,
): string | undefined {
	const names = executableNames(command, platform);
	if (command.includes("/") || command.includes("\\")) {
		for (const name of names) {
			const candidate = findExecutableFile(resolve(name), { platform });
			if (candidate) return candidate;
		}
		return undefined;
	}
	return findExecutable(names, {
		path: envValue(env, "PATH", platform) ?? "",
		platform,
	});
}

function existingFile(rawPath: string, cwd: string): string | undefined {
	const trimmed = rawPath.trim().replace(/^(["'])(.*)\1$/u, "$2");
	if (!trimmed) return undefined;
	const path = resolve(cwd, trimmed);
	try {
		return statSync(path).isFile() ? path : undefined;
	} catch {
		return undefined;
	}
}

function displayPath(path: string, cwd: string): string {
	const local = relative(cwd, path);
	return local && !local.startsWith(`..${sep}`) && !isAbsolute(local)
		? local
		: path;
}

function pathLike(token: string): boolean {
	return (
		token.includes("/") ||
		token.includes("\\") ||
		token.startsWith(".") ||
		/^[^\s.][^/\\\s]*\.[A-Za-z0-9_-]+$/u.test(token)
	);
}

function trimToken(token: string): string {
	return token
		.trim()
		.replace(/^[@`'"<([{]+/u, "")
		.replace(/[`'">)\]},;!?]+$/u, "")
		.replace(/[.:]+$/u, "");
}

function extractTextPaths(text: string, cwd: string): string[] {
	const raw: Array<{ value: string; structured: boolean; index: number }> = [];
	for (const match of text.matchAll(/\]\(([^)\n]+)\)/gu)) {
		raw.push({ value: match[1], structured: true, index: match.index });
	}
	for (const match of text.matchAll(/`([^`\n]+)`/gu)) {
		raw.push({ value: match[1], structured: true, index: match.index });
	}
	for (const match of text.matchAll(/\S+/gu)) {
		raw.push({ value: match[0], structured: false, index: match.index });
	}
	raw.sort(
		(left, right) =>
			left.index - right.index ||
			Number(right.structured) - Number(left.structured),
	);

	const paths: string[] = [];
	const seen = new Set<string>();
	for (const item of raw) {
		const token = trimToken(item.value);
		if (!token || (!item.structured && !pathLike(token))) continue;
		const path = existingFile(token, cwd);
		if (!path || seen.has(path)) continue;
		seen.add(path);
		paths.push(path);
	}
	return paths;
}

function textFromContent(content: unknown): string[] {
	if (typeof content === "string") return [content];
	if (!Array.isArray(content)) return [];
	return content.flatMap((part) =>
		part &&
		typeof part === "object" &&
		"type" in part &&
		part.type === "text" &&
		"text" in part &&
		typeof part.text === "string"
			? [part.text]
			: [],
	);
}

function toolPath(
	block: unknown,
): { source: "edit" | "write"; path: string } | undefined {
	if (!block || typeof block !== "object" || !("type" in block))
		return undefined;
	if (block.type !== "toolCall" || !("name" in block)) return undefined;
	if (block.name !== "edit" && block.name !== "write") return undefined;
	if (
		!("arguments" in block) ||
		!block.arguments ||
		typeof block.arguments !== "object"
	)
		return undefined;
	if (!("path" in block.arguments) || typeof block.arguments.path !== "string")
		return undefined;
	return { source: block.name, path: block.arguments.path };
}

function collectCandidates(
	entries: readonly SessionEntry[] | readonly unknown[],
	cwd: string,
): FileCandidate[] {
	const buckets = new Map(
		SOURCE_ORDER.map((source) => [source, new Map<string, FileCandidate>()]),
	);
	const found = new Set<string>();
	const add = (path: string, source: CandidateSource): void => {
		const bucket = buckets.get(source);
		if (!bucket || bucket.has(path)) return;
		bucket.set(path, { path, displayPath: displayPath(path, cwd), source });
		found.add(path);
	};

	for (let index = entries.length - 1; index >= 0; index--) {
		const entry = entries[index];
		if (!entry || typeof entry !== "object" || !("type" in entry)) continue;
		if (entry.type !== "message" || !("message" in entry)) continue;
		const message = entry.message;
		if (!message || typeof message !== "object" || !("role" in message))
			continue;
		if (message.role === "user" && "content" in message) {
			for (const text of textFromContent(message.content).toReversed()) {
				for (const path of extractTextPaths(text, cwd).toReversed())
					add(path, "user");
			}
			if (found.size >= MAX_CANDIDATES) break;
			continue;
		}
		if (message.role !== "assistant" || !("content" in message)) continue;
		for (const block of Array.isArray(message.content)
			? message.content.toReversed()
			: []) {
			const touched = toolPath(block);
			if (touched) {
				const path = existingFile(touched.path, cwd);
				if (path) add(path, touched.source);
				continue;
			}
			for (const text of textFromContent([block])) {
				for (const path of extractTextPaths(text, cwd).toReversed())
					add(path, "assistant");
			}
		}
		if (found.size >= MAX_CANDIDATES) break;
	}

	const candidates: FileCandidate[] = [];
	const seen = new Set<string>();
	for (const source of SOURCE_ORDER) {
		for (const candidate of buckets.get(source)?.values() ?? []) {
			if (seen.has(candidate.path)) continue;
			seen.add(candidate.path);
			candidates.push(candidate);
			if (candidates.length === MAX_CANDIDATES) return candidates;
		}
	}
	return candidates;
}

function fuzzyCandidates(
	candidates: readonly FileCandidate[],
	query: string,
): FileCandidate[] {
	return fuzzyFilter(
		[...candidates],
		query,
		(candidate) => candidate.displayPath,
	);
}

async function chooseCandidate(
	candidates: readonly FileCandidate[],
	ctx: Pick<ExtensionContext, "mode" | "ui">,
): Promise<FileCandidate | undefined> {
	if (ctx.mode !== "tui") {
		const selected = await ctx.ui.select(
			"Open",
			candidates.map((candidate) => candidate.displayPath),
		);
		return candidates.find((candidate) => candidate.displayPath === selected);
	}
	const selectedPath = await ctx.ui.custom<string | undefined>(
		(tui, theme, keybindings, done) => {
			const container = new Container();
			const input = new Input();
			input.focused = true;
			let list: SelectList;

			function rebuild(): void {
				const filtered = fuzzyCandidates(candidates, input.getValue());
				list = new SelectList(
					filtered.map((candidate) => ({
						value: candidate.path,
						label: candidate.displayPath,
						description: candidate.source,
					})),
					Math.min(Math.max(filtered.length, 1), 8),
					{
						selectedPrefix: (text) => theme.fg("accent", text),
						selectedText: (text) => theme.fg("accent", text),
						description: (text) => theme.fg("dim", text),
						scrollInfo: (text) => theme.fg("dim", text),
						noMatch: (text) => theme.fg("warning", text),
					},
				);
				list.onSelect = (item) => done(item.value);
				list.onCancel = () => done(undefined);
				container.clear();
				container.addChild(input);
				container.addChild(list);
			}

			rebuild();
			let focused = true;
			return {
				get focused() {
					return focused;
				},
				set focused(value: boolean) {
					focused = value;
					input.focused = value;
				},
				render: (width: number) => container.render(width),
				invalidate: () => container.invalidate(),
				handleInput(data: string) {
					if (
						keybindings.matches(data, "tui.select.up") ||
						keybindings.matches(data, "tui.select.down")
					) {
						list.handleInput(data);
						tui.requestRender();
						return;
					}
					if (keybindings.matches(data, "tui.select.confirm")) {
						const selected = list.getSelectedItem();
						if (selected) done(selected.value);
						return;
					}
					if (keybindings.matches(data, "tui.select.cancel")) {
						done(undefined);
						return;
					}
					const previous = input.getValue();
					input.handleInput(data);
					if (input.getValue() !== previous) rebuild();
					tui.requestRender();
				},
			};
		},
	);
	return candidates.find((candidate) => candidate.path === selectedPath);
}

function splitArguments(
	paneId: string,
	cwd: string,
	editor: string,
	path: string,
): string[] {
	return [
		"cli",
		"split-pane",
		"--pane-id",
		paneId,
		"--right",
		"--percent",
		SPLIT_PERCENT,
		"--cwd",
		cwd,
		"--",
		editor,
		path,
	];
}

function termuxArguments(
	userId: string | undefined,
	cwd: string,
	editor: string,
	path: string,
): string[] {
	const androidUser = /^(?:0|[1-9][0-9]*)$/u.test(userId ?? "")
		? (userId as string)
		: "0";
	return [
		"startservice",
		"--user",
		androidUser,
		"-n",
		"com.termux/com.termux.app.RunCommandService",
		"-a",
		"com.termux.RUN_COMMAND",
		"--es",
		"com.termux.RUN_COMMAND_PATH",
		editor,
		"--esa",
		"com.termux.RUN_COMMAND_ARGUMENTS",
		path,
		"--es",
		"com.termux.RUN_COMMAND_WORKDIR",
		cwd,
		"--ez",
		"com.termux.RUN_COMMAND_BACKGROUND",
		"false",
		"--es",
		"com.termux.RUN_COMMAND_SESSION_ACTION",
		"0",
		"--es",
		"com.termux.RUN_COMMAND_LABEL",
		"Pi open",
	];
}

async function launchDetached(
	command: string,
	args: string[],
	cwd: string,
): Promise<void> {
	await new Promise<void>((resolveLaunch, reject) => {
		const child = spawn(command, args, {
			cwd,
			detached: true,
			stdio: "ignore",
			windowsHide: true,
		});
		child.once("error", reject);
		child.once("spawn", () => {
			child.unref();
			resolveLaunch();
		});
	});
}

export default function openExtension(pi: ExtensionAPI): void {
	const wezterm = executableOnPath("wezterm");
	const am = executableOnPath("am");
	const editorName = process.env.EDITOR?.trim();
	const editor = editorName ? executableOnPath(editorName) : undefined;
	const paneId = process.env.WEZTERM_PANE;
	const inTermux = Boolean(process.env.TERMUX_VERSION);
	const weztermReady = Boolean(wezterm && paneId);
	const termuxReady = Boolean(inTermux && am);

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		if (!editorName) ctx.ui.notify("/open: EDITOR is not set", "warning");
		else if (!editor)
			ctx.ui.notify(`/open: EDITOR '${editorName}' is not on PATH`, "warning");
		if (!weztermReady && !termuxReady) {
			if (!wezterm) ctx.ui.notify("/open: wezterm is not on PATH", "warning");
			if (!paneId) ctx.ui.notify("/open: WEZTERM_PANE is not set", "warning");
			if (inTermux && !am) ctx.ui.notify("/open: am is not on PATH", "warning");
		}
	});

	pi.registerCommand("open", {
		description: "Open a file beside Pi",
		async handler(args, ctx) {
			const finish = span?.("open.launch", { platform: process.platform });
			try {
				if (!editor || (!weztermReady && !termuxReady)) {
					ctx.ui.notify("/open unavailable; see startup warnings", "error");
					finish?.("finish", { outcome: "unavailable" });
					return;
				}

				let path: string | undefined;
				if (args.trim()) {
					path = existingFile(args, ctx.cwd);
					if (!path) {
						ctx.ui.notify(`Not a file: ${args.trim()}`, "error");
						finish?.("finish", { outcome: "invalid" });
						return;
					}
				} else {
					const candidates = collectCandidates(
						ctx.sessionManager.getBranch(),
						ctx.cwd,
					);
					if (candidates.length === 0) {
						ctx.ui.notify("No mentioned files", "info");
						finish?.("finish", { outcome: "no_candidates" });
						return;
					}
					path =
						candidates.length === 1
							? candidates[0].path
							: (await chooseCandidate(candidates, ctx))?.path;
				}
				if (!path) {
					finish?.("finish", { outcome: "cancelled" });
					return;
				}

				try {
					if (wezterm && paneId) {
						await launchDetached(
							wezterm,
							splitArguments(paneId, ctx.cwd, editor, path),
							ctx.cwd,
						);
					} else if (am) {
						await launchDetached(
							am,
							termuxArguments(
								process.env.TERMUX__USER_ID,
								ctx.cwd,
								editor,
								path,
							),
							ctx.cwd,
						);
					}
					finish?.("finish", { kind: "launch", outcome: "launched" });
				} catch (error) {
					finish?.("error", { kind: "launch", outcome: "failed" });
					ctx.ui.notify(
						`/open failed: ${error instanceof Error ? error.message : String(error)}`,
						"error",
					);
				}
			} catch (error) {
				finish?.("error", { outcome: "failed" });
				throw error;
			}
		},
	});

	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
	});
}

export const __test = {
	chooseCandidate,
	collectCandidates,
	executableOnPath,
	extractTextPaths,
	fuzzyCandidates,
	splitArguments,
	termuxArguments,
};