pi-ext

/**
 * VCS status footer segment producer.
 *
 * Detects colocated jj before git, jj-only repos after git, and git status via porcelain v2.
 * jj queries use --ignore-working-copy to avoid snapshots.
 */

import { execFile } from "node:child_process";
import { readFile, stat } from "node:fs/promises";
import { basename, dirname, isAbsolute, join, resolve } from "node:path";
import { promisify } from "node:util";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";
import {
	type FooterSegment,
	offerFooterSegment,
	removeFooterSegment,
} from "./src/pi-ext-footer-segment.ts";
import {
	MAX_TIMER_MS,
	parseNumberSetting,
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";
import { isStaleContextError } from "./src/pi-ext-stale-context.ts";

const execFileAsync = promisify(execFile);
const CWD_STATUS_ID = "cwd";
const VCS_STATUS_ID = "vcs";
const TIMEOUT_SETTING: SettingDeclaration<number> = {
	key: "vcs-status.timeoutMs",
	env: "PI_VCS_STATUS_TIMEOUT_MS",
	default: 2500,
	parse: (raw, source) => {
		const value = parseNumberSetting(raw, source);
		return value !== undefined &&
			Number.isInteger(value) &&
			value > 0 &&
			value <= MAX_TIMER_MS
			? value
			: undefined;
	},
};

const VCS_GIT = "git";
const VCS_JJ = "jj";
const DOT_GIT = ".git";
const DOT_JJ = ".jj";
const HEAD_FILE = "HEAD";
const OBJECTS_DIR = "objects";
const REFS_DIR = "refs";
const GITDIR_PREFIX = "gitdir: ";
const REF_HEADS_PREFIX = "ref: refs/heads/";
const DETACHED_LABEL = "HEAD";
const FLAG_EMPTY = "∅";
const FLAG_CONFLICT = "×";
const FLAG_DIVERGENT = "≈";
const JJ_TEMPLATE =
	'"{\\"change\\":" ++ change_id.short().escape_json() ++ ",\\"bookmarks\\":" ++ stringify(local_bookmarks.join(",")).escape_json() ++ ",\\"empty\\":" ++ json(empty) ++ ",\\"conflict\\":" ++ json(conflict) ++ ",\\"divergent\\":" ++ json(divergent) ++ "}\\n"';

const ICONS = {
	git: "",
	jj: String.fromCodePoint(0xf329),
	ahead: "↑",
	behind: "↓",
	diverged: "⇕",
	conflicted: "=",
	untracked: "?",
	stashed: "$",
	modified: "!",
	staged: "+",
	renamed: "»",
	deleted: "✘",
};

type GitStatus = {
	branch?: string;
	ahead: number;
	behind: number;
	conflicted: number;
	untracked: number;
	stashed: boolean;
	modified: number;
	staged: number;
	renamed: number;
	deleted: number;
};

type JjStatus = {
	available: boolean;
	change?: string;
	bookmarks?: string;
	empty?: boolean;
	conflict?: boolean;
	divergent?: boolean;
};

type GitVcs = { kind: typeof VCS_GIT; gitDir: string };
type JjVcs = { kind: typeof VCS_JJ; root: string };
type Vcs = GitVcs | JjVcs;
type VcsSegment = { icon: string; text?: string };

const emptyGitStatus = (): GitStatus => ({
	branch: undefined,
	ahead: 0,
	behind: 0,
	conflicted: 0,
	untracked: 0,
	stashed: false,
	modified: 0,
	staged: 0,
	renamed: 0,
	deleted: 0,
});

async function pathType(path: string): Promise<"file" | "dir" | undefined> {
	try {
		const s = await stat(path);
		if (s.isDirectory()) return "dir";
		if (s.isFile()) return "file";
	} catch {
		return undefined;
	}
}

function expandHome(path: string): string {
	if (!path.startsWith("~")) return path;
	const home = process.env.HOME;
	return home ? join(home, path.slice(1)) : path;
}

function resolveGitPath(path: string, base: string): string {
	const expanded = expandHome(path);
	return isAbsolute(expanded) ? resolve(expanded) : resolve(base, expanded);
}

async function gitDirFromFile(
	dotgit: string,
	base: string,
): Promise<string | undefined> {
	try {
		const line = (await readFile(dotgit, "utf8")).trim();
		if (!line.startsWith(GITDIR_PREFIX)) return undefined;
		return resolveGitPath(line.slice(GITDIR_PREFIX.length), base);
	} catch {
		return undefined;
	}
}

async function isConventionalBareGitDir(dir: string): Promise<boolean> {
	return (
		basename(dir) === DOT_GIT &&
		(await pathType(join(dir, HEAD_FILE))) === "file" &&
		(await pathType(join(dir, OBJECTS_DIR))) === "dir" &&
		(await pathType(join(dir, REFS_DIR))) === "dir"
	);
}

function envGitVcs(cwd: string): GitVcs | undefined {
	const gitDir = process.env.GIT_DIR;
	if (!gitDir) return undefined;
	return { kind: VCS_GIT, gitDir: resolveGitPath(gitDir, cwd) };
}

async function findGitOrColocatedJj(cwd: string): Promise<Vcs | undefined> {
	let dir = resolve(cwd);

	while (true) {
		const dotgit = join(dir, DOT_GIT);
		const type = await pathType(dotgit);

		if (type === "dir") {
			if ((await pathType(join(dir, DOT_JJ))) === "dir")
				return { kind: VCS_JJ, root: dir };
			return { kind: VCS_GIT, gitDir: dotgit };
		}

		if (type === "file") {
			const gitDir = await gitDirFromFile(dotgit, dir);
			if (gitDir) {
				if ((await pathType(join(dir, DOT_JJ))) === "dir")
					return { kind: VCS_JJ, root: dir };
				return { kind: VCS_GIT, gitDir };
			}
		}

		if (await isConventionalBareGitDir(dir))
			return { kind: VCS_GIT, gitDir: dir };

		const parent = dirname(dir);
		if (parent === dir) return envGitVcs(cwd);
		dir = parent;
	}
}

async function findJjOnly(cwd: string): Promise<JjVcs | undefined> {
	let dir = resolve(cwd);

	while (true) {
		if ((await pathType(join(dir, DOT_JJ))) === "dir")
			return { kind: VCS_JJ, root: dir };
		const parent = dirname(dir);
		if (parent === dir) return undefined;
		dir = parent;
	}
}

async function findVcs(cwd: string): Promise<Vcs | undefined> {
	return (await findGitOrColocatedJj(cwd)) ?? (await findJjOnly(cwd));
}

function branchFromHead(head: string): string {
	return head.startsWith(REF_HEADS_PREFIX)
		? head.slice(REF_HEADS_PREFIX.length)
		: DETACHED_LABEL;
}

async function gitBranchFromHead(gitDir: string): Promise<string | undefined> {
	try {
		return branchFromHead(
			(await readFile(join(gitDir, HEAD_FILE), "utf8")).trim(),
		);
	} catch {
		return undefined;
	}
}

function parsePorcelain(stdout: string, hasStash = false): GitStatus {
	const s = emptyGitStatus();
	s.stashed = hasStash;

	for (const line of stdout.split(/\r?\n/)) {
		if (!line) continue;
		if (line.startsWith("# branch.head ")) {
			const b = line.slice("# branch.head ".length).trim();
			s.branch = b === "(detached)" ? DETACHED_LABEL : b || undefined;
			continue;
		}
		if (line.startsWith("# branch.ab ")) {
			const m = line.match(/\+(\d+)\s+-(\d+)/);
			if (m) {
				s.ahead = Number(m[1]);
				s.behind = Number(m[2]);
			}
			continue;
		}
		if (line.startsWith("# stash ")) {
			s.stashed = Number(line.slice("# stash ".length).trim()) > 0;
			continue;
		}
		if (line.startsWith("#")) continue;
		if (line.startsWith("? ")) {
			s.untracked++;
			continue;
		}
		if (line.startsWith("u ")) {
			s.conflicted++;
			continue;
		}
		if (!(line.startsWith("1 ") || line.startsWith("2 "))) continue;
		const xy = line.split(" ")[1] ?? "..";
		const x = xy[0] ?? ".",
			y = xy[1] ?? ".";
		if (x === "R") s.renamed++;
		else if (x === "D") s.deleted++;
		else if (x !== "." && x !== " ") s.staged++;
		if (y === "M") s.modified++;
		else if (y === "D") s.deleted++;
	}
	return s;
}

async function readGitStatus(
	cwd: string,
	gitDir: string,
	timeout: number,
): Promise<GitStatus> {
	const headBranch = await gitBranchFromHead(gitDir);
	try {
		const { stdout } = await execFileAsync(
			"git",
			["status", "--porcelain=2", "--branch", "--show-stash"],
			{
				cwd,
				timeout,
			},
		);
		const status = parsePorcelain(
			typeof stdout === "string" ? stdout : String(stdout),
		);
		status.branch = headBranch ?? status.branch;
		return status;
	} catch {
		return { ...emptyGitStatus(), branch: headBranch };
	}
}

async function readJjStatus(
	root: string,
	timeout: number,
): Promise<JjStatus | undefined> {
	try {
		const { stdout } = await execFileAsync(
			"jj",
			[
				"-R",
				root,
				"--ignore-working-copy",
				"--no-pager",
				"--color",
				"never",
				"log",
				"--no-graph",
				"-r",
				"@",
				"-T",
				JJ_TEMPLATE,
			],
			{ timeout },
		);
		return {
			...JSON.parse(typeof stdout === "string" ? stdout : String(stdout)),
			available: true,
		};
	} catch (error: any) {
		return error?.code === "ENOENT" ? { available: false } : undefined;
	}
}

function gitLabel(status: GitStatus): string | undefined {
	if (!status.branch) return undefined;
	const allStatus = [
		status.conflicted > 0 ? ICONS.conflicted : "",
		status.stashed ? ICONS.stashed : "",
		status.deleted > 0 ? ICONS.deleted : "",
		status.renamed > 0 ? ICONS.renamed : "",
		status.modified > 0 ? ICONS.modified : "",
		status.staged > 0 ? ICONS.staged : "",
		status.untracked > 0 ? ICONS.untracked : "",
	].join("");
	const ab =
		status.ahead > 0 && status.behind > 0
			? ICONS.diverged
			: status.ahead > 0
				? ICONS.ahead
				: status.behind > 0
					? ICONS.behind
					: "";
	const statusBlock = allStatus || ab ? `[${allStatus}${ab}]` : "";
	return `${status.branch}${statusBlock ? ` ${statusBlock}` : ""}`;
}

function suffix(value?: string): string {
	return value ? ` ${value}` : "";
}

function jjFlags(status: JjStatus): string {
	return [
		status.empty ? FLAG_EMPTY : "",
		status.conflict ? FLAG_CONFLICT : "",
		status.divergent ? FLAG_DIVERGENT : "",
	].join("");
}

function jjLabel(status?: JjStatus): string | undefined {
	if (!status) return undefined;
	if (!status.available) return "?";
	if (!status.change) return undefined;
	return `@ ${status.change}${suffix(status.bookmarks)}${suffix(jjFlags(status))}`;
}

async function readVcsStatus(
	cwd: string,
	timeout: number,
): Promise<VcsSegment | undefined> {
	const vcs = await findVcs(cwd);
	if (!vcs) return undefined;
	if (vcs.kind === VCS_JJ)
		return {
			icon: ICONS.jj,
			text: jjLabel(await readJjStatus(vcs.root, timeout)),
		};
	return {
		icon: ICONS.git,
		text: gitLabel(await readGitStatus(cwd, vcs.gitDir, timeout)),
	};
}

function ignoreStaleContext(fn: () => void) {
	try {
		fn();
	} catch (error) {
		if (!isStaleContextError(error)) throw error;
	}
}

function clearSegments(pi: ExtensionAPI) {
	ignoreStaleContext(() => removeFooterSegment(pi, "cwd"));
	ignoreStaleContext(() => removeFooterSegment(pi, "vcs"));
}

function offerSegment(pi: ExtensionAPI, segment: FooterSegment): boolean {
	let accepted = false;
	ignoreStaleContext(() => {
		accepted = offerFooterSegment(pi, segment);
	});
	return accepted;
}

function removeSegment(pi: ExtensionAPI, id: string): boolean {
	let accepted = false;
	ignoreStaleContext(() => {
		accepted = removeFooterSegment(pi, id);
	});
	return accepted;
}

function cwdLabel(cwd: string): string {
	const home = process.env.HOME;
	return home && cwd.startsWith(home) ? `~${cwd.slice(home.length)}` : cwd;
}

function emitVcs(
	pi: ExtensionAPI,
	ctx: ExtensionContext,
	cwd: string,
	vcs: VcsSegment | undefined,
	mode: ExtensionContext["mode"],
) {
	const path = cwdLabel(cwd);
	let accepted = false;

	if (mode === "tui") {
		const cwdAccepted = offerSegment(pi, {
			id: "cwd",
			text: path,
			icon: "󰝰",
			color: "syntaxOperator",
			zone: "workspace",
			order: 0,
		});
		const vcsAccepted = vcs?.text
			? offerSegment(pi, {
					id: "vcs",
					text: vcs.text,
					icon: vcs.icon,
					color: "syntaxKeyword",
					zone: "workspace",
					order: 1,
				})
			: removeSegment(pi, "vcs");
		accepted = cwdAccepted && vcsAccepted;
	}

	if (!accepted) clearSegments(pi);
	ignoreStaleContext(() =>
		ctx.ui.setStatus(CWD_STATUS_ID, accepted ? undefined : path),
	);
	ignoreStaleContext(() =>
		ctx.ui.setStatus(
			VCS_STATUS_ID,
			accepted || !vcs?.text ? undefined : `${vcs.icon} ${vcs.text}`,
		),
	);
}

export const __test = {
	branchFromHead,
	findVcs,
	gitLabel,
	jjLabel,
	parsePorcelain,
};

export default function vcsStatusExtension(pi: ExtensionAPI) {
	let generation = 0;
	let hasUI = false;
	let timeoutMs = TIMEOUT_SETTING.default;

	const refresh = async (ctx: ExtensionContext) => {
		const finish = span?.("status.refresh");
		let outcome: "success" | "no_vcs" | "stale" | "no_ui" | "error" = "success";
		try {
			const ticket = generation;
			let cwd: string;
			let mode: ExtensionContext["mode"];
			try {
				cwd = ctx.cwd;
				mode = ctx.mode;
			} catch (error) {
				if (isStaleContextError(error)) {
					outcome = "stale";
					finish?.("finish", { outcome });
					return;
				}
				throw error;
			}
			if (!hasUI) {
				outcome = "no_ui";
				finish?.("finish", { outcome });
				return;
			}

			const vcs = await readVcsStatus(cwd, timeoutMs);
			if (!vcs) outcome = "no_vcs";
			if (ticket !== generation) {
				outcome = "stale";
				finish?.("finish", { outcome });
				return;
			}
			emitVcs(pi, ctx, cwd, vcs, mode);
			finish?.("finish", { outcome });
		} catch (error) {
			finish?.("error", { outcome: "error" });
			throw error;
		}
	};

	pi.on("session_start", async (_e, ctx) => {
		dbg?.("session.start");
		generation++;
		hasUI = ctx.hasUI;
		const setting = resolveSetting(pi, ctx, TIMEOUT_SETTING);
		timeoutMs = setting.ok ? setting.value : TIMEOUT_SETTING.default;
		if (!setting.ok)
			ctx.ui.notify(
				`${setting.error}; using default ${TIMEOUT_SETTING.default} ms`,
				"warning",
			);
		void refresh(ctx);
	});
	pi.on("agent_end", async (_e, ctx) => {
		if (hasUI) void refresh(ctx);
	});
	pi.on("tool_execution_end", async (_e, ctx) => {
		if (hasUI) void refresh(ctx);
	});
	pi.on("session_shutdown", async (_e, ctx) => {
		dbg?.("session.shutdown");
		closeDebug();
		generation++;
		clearSegments(pi);
		ignoreStaleContext(() => {
			if (ctx.hasUI) {
				ctx.ui.setStatus(CWD_STATUS_ID, undefined);
				ctx.ui.setStatus(VCS_STATUS_ID, undefined);
			}
		});
	});
}