pi-ext

import { execFile } from "node:child_process";
import { existsSync, readFileSync, statSync } from "node:fs";
import { dirname, isAbsolute, join, resolve } from "node:path";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import * as focus from "./focus.js";
import { closeDebug, dbg, span } from "./src/debug.ts";

export type NotificationSeverity = "info" | "warning" | "error";
export type OscBackend = "osc777" | "osc99";

export interface NativeCommand {
	command: string;
	args: string[];
	env?: NodeJS.ProcessEnv;
}

type AgentDoneNotification = {
	title: string;
	message: string;
	severity: NotificationSeverity;
};

const PROMPT_LIMIT = 58;
const ERROR_LIMIT = 120;

function compact(value: unknown, limit = Number.POSITIVE_INFINITY): string {
	if (typeof value !== "string") return "";
	const text = value.replace(/[\x00-\x1f\x7f-\x9f]+/g, " ").trim();
	if (text.length <= limit) return text;
	return `${text.slice(0, limit).trimEnd()}…`;
}

function formatDuration(ms: number | undefined): string {
	if (ms === undefined || !Number.isFinite(ms) || ms < 0) return "";
	if (ms < 1000) return "<1s";
	const seconds = Math.round(ms / 1000);
	if (seconds < 60) return `${seconds}s`;
	const minutes = Math.floor(seconds / 60);
	const rest = seconds % 60;
	return rest === 0 ? `${minutes}m` : `${minutes}m ${rest}s`;
}

function basename(path: string): string {
	return (
		path
			.replace(/[/\\]+$/, "")
			.split(/[/\\]/)
			.pop() || path
	);
}

function textContent(value: unknown): string {
	if (typeof value === "string") return value;
	if (!Array.isArray(value)) return "";
	return value
		.map((part) => {
			const item = part as { type?: unknown; text?: unknown } | undefined;
			return item?.type === "text" && typeof item.text === "string"
				? item.text
				: "";
		})
		.filter(Boolean)
		.join(" ");
}

function promptSnippet(messages: unknown[]): string {
	for (let i = messages.length - 1; i >= 0; i--) {
		const message = messages[i] as
			| { role?: unknown; content?: unknown }
			| undefined;
		if (message?.role === "user")
			return compact(textContent(message.content), PROMPT_LIMIT);
	}
	return "";
}

function formatMessage(
	head: string,
	context: string[],
	prompt: string,
): string {
	const line = [head, ...context].filter(Boolean).join(" • ");
	return prompt ? `${line}\n↳ ${prompt}` : line;
}

export function currentGitBranch(cwd: string): string {
	try {
		let dir = resolve(cwd);
		while (true) {
			const dotGit = join(dir, ".git");
			if (existsSync(dotGit)) {
				const stat = statSync(dotGit);
				const gitDir = stat.isDirectory()
					? dotGit
					: resolveGitFile(dotGit, dir);
				const head = readFileSync(join(gitDir, "HEAD"), "utf8").trim();
				if (head.startsWith("ref: refs/heads/")) {
					return head.slice("ref: refs/heads/".length);
				}
				return /^[0-9a-f]{7,}$/i.test(head) ? head.slice(0, 7) : "";
			}
			const parent = dirname(dir);
			if (parent === dir) return "";
			dir = parent;
		}
	} catch {
		return "";
	}
}

function resolveGitFile(dotGit: string, worktreeDir: string): string {
	const text = readFileSync(dotGit, "utf8").trim();
	const prefix = "gitdir:";
	if (!text.toLowerCase().startsWith(prefix)) return dotGit;
	const target = text.slice(prefix.length).trim();
	return isAbsolute(target) ? target : resolve(worktreeDir, target);
}

export function agentWaitingNotification(
	event: {
		kind: "select" | "confirm" | "input" | "editor" | "custom";
		title?: string;
	},
	meta: { cwd?: string; branch?: string; sessionName?: string },
): AgentDoneNotification {
	const kind = {
		select: "Selection",
		confirm: "Confirmation",
		input: "Input",
		editor: "Editor",
		custom: "Request",
	}[event.kind];
	return {
		title: "Pi needs input",
		message: [
			kind,
			compact(event.title, ERROR_LIMIT),
			meta.cwd ? compact(basename(meta.cwd)) : "",
			compact(meta.branch),
			compact(meta.sessionName),
		]
			.filter(Boolean)
			.join(" • "),
		severity: "warning",
	};
}

export function agentDoneNotification(
	event: { messages?: unknown[] },
	meta: {
		cwd?: string;
		model?: { provider?: string; id?: string; name?: string };
		durationMs?: number;
		branch?: string;
		sessionName?: string;
		prompt?: string;
	},
): AgentDoneNotification {
	const messages = Array.isArray(event.messages) ? event.messages : [];
	const duration = formatDuration(meta.durationMs);
	const model = meta.model
		? compact(meta.model.name) ||
			[compact(meta.model.provider), compact(meta.model.id)]
				.filter(Boolean)
				.join("/")
		: "";
	const context = [
		meta.cwd ? compact(basename(meta.cwd)) : "",
		compact(meta.branch),
		compact(meta.sessionName),
		model,
	].filter(Boolean);
	const prompt = promptSnippet(messages) || compact(meta.prompt, PROMPT_LIMIT);

	for (let i = messages.length - 1; i >= 0; i--) {
		const message = messages[i] as
			| { role?: unknown; stopReason?: unknown; errorMessage?: unknown }
			| undefined;
		if (message?.role !== "assistant") continue;
		if (message.stopReason === "aborted") {
			return {
				title: "Pi stopped",
				message: formatMessage(
					`Stopped${duration ? ` after ${duration}` : ""}`,
					context,
					prompt,
				),
				severity: "warning",
			};
		}
		if (message.stopReason === "error") {
			return {
				title: "Pi failed",
				message: formatMessage(
					`Failed${duration ? ` after ${duration}` : ""}`,
					[compact(message.errorMessage, ERROR_LIMIT), ...context],
					prompt,
				),
				severity: "error",
			};
		}
		break;
	}

	return {
		title: "Pi finished",
		message: formatMessage(
			`Ready${duration ? ` after ${duration}` : ""}`,
			context,
			prompt,
		),
		severity: "info",
	};
}

function urgency(
	severity: NotificationSeverity,
): "low" | "normal" | "critical" {
	if (severity === "info") return "low";
	if (severity === "warning") return "normal";
	return "critical";
}

const MACOS_SCRIPT = [
	"on run argv",
	"display notification (item 2 of argv) with title (item 1 of argv)",
	"end run",
].join("\n");

const WINDOWS_SCRIPT = [
	"[Windows.UI.Notifications.ToastNotificationManager, Windows.UI.Notifications, ContentType = WindowsRuntime] > $null",
	"$xml = [Windows.UI.Notifications.ToastNotificationManager]::GetTemplateContent([Windows.UI.Notifications.ToastTemplateType]::ToastText02)",
	"$xml.GetElementsByTagName('text')[0].AppendChild($xml.CreateTextNode($env:PI_OS_NOTIFIER_TITLE)) > $null",
	"$xml.GetElementsByTagName('text')[1].AppendChild($xml.CreateTextNode($env:PI_OS_NOTIFIER_BODY)) > $null",
	"$toast = [Windows.UI.Notifications.ToastNotification]::new($xml)",
	"[Windows.UI.Notifications.ToastNotificationManager]::CreateToastNotifier('Pi').Show($toast)",
].join("; ");

export function nativeCommand(
	platform: "linux" | "darwin" | "win32",
	title: string,
	body: string,
	severity: NotificationSeverity,
): NativeCommand;
export function nativeCommand(
	platform: string,
	title: string,
	body: string,
	severity: NotificationSeverity,
): NativeCommand | undefined;
export function nativeCommand(
	platform: string,
	title: string,
	body: string,
	severity: NotificationSeverity,
): NativeCommand | undefined {
	if (platform === "linux") {
		return {
			command: "notify-send",
			args: ["--app-name=Pi", `--urgency=${urgency(severity)}`, title, body],
		};
	}
	if (platform === "darwin") {
		return {
			command: "osascript",
			args: ["-e", MACOS_SCRIPT, "--", title, body],
		};
	}
	if (platform === "win32") {
		return {
			command: "powershell.exe",
			args: [
				"-NoProfile",
				"-NonInteractive",
				"-EncodedCommand",
				Buffer.from(WINDOWS_SCRIPT, "utf16le").toString("base64"),
			],
			env: {
				...process.env,
				PI_OS_NOTIFIER_TITLE: title,
				PI_OS_NOTIFIER_BODY: body,
			},
		};
	}
	return undefined;
}

export function oscBackend(
	env: Record<string, string | undefined>,
): OscBackend {
	return env.KITTY_WINDOW_ID ? "osc99" : "osc777";
}

export function oscUserVar(name: string, value: string): string {
	return `\x1b]1337;SetUserVar=${name}=${Buffer.from(value).toString("base64")}\x07`;
}

function writeUserVar(name: string, value: string): void {
	try {
		process.stdout.write(oscUserVar(name, value));
	} catch {
		// Terminal metadata is best-effort.
	}
}

function severityPrefix(severity: NotificationSeverity): string {
	return severity === "warning" ? "⚠️ " : severity === "error" ? "❌ " : "";
}

function oscNotification(
	backend: OscBackend,
	title: string,
	body: string,
	severity: NotificationSeverity,
): string {
	const message = `${severityPrefix(severity)}${body}`;
	if (backend === "osc99") {
		return `\x1b]99;i=1:d=0;${title}\x1b\\\x1b]99;i=1:p=body;${message}\x1b\\`;
	}
	return `\x1b]777;notify;${title};${message}\x07`;
}

export function deliverNative(
	command: NativeCommand | undefined,
	title: string,
	body: string,
	severity: NotificationSeverity,
	fallback: OscBackend,
	write: (value: string) => unknown = (value) => process.stdout.write(value),
	timeoutMs = 5000,
): Promise<"native" | "fallback" | "failed"> {
	return new Promise((resolveDelivery) => {
		const useFallback = () => {
			try {
				write(oscNotification(fallback, title, body, severity));
				resolveDelivery("fallback");
			} catch {
				// No backend remains.
				resolveDelivery("failed");
			}
		};
		if (!command) {
			useFallback();
			return;
		}
		try {
			execFile(
				command.command,
				command.args,
				{ env: command.env, timeout: timeoutMs, killSignal: "SIGKILL" },
				(error) => {
					if (error) useFallback();
					else resolveDelivery("native");
				},
			);
		} catch {
			useFallback();
		}
	});
}

type InputSource = "interactive" | "rpc" | "extension" | string;
type StreamingBehavior = "steer" | "followUp" | undefined;

export class TurnEligibility {
	private active = false;
	private activeTui = false;
	private current = false;
	private pending = false;
	private prepared = false;

	input(source: InputSource, _behavior?: StreamingBehavior): void {
		if (this.active) return;
		this.pending = source === "interactive";
	}

	prepare(): void {
		if (this.active) {
			this.current ||= this.activeTui && this.pending;
			this.pending = false;
			return;
		}
		this.prepared = this.pending;
		this.pending = false;
	}

	start(mode: string): void {
		if (this.active) return;
		this.active = true;
		this.activeTui = mode === "tui";
		this.current = this.activeTui && this.prepared;
		this.prepared = false;
	}

	settle(): boolean {
		const eligible = this.isEligible();
		this.reset();
		return eligible;
	}

	isEligible(): boolean {
		return this.current;
	}

	clear(): void {
		this.reset();
	}

	private reset(): void {
		this.active = false;
		this.activeTui = false;
		this.current = false;
		this.pending = false;
		this.prepared = false;
	}
}

export default function (pi: ExtensionAPI) {
	const eligibility = new TurnEligibility();
	let agentStartedAt: number | undefined;
	let candidatePrompt: string | undefined;
	let acceptedPrompt: string | undefined;
	let lastAgentEnd: { messages?: unknown[] } | undefined;
	let focusAttached = false;
	let awaitingInput = false;
	let pendingDeliveries = 0;
	let shuttingDown = false;

	const closeWhenIdle = (): void => {
		if (shuttingDown && pendingDeliveries === 0) closeDebug();
	};

	const deliver = (notification: AgentDoneNotification): void => {
		const finish = span?.("notification.deliver", {
			platform: process.platform,
			status: notification.severity,
		});
		pendingDeliveries++;
		void deliverNative(
			nativeCommand(
				process.platform,
				notification.title,
				notification.message,
				notification.severity,
			),
			notification.title,
			notification.message,
			notification.severity,
			oscBackend(process.env),
		)
			.then((outcome) => finish?.("finish", { kind: "delivery", outcome }))
			.catch(() => finish?.("finish", { kind: "delivery", outcome: "failed" }))
			.finally(() => {
				pendingDeliveries--;
				closeWhenIdle();
			});
	};

	pi.on("session_start", async (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		if (ctx.mode === "tui") {
			focus.attach();
			writeUserVar("agent", "pi");
			writeUserVar("agent_session", pi.getSessionName() ?? "");
			writeUserVar("agent_project", basename(ctx.cwd));
			focusAttached = true;
		}
	});

	pi.on("session_info_changed", async (event) => {
		if (focusAttached) writeUserVar("agent_session", event.name ?? "");
	});

	pi.on("input", async (event) => {
		eligibility.input(event.source, event.streamingBehavior);
		if (event.streamingBehavior === undefined) {
			candidatePrompt = event.source === "interactive" ? event.text : undefined;
		}
	});

	pi.on("before_agent_start", async () => {
		eligibility.prepare();
		acceptedPrompt = candidatePrompt;
		candidatePrompt = undefined;
	});

	pi.on("agent_start", async (_event, ctx) => {
		eligibility.start(ctx.mode);
		agentStartedAt ??= Date.now();
	});

	pi.on("agent_end", async (event) => {
		lastAgentEnd = event;
	});

	type PromptEvent = {
		kind: "select" | "confirm" | "input" | "editor" | "custom";
		title?: string;
	};
	const onPromptStart = pi.on as unknown as (
		event: "ui_prompt_start",
		handler: (
			event: PromptEvent,
			ctx: ExtensionContext,
		) => Promise<void> | void,
	) => void;
	const onPromptEnd = pi.on as unknown as (
		event: "ui_prompt_end",
		handler: (
			event: PromptEvent,
			ctx: ExtensionContext,
		) => Promise<void> | void,
	) => void;
	onPromptStart("ui_prompt_start", async (event, ctx) => {
		if (awaitingInput) return;
		awaitingInput = true;
		if (!eligibility.isEligible() || focus.isFocused()) return;
		deliver(
			agentWaitingNotification(event, {
				cwd: ctx.cwd,
				branch: currentGitBranch(ctx.cwd),
				sessionName: pi.getSessionName(),
			}),
		);
	});
	onPromptEnd("ui_prompt_end", async () => {
		awaitingInput = false;
	});

	pi.on("agent_settled", async (_event, ctx) => {
		const shouldNotify = eligibility.settle();
		const focused = shouldNotify && lastAgentEnd ? focus.isFocused() : false;
		if (!shouldNotify) dbg?.("notification.skip", { reason: "ineligible" });
		else if (!lastAgentEnd)
			dbg?.("notification.skip", { reason: "missing_result" });
		else if (focused) dbg?.("notification.skip", { reason: "focused" });
		if (shouldNotify && lastAgentEnd && !focused) {
			const notification = agentDoneNotification(lastAgentEnd, {
				cwd: ctx.cwd,
				model: ctx.model,
				durationMs:
					agentStartedAt === undefined
						? undefined
						: Date.now() - agentStartedAt,
				branch: currentGitBranch(ctx.cwd),
				sessionName: pi.getSessionName(),
				prompt: acceptedPrompt,
			});
			deliver(notification);
		}
		awaitingInput = false;
		agentStartedAt = undefined;
		candidatePrompt = undefined;
		acceptedPrompt = undefined;
		lastAgentEnd = undefined;
	});

	pi.on("session_shutdown", async () => {
		dbg?.("session.shutdown");
		shuttingDown = true;
		try {
			eligibility.clear();
			agentStartedAt = undefined;
			candidatePrompt = undefined;
			acceptedPrompt = undefined;
			lastAgentEnd = undefined;
			awaitingInput = false;
			if (focusAttached) {
				writeUserVar("agent", "");
				writeUserVar("agent_session", "");
				writeUserVar("agent_project", "");
				focus.detach();
				focusAttached = false;
			}
		} finally {
			closeWhenIdle();
		}
	});
}