pi-ext

import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import {
	type BuildSystemPromptOptions,
	CONFIG_DIR_NAME,
	type ExtensionAPI,
	type ExtensionContext,
	getAgentDir,
	getDocsPath,
	getExamplesPath,
	getPackageDir,
	getReadmePath,
	loadSkillsFromDir,
	VERSION,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg } from "./src/debug.ts";

const PFUI_DIR = import.meta.dirname;
const FALLBACK_PATH = join(PFUI_DIR, "SYSTEM.md");
const SKILL_DIR = join(PFUI_DIR, "skills", "pi-knowledge");
const SKILL_PATH = join(SKILL_DIR, "SKILL.md");
const GLOBAL_STATE_KEY = Symbol.for("@bugabinga/pi-ext-pfui/state");
const KNOWN_PROMPT_OPTIONS = new Set([
	"customPrompt",
	"forceSystemPrompt",
	"selectedTools",
	"toolSnippets",
	"toolGuidelines",
	"promptGuidelines",
	"appendSystemPrompt",
	"cwd",
	"contextFiles",
	"skills",
	"sections",
]);

interface PfuiGlobalState {
	warnings: Set<string>;
	deliveredWarnings: Set<string>;
}

type PfuiGlobal = typeof globalThis & {
	[GLOBAL_STATE_KEY]?: PfuiGlobalState;
};

interface EnvironmentValue {
	name: string;
	value?: string;
	path?: boolean;
}

function globalState(): PfuiGlobalState {
	const scope = globalThis as PfuiGlobal;
	const existing = scope[GLOBAL_STATE_KEY];
	if (existing) return existing;
	const created = {
		warnings: new Set<string>(),
		deliveredWarnings: new Set<string>(),
	};
	scope[GLOBAL_STATE_KEY] = created;
	return created;
}

function warn(message: string): void {
	globalState().warnings.add(message);
}

function deliverWarnings(ctx: ExtensionContext): void {
	const state = globalState();
	const pending = [...state.warnings].filter(
		(message) => !state.deliveredWarnings.has(message),
	);
	if (pending.length === 0) return;
	for (const message of pending) state.deliveredWarnings.add(message);
	ctx.ui.notify(`PFUI warnings:\n- ${pending.join("\n- ")}`, "warning");
}

function resolveEnvironmentValues(): EnvironmentValue[] {
	let packageDir: string | undefined;
	try {
		packageDir = getPackageDir();
	} catch {
		warn("PI_CONTEXT_PACKAGE_DIR could not be resolved.");
	}

	const resolvePath = (
		name: string,
		getter: () => string,
	): EnvironmentValue => {
		try {
			return { name, value: getter(), path: true };
		} catch {
			warn(`${name} could not be resolved.`);
			return { name };
		}
	};

	return [
		{ name: "PI_CONTEXT_PACKAGE_DIR", value: packageDir, path: true },
		resolvePath("PI_CONTEXT_AGENT_DIR", getAgentDir),
		resolvePath("PI_CONTEXT_README_PATH", getReadmePath),
		resolvePath("PI_CONTEXT_DOCS_PATH", getDocsPath),
		resolvePath("PI_CONTEXT_EXAMPLES_PATH", getExamplesPath),
		{
			name: "PI_CONTEXT_CHANGELOG_PATH",
			value: packageDir && join(packageDir, "CHANGELOG.md"),
			path: true,
		},
		{
			name: "PI_CONTEXT_PACKAGE_JSON_PATH",
			value: packageDir && join(packageDir, "package.json"),
			path: true,
		},
		{ name: "PI_CONTEXT_VERSION", value: VERSION },
		{ name: "PI_CONTEXT_CONFIG_DIR_NAME", value: CONFIG_DIR_NAME },
	];
}

function exportEnvironment(): Map<string, string> {
	const owned = new Map<string, string>();
	for (const resource of resolveEnvironmentValues()) {
		if (!resource.value) continue;
		if (resource.path && !existsSync(resource.value)) {
			warn(
				`${resource.name} was omitted because its resource path is missing.`,
			);
			continue;
		}

		const existing = process.env[resource.name];
		if (existing === undefined) {
			process.env[resource.name] = resource.value;
			owned.set(resource.name, resource.value);
		} else if (existing !== resource.value) {
			warn(
				`${resource.name} already has a different value; PFUI preserved it.`,
			);
		}
	}
	return owned;
}

function loadFallback(): string | undefined {
	try {
		return readFileSync(FALLBACK_PATH, "utf8").trim();
	} catch {
		warn(
			"Bundled SYSTEM.md is unavailable; PFUI will preserve Pi's system prompt.",
		);
		return undefined;
	}
}

function discoverBundledSkill(): string | undefined {
	if (!existsSync(SKILL_PATH)) {
		warn("Bundled pi-knowledge skill is missing and was omitted.");
		return undefined;
	}
	try {
		const loaded = loadSkillsFromDir({ dir: SKILL_DIR, source: "pfui" });
		const errors = loaded.diagnostics.filter(
			(diagnostic) => diagnostic.type === "error",
		);
		const skill = loaded.skills.find(
			(candidate) => candidate.name === "pi-knowledge",
		);
		if (skill && errors.length === 0) {
			for (const diagnostic of loaded.diagnostics) warn(diagnostic.message);
			return skill.filePath;
		}
		const detail = errors.map((diagnostic) => diagnostic.message).join("; ");
		warn(
			detail
				? `Bundled pi-knowledge skill is invalid and was omitted: ${detail}`
				: "Bundled pi-knowledge skill is invalid and was omitted.",
		);
	} catch (error) {
		warn(
			`Bundled pi-knowledge skill could not be validated and was omitted: ${error instanceof Error ? error.message : String(error)}`,
		);
	}
	return undefined;
}

function selectedTools(options: BuildSystemPromptOptions): string[] {
	return options.selectedTools ?? ["read", "bash", "edit", "write"];
}

function toolSection(
	options: BuildSystemPromptOptions,
	tools: string[],
): string {
	const visibleTools = [...new Set(tools)].filter(
		(name) => !!options.toolSnippets?.[name],
	);
	const toolLines =
		visibleTools.length > 0
			? visibleTools
					.map((name) => `- ${name}: ${options.toolSnippets?.[name]}`)
					.join("\n")
			: "(none)";
	return `${toolLines}\n\nIn addition to the tools above, you may have access to other custom tools depending on the project.`;
}

function uniqueGuidelines(guidelines: readonly string[]): string[] {
	const seen = new Set<string>();
	const result: string[] = [];
	for (const guideline of guidelines) {
		const normalized = guideline.trim();
		if (!normalized || seen.has(normalized)) continue;
		seen.add(normalized);
		result.push(normalized);
	}
	return result;
}

function generalGuidelinesSection(
	options: BuildSystemPromptOptions,
	tools: string[],
): string {
	const guidelines: string[] = [];
	const hasBash = tools.includes("bash");
	const hasPowerShell = tools.includes("powershell");
	if (
		(hasBash || hasPowerShell) &&
		!tools.includes("grep") &&
		!tools.includes("find") &&
		!tools.includes("ls")
	) {
		if (hasBash && hasPowerShell) {
			guidelines.push(
				"Use bash or PowerShell for file operations like listing, searching, and finding files",
			);
		} else if (hasPowerShell) {
			guidelines.push(
				"Use PowerShell for file operations like listing, searching, and finding files",
			);
		} else {
			guidelines.push(
				"Use bash for file operations like listing, searching, and finding files",
			);
		}
	}
	guidelines.push(...(options.promptGuidelines ?? []));
	return uniqueGuidelines(guidelines).join("\n");
}

function toolGuidelinesSection(
	options: BuildSystemPromptOptions,
	tools: string[],
): string {
	const groups: string[] = [];
	for (const name of new Set(tools)) {
		const guidelines = uniqueGuidelines(options.toolGuidelines?.[name] ?? []);
		if (guidelines.length === 0) continue;
		groups.push(`<tool name="${name}">\n${guidelines.join("\n")}\n</tool>`);
	}
	return groups.join("\n\n");
}

function pfuiSections(
	options: BuildSystemPromptOptions,
): Record<string, string> {
	const sections = Object.fromEntries(
		Object.entries(options.sections ?? {}).filter(
			([name, content]) =>
				content &&
				![
					"cwd",
					"docs",
					"general_guidelines",
					"preamble",
					"rules",
					"tool_guidelines",
					"tools",
				].includes(name),
		),
	);
	return Object.defineProperties(sections, {
		tools: {
			configurable: true,
			enumerable: true,
			get: () => toolSection(options, selectedTools(options)),
		},
		general_guidelines: {
			configurable: true,
			enumerable: true,
			get: () => generalGuidelinesSection(options, selectedTools(options)),
		},
		tool_guidelines: {
			configurable: true,
			enumerable: true,
			get: () => toolGuidelinesSection(options, selectedTools(options)),
		},
	});
}

function warnUnknownOptions(options: BuildSystemPromptOptions): void {
	for (const key of Object.keys(options)) {
		if (!KNOWN_PROMPT_OPTIONS.has(key)) {
			warn(
				`Unknown systemPromptOptions key ${JSON.stringify(key)} was ignored.`,
			);
		}
	}
}

export default function pfui(pi: ExtensionAPI): void {
	const ownedEnvironment = exportEnvironment();
	const fallback = loadFallback();
	const bundledSkill = discoverBundledSkill();

	pi.on("resources_discover", () =>
		bundledSkill ? { skillPaths: [bundledSkill] } : undefined,
	);
	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		deliverWarnings(ctx);
	});

	pi.on("before_agent_start", (event, ctx) => {
		if (
			event.systemPromptOptions.customPrompt ||
			event.systemPromptOptions.forceSystemPrompt
		) {
			dbg?.("prompt.apply", { outcome: "custom_prompt" });
			deliverWarnings(ctx);
			return;
		}
		warnUnknownOptions(event.systemPromptOptions);
		if (!fallback) {
			dbg?.("prompt.apply", { outcome: "fallback_missing" });
			deliverWarnings(ctx);
			return;
		}

		event.systemPromptOptions.customPrompt = fallback;
		event.systemPromptOptions.sections = pfuiSections(
			event.systemPromptOptions,
		);
		dbg?.("prompt.apply", { outcome: "applied" });
		deliverWarnings(ctx);
	});

	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		try {
			for (const [name, value] of ownedEnvironment) {
				if (process.env[name] === value) delete process.env[name];
			}
		} finally {
			closeDebug();
		}
	});
}