pi-ext

import type {
	ContextUsage,
	ExtensionAPI,
	ExtensionContext,
	ThemeColor,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg } from "./src/debug.ts";
import {
	offerFooterSegment,
	removeFooterSegment,
} from "./src/pi-ext-footer-segment.ts";
import {
	parseNumberSetting,
	resolveSetting,
	type SettingDeclaration,
	type SettingSource,
} from "./src/pi-ext-settings.ts";

const STATUS_ID = "context-size";

const FOOTER_SEGMENT = {
	id: "context-size",
	icon: "",
	zone: "llm",
	order: 1,
} as const;

type Limits = { warningPercent: number; errorPercent: number };

const DEFAULT_LIMITS: Limits = { warningPercent: 70, errorPercent: 85 };

function parsePercent(raw: unknown, source: SettingSource): number | undefined {
	const value = parseNumberSetting(raw, source);
	if (value !== undefined && (value <= 0 || value > 100))
		throw new Error("must be greater than 0 and at most 100");
	return value;
}

const WARNING_SETTING: SettingDeclaration<number> = {
	key: "context-size.warningPercent",
	env: "PI_CONTEXT_SIZE_WARNING_PERCENT",
	parse: parsePercent,
	default: DEFAULT_LIMITS.warningPercent,
};

const ERROR_SETTING: SettingDeclaration<number> = {
	key: "context-size.errorPercent",
	env: "PI_CONTEXT_SIZE_ERROR_PERCENT",
	parse: parsePercent,
	default: DEFAULT_LIMITS.errorPercent,
};

const BAR = {
	width: 10,
	empty: "─",
	full: "━",
} as const;

const UPDATE_EVENTS = [
	"session_start",
	"agent_start",
	"turn_end",
	"agent_end",
	"model_select",
	"session_compact",
] as const;

type SegmentView = {
	text: string;
	color: ThemeColor;
};

function clamp(value: number, min: number, max: number): number {
	return Math.min(max, Math.max(min, value));
}

function toneForPercent(percent: number | null, limits: Limits): ThemeColor {
	if (percent === null) return "muted";
	if (percent >= limits.errorPercent) return "error";
	if (percent >= limits.warningPercent) return "warning";
	return "muted";
}

function barTone(textTone: ThemeColor): ThemeColor {
	return textTone === "muted" ? "dim" : textTone;
}

function renderThinBar(percent: number | null): string {
	const filled =
		percent === null
			? 0
			: clamp(Math.round((percent / 100) * BAR.width), 0, BAR.width);
	return BAR.full.repeat(filled) + BAR.empty.repeat(BAR.width - filled);
}

/** Human-readable token count: 940, 1.2k, 262k, 1.2M, 12M. */
function formatTokens(tokens: number): string {
	if (tokens < 1_000) return `${Math.round(tokens)}`;
	if (tokens < 10_000) return `${(tokens / 1_000).toFixed(1)}k`;
	if (tokens < 1_000_000) return `${Math.round(tokens / 1_000)}k`;
	if (tokens < 10_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`;
	return `${Math.round(tokens / 1_000_000)}M`;
}

/**
 * TUI view: the bar alone. Its length encodes the percentage and its colour
 * encodes the pressure, so percent and token counts would only repeat it.
 */
function buildSegment(
	ctx: ExtensionContext,
	usage: ContextUsage,
	limits: Limits,
): SegmentView {
	const percent = usage.percent ?? null;
	const textTone = toneForPercent(percent, limits);

	return {
		text: ctx.ui.theme.fg(barTone(textTone), renderThinBar(percent)),
		color: textTone,
	};
}

/**
 * Non-TUI view: no bar and no glyphs, because RPC clients render the status
 * line as plain text in an unknown font.
 */
function buildPlainLabel(usage: ContextUsage): string {
	const percent = usage.percent ?? null;
	if (percent === null || usage.tokens === null) return "context unknown";
	const left = Math.max(0, usage.contextWindow - usage.tokens);
	return `context ${Math.round(percent)}% used, ${formatTokens(left)} left`;
}

export default function contextSizeExtension(pi: ExtensionAPI) {
	let limits = DEFAULT_LIMITS;

	function resolveLimits(ctx: ExtensionContext): Limits {
		const warning = resolveSetting(pi, ctx, WARNING_SETTING);
		const error = resolveSetting(pi, ctx, ERROR_SETTING);
		let problem: string;
		if (!warning.ok || !error.ok)
			problem = [warning, error]
				.flatMap((result) => (result.ok ? [] : [result.error]))
				.join("; ");
		else if (warning.value < error.value)
			return { warningPercent: warning.value, errorPercent: error.value };
		else
			problem = `${WARNING_SETTING.key} (${warning.value}) must be below ${ERROR_SETTING.key} (${error.value})`;
		if (ctx.hasUI)
			ctx.ui.notify(
				`${problem}; using defaults ${DEFAULT_LIMITS.warningPercent}% and ${DEFAULT_LIMITS.errorPercent}%`,
				"warning",
			);
		return DEFAULT_LIMITS;
	}

	function clearSegment() {
		removeFooterSegment(pi, FOOTER_SEGMENT.id);
	}

	function clearPublished(ctx: ExtensionContext) {
		clearSegment();
		ctx.ui.setStatus(STATUS_ID, undefined);
	}

	function updateSegment(ctx: ExtensionContext) {
		try {
			if (!ctx.hasUI) {
				dbg?.("update", { outcome: "no-ui" });
				return;
			}

			const usage = ctx.getContextUsage();
			if (!usage) {
				clearPublished(ctx);
				dbg?.("update", { outcome: "unknown" });
				return;
			}

			if (ctx.mode !== "tui") {
				clearSegment();
				ctx.ui.setStatus(STATUS_ID, buildPlainLabel(usage));
				dbg?.("update", { outcome: "plain" });
				return;
			}

			const view = buildSegment(ctx, usage, limits);
			const accepted = offerFooterSegment(pi, {
				...FOOTER_SEGMENT,
				text: view.text,
				color: view.color,
			});

			if (!accepted) clearSegment();
			ctx.ui.setStatus(STATUS_ID, accepted ? undefined : view.text);
			dbg?.("update", { outcome: accepted ? "footer" : "fallback" });
		} catch {
			// Ignore late events during reload/switch/shutdown; Pi marks old ctx/pi stale.
			dbg?.("update", { outcome: "error" });
		}
	}

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start");
		limits = resolveLimits(ctx);
	});

	for (const eventName of UPDATE_EVENTS) {
		pi.on(eventName as any, async (_event: any, ctx: any) =>
			updateSegment(ctx),
		);
	}

	pi.on("session_shutdown", async (_event, ctx) => {
		if (ctx.hasUI) ctx.ui.setStatus(STATUS_ID, undefined);
		dbg?.("session.shutdown");
		closeDebug();
	});

	// Footer owns its own teardown. Emitting through `pi.events` during session_shutdown
	// can trip Pi's stale-context guard after reload/switch.
}

export const __test = { buildPlainLabel, formatTokens, renderThinBar };