pi-ext

import { spawnSync } from "node:child_process";
import { mkdtempSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Key } from "@earendil-works/pi-tui";
import { closeDebug, dbg, span } from "./src/debug.ts";

const OUTPUT_PREFIX = "pi-prompt-autopsy-";
const OUTPUT_FILENAME = "system-prompt.md";
const LABEL = {
	preamble: "[PREAMBLE]",
	core: "[CORE]",
	tools: "[TOOLS]",
	rules: "[RULES]",
	generalGuidelines: "[GENERAL GUIDELINES]",
	toolGuidelines: "[TOOL GUIDELINES]",
	docs: "[DOCS]",
	append: "[APPEND]",
	context: "[CTX]",
	skills: "[SKILLS]",
	cwd: "[CWD]",
} as const;

const SECTION_LABEL: Readonly<Record<string, string>> = {
	tools: LABEL.tools,
	rules: LABEL.rules,
	general_guidelines: LABEL.generalGuidelines,
	tool_guidelines: LABEL.toolGuidelines,
	docs: LABEL.docs,
	addendum: LABEL.append,
	project_context: LABEL.context,
	skills: LABEL.skills,
	cwd: LABEL.cwd,
};
const SECTION_START = /(?:^|\n)<([a-z][a-z0-9_-]*)>\n/g;

type Range = { start: number; end: number; label: string };
type CapturedPrompt = { cwd: string; prompt: string };

export function groupSystemPrompt(prompt: string): string {
	const ranges = findSectionRanges(prompt);
	let pos = 0;
	const out: string[] = [];
	for (const [index, range] of ranges.entries()) {
		appendSection(
			out,
			index === 0 ? LABEL.preamble : LABEL.core,
			prompt.slice(pos, range.start),
		);
		appendSection(out, range.label, prompt.slice(range.start, range.end));
		pos = range.end;
	}
	appendSection(out, LABEL.core, prompt.slice(pos));
	return out.join("\n\n");
}

function findSectionRanges(prompt: string): Range[] {
	const ranges: Range[] = [];
	const pattern = new RegExp(SECTION_START.source, "g");
	let consumed = 0;
	for (let match = pattern.exec(prompt); match; match = pattern.exec(prompt)) {
		const start = match.index + (match[0].startsWith("\n") ? 1 : 0);
		if (start < consumed) continue;
		const name = match[1];
		const end = findSectionEnd(prompt, name, pattern.lastIndex);
		if (end === undefined) continue;
		ranges.push({
			start,
			end,
			label: SECTION_LABEL[name] ?? `[SECTION ${name}]`,
		});
		consumed = end;
		pattern.lastIndex = end;
	}
	return ranges;
}

function findSectionEnd(
	prompt: string,
	name: string,
	contentStart: number,
): number | undefined {
	const opening = `<${name}>\n`;
	const closing = `\n</${name}>`;
	let depth = 1;
	let cursor = contentStart;

	while (depth > 0) {
		const nextOpening = findLineOpening(prompt, opening, cursor);
		const closeStart = prompt.indexOf(closing, cursor);
		if (closeStart < 0) return undefined;

		if (nextOpening >= 0 && nextOpening < closeStart) {
			depth++;
			cursor = nextOpening + opening.length;
			continue;
		}

		depth--;
		cursor = closeStart + closing.length;
		if (depth === 0) return cursor;
	}

	return undefined;
}

function findLineOpening(
	prompt: string,
	opening: string,
	from: number,
): number {
	for (
		let index = prompt.indexOf(opening, from);
		index >= 0;
		index = prompt.indexOf(opening, index + opening.length)
	) {
		if (index === 0 || prompt[index - 1] === "\n") return index;
	}
	return -1;
}

function appendSection(out: string[], label: string, text: string): void {
	const body = text.trim();
	if (body) out.push(`${label}\n${body}`);
}

function writePromptFile(prompt: string): string {
	const directory = mkdtempSync(join(tmpdir(), OUTPUT_PREFIX));
	const outputPath = join(directory, OUTPUT_FILENAME);
	writeFileSync(outputPath, prompt, {
		encoding: "utf8",
		mode: 0o600,
		flag: "wx",
	});
	return outputPath;
}

function openEditor(path: string): boolean {
	const editor = process.env.VISUAL || process.env.EDITOR || "vi";
	const result = spawnSync(`${editor} ${JSON.stringify(path)}`, {
		shell: true,
		stdio: "inherit",
	});
	return !result.error && result.status === 0;
}

export default function (pi: ExtensionAPI) {
	let last: CapturedPrompt | undefined;

	pi.on("session_start", () => {
		dbg?.("session.start");
		last = undefined;
	});
	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		last = undefined;
		closeDebug();
	});
	pi.on("before_agent_start", (event, ctx) => {
		last = { cwd: ctx.cwd, prompt: event.systemPrompt };
	});
	pi.on("agent_start", (_event, ctx) => {
		// Overwrite prompt with final version (post all before_agent_start modifications).
		if (last?.cwd === ctx.cwd) last.prompt = ctx.getSystemPrompt();
	});

	const run = async (ctx: ExtensionContext) => {
		const captured = last?.cwd === ctx.cwd ? last : undefined;
		if (!captured) {
			dbg?.("prompt.skip", { reason: "not_captured" });
			ctx.ui.notify(
				"No provider-effective prompt captured since session start or reload. Send a prompt, then retry.",
				"warning",
			);
			return;
		}
		const finish = span?.("prompt.open");
		try {
			const outputPath = writePromptFile(groupSystemPrompt(captured.prompt));
			const ok = openEditor(outputPath);
			ctx.ui.notify(
				ok ? `wrote ${outputPath}` : `wrote ${outputPath}; editor failed`,
				ok ? "info" : "warning",
			);
			finish?.("finish", { outcome: ok ? "opened" : "editor_failed" });
		} catch (error) {
			finish?.("error", { outcome: "failed" });
			throw error;
		}
	};

	pi.registerCommand("prompt-autopsy", {
		description: "Autopsy current system prompt in VISUAL",
		handler: async (_args, ctx) => run(ctx),
	});
	pi.registerShortcut(Key.ctrlAlt("a"), {
		description: "Autopsy current system prompt in VISUAL",
		handler: run,
	});
}