pi-ext

/**
 * Model route footer segments.
 *
 * Owns executor model, fast transport, and Angel advisor presentation.
 */

import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { shortModelLabel } from "./model-label.ts";
import { closeDebug, dbg } from "./src/debug.ts";
import {
	type FooterSegment,
	offerFooterSegment,
	removeFooterSegment,
} from "./src/pi-ext-footer-segment.ts";

const STATUS_ID = "model-info";
const GROUP_ID = "model-route";
const PROVIDER_SEGMENT_ID = "provider";
const SUBSCRIPTION_SEGMENT_ID = "subscription";
const MODEL_SEGMENT_ID = "model";
const FAST_SEGMENT_ID = "fast";
const ANGEL_SEGMENT_ID = "angel";
const SERVED_SEGMENT_ID = "served";
const FAST_ICON = "󱐋";
const ANGEL_ICON = "󰧑";
const SUBSCRIPTION_ICON = "󱏐";
const PROVIDER_ICONS: Readonly<Record<string, string>> = {
	anthropic: "",
	klaus: "",
	openai: "",
	"openai-codex": "",
	"github-copilot": "",
	google: "",
	"google-gemini-cli": "",
	zai: "󰬇",
	hetzner: "󰰀",
	grts: "󰙺",
};

type AngelState = { modelId: string; thinkingLevel: string };

type ReplyMessage = {
	role?: unknown;
	provider?: unknown;
	model?: unknown;
	responseModel?: unknown;
};

/**
 * The model that actually answered, when a router model (grts/auto,
 * OpenRouter auto, kilo-auto/free) resolved to a different one.
 */
function servedModel(
	message: ReplyMessage | undefined,
	ctx: Pick<ExtensionContext, "model">,
): string | undefined {
	if (
		message?.role !== "assistant" ||
		message.provider !== ctx.model?.provider ||
		message.model !== ctx.model?.id ||
		typeof message.responseModel !== "string" ||
		message.responseModel === message.model
	)
		return undefined;
	return message.responseModel;
}

function servedFromBranch(
	ctx: Pick<ExtensionContext, "model" | "sessionManager">,
): string | undefined {
	const branch = ctx.sessionManager.getBranch();
	for (let i = branch.length - 1; i >= 0; i--) {
		const entry = branch[i] as
			| { type?: unknown; message?: ReplyMessage }
			| undefined;
		if (entry?.type === "message" && entry.message?.role === "assistant")
			return servedModel(entry.message, ctx);
	}
	return undefined;
}

function modelLabel(ctx: Pick<ExtensionContext, "model">): string {
	return shortModelLabel(ctx.model?.id ?? "no-model");
}

function providerLabel(ctx: Pick<ExtensionContext, "model">): string {
	return ctx.model?.provider ?? "no-provider";
}

function providerIcon(
	ctx: Pick<ExtensionContext, "model">,
): string | undefined {
	return PROVIDER_ICONS[providerLabel(ctx)];
}

function usesSubscription(
	ctx: Pick<ExtensionContext, "model" | "modelRegistry">,
): boolean {
	const model = ctx.model;
	if (!model) return false;
	if (model.provider === "kimi-coding") return true;
	return (
		ctx.modelRegistry.isUsingOAuth(model) &&
		ctx.modelRegistry.getProvider(model.provider)?.auth.oauth
			?.isSubscription === true
	);
}

function thinkingLevelFromBranch(
	ctx: Pick<ExtensionContext, "sessionManager">,
): string {
	const branch = ctx.sessionManager.getBranch();
	for (let i = branch.length - 1; i >= 0; i--) {
		const entry = branch[i] as
			| { type?: unknown; thinkingLevel?: unknown }
			| undefined;
		if (
			entry?.type === "thinking_level_change" &&
			typeof entry.thinkingLevel === "string"
		)
			return entry.thinkingLevel;
	}
	return "off";
}

function modelInfoLabel(
	ctx: Pick<ExtensionContext, "model" | "sessionManager">,
	thinking?: string,
): string {
	return `${modelLabel(ctx)}:${thinking ?? thinkingLevelFromBranch(ctx)}`;
}

function clearSegments(pi: ExtensionAPI): void {
	for (const id of [
		PROVIDER_SEGMENT_ID,
		SUBSCRIPTION_SEGMENT_ID,
		MODEL_SEGMENT_ID,
		SERVED_SEGMENT_ID,
		FAST_SEGMENT_ID,
		ANGEL_SEGMENT_ID,
	])
		removeFooterSegment(pi, id);
}

function statusLabel(
	model: string,
	fast: boolean,
	angel: AngelState | undefined,
	served?: string,
): string {
	const parts = [model];
	if (served) parts.push(`→ ${shortModelLabel(served)}`);
	if (fast) parts.push(FAST_ICON);
	if (angel)
		parts.push(
			`→ ${ANGEL_ICON} ${shortModelLabel(angel.modelId)}:${angel.thinkingLevel}`,
		);
	return parts.join(" ");
}

export default function modelInfoExtension(pi: ExtensionAPI): void {
	let currentCtx: ExtensionContext | undefined;
	let fast = false;
	let angel: AngelState | undefined;
	let served: string | undefined;

	function publish(ctx: ExtensionContext, thinking?: string): void {
		const model = modelInfoLabel(ctx, thinking);
		ctx.ui.setStatus(FAST_SEGMENT_ID, undefined);
		ctx.ui.setStatus(ANGEL_SEGMENT_ID, undefined);
		if (ctx.mode !== "tui") {
			clearSegments(pi);
			ctx.ui.setStatus(STATUS_ID, statusLabel(model, fast, angel, served));
			dbg?.("route.refresh", { mode: ctx.mode, outcome: "status" });
			return;
		}

		let accepted = true;
		const offer = (segment: FooterSegment) => {
			accepted = offerFooterSegment(pi, segment) && accepted;
		};
		const remove = (id: string) => {
			accepted = removeFooterSegment(pi, id) && accepted;
		};
		const provider = providerIcon(ctx);
		offer({
			id: PROVIDER_SEGMENT_ID,
			text: provider ? "" : providerLabel(ctx),
			icon: provider,
			color: "dim",
			zone: "llm",
			group: GROUP_ID,
			order: -0.5,
		});
		if (usesSubscription(ctx))
			offer({
				id: SUBSCRIPTION_SEGMENT_ID,
				text: "",
				icon: SUBSCRIPTION_ICON,
				color: "success",
				zone: "llm",
				group: GROUP_ID,
				gap: 2,
				order: -0.25,
			});
		else remove(SUBSCRIPTION_SEGMENT_ID);
		offer({
			id: MODEL_SEGMENT_ID,
			text: model,
			color: "syntaxKeyword",
			zone: "llm",
			group: GROUP_ID,
			gap: 2,
			order: 0,
		});
		if (served)
			offer({
				id: SERVED_SEGMENT_ID,
				text: shortModelLabel(served),
				color: "dim",
				zone: "llm",
				group: GROUP_ID,
				separator: "→",
				order: 0.25,
			});
		else remove(SERVED_SEGMENT_ID);
		if (fast)
			offer({
				id: FAST_SEGMENT_ID,
				text: "",
				icon: FAST_ICON,
				color: "warning",
				zone: "llm",
				group: GROUP_ID,
				order: 0.5,
			});
		else remove(FAST_SEGMENT_ID);
		if (angel)
			offer({
				id: ANGEL_SEGMENT_ID,
				text: `${shortModelLabel(angel.modelId)}:${angel.thinkingLevel}`,
				icon: ANGEL_ICON,
				color: "accent",
				zone: "llm",
				group: GROUP_ID,
				separator: "→",
				order: 0.75,
			});
		else remove(ANGEL_SEGMENT_ID);

		if (!accepted) clearSegments(pi);
		ctx.ui.setStatus(
			STATUS_ID,
			accepted ? undefined : statusLabel(model, fast, angel, served),
		);
		dbg?.("route.refresh", {
			mode: ctx.mode,
			outcome: accepted ? "footer" : "fallback",
		});
	}

	function refresh(ctx: ExtensionContext, thinking?: string): void {
		currentCtx = ctx;
		if (!ctx.hasUI) {
			dbg?.("route.refresh", { mode: ctx.mode, outcome: "no_ui" });
			return;
		}
		publish(ctx, thinking ?? ctx.thinkingLevel);
	}

	pi.events.on("model-route:fast", (data: unknown) => {
		if (typeof (data as { active?: unknown })?.active !== "boolean") return;
		fast = (data as { active: boolean }).active;
		if (currentCtx) refresh(currentCtx);
	});
	pi.events.on("model-route:angel", (data: unknown) => {
		const candidate = (data as { advisor?: unknown })?.advisor;
		if (candidate === undefined) angel = undefined;
		else if (
			typeof (candidate as AngelState)?.modelId === "string" &&
			typeof (candidate as AngelState)?.thinkingLevel === "string"
		)
			angel = candidate as AngelState;
		else return;
		if (currentCtx) refresh(currentCtx);
	});

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		served = servedFromBranch(ctx);
		refresh(ctx);
	});
	pi.on("model_select", (_event, ctx) => {
		served = undefined;
		refresh(ctx);
	});
	pi.on("message_end", (event, ctx) => {
		if ((event.message as ReplyMessage).role !== "assistant") return;
		const next = servedModel(event.message as ReplyMessage, ctx);
		if (next === served) return;
		served = next;
		refresh(ctx);
	});
	pi.on("thinking_level_select", (event, ctx) => refresh(ctx, event.level));
	pi.on("session_compact", (_event, ctx) => refresh(ctx));
	pi.on("session_shutdown", (_event, ctx) => {
		dbg?.("session.shutdown");
		clearSegments(pi);
		if (ctx.hasUI) ctx.ui.setStatus(STATUS_ID, undefined);
		currentCtx = undefined;
		fast = false;
		angel = undefined;
		served = undefined;
		closeDebug();
	});
}

export const __test = {
	modelInfoLabel,
	modelLabel,
	providerIcon,
	providerLabel,
	servedFromBranch,
	servedModel,
	statusLabel,
	usesSubscription,
	thinkingLevelFromBranch,
};