pi-ext

import type { Model, UserMessage } from "@earendil-works/pi-ai";
import type { ExtensionContext } from "@earendil-works/pi-coding-agent";

export interface ActionSummaryInput {
	agentId: string;
	toolName: string;
	args: unknown;
	taskSummary: string;
	signal?: AbortSignal;
}

export type ActionSummarizer = ((
	input: ActionSummaryInput,
) => Promise<string | undefined>) & {
	finish(agentId: string): void;
	cancel(agentId: string): void;
	dispose(): void;
};

type GenerateActionSummary = ExtensionContext["modelRegistry"]["complete"];

type TimerScheduler = {
	setTimeout(callback: () => void, delayMs: number): unknown;
	clearTimeout(handle: unknown): void;
};

type Job = {
	input: ActionSummaryInput;
	resolve(summary: string | undefined): void;
	controller?: AbortController;
	cancelled?: boolean;
};

type ActionSummarizerOptions = {
	generate?: GenerateActionSummary;
	scheduler?: TimerScheduler;
};

const MAX_INPUT_CHARS = 1_500;
const MAX_SUMMARY_CHARS = 80;
const MAX_TOKENS = 40;
const TIMEOUT_MS = 2_000;
const COALESCE_MS = 350;

const SYSTEM_PROMPT = `Summarize one coding-agent tool action as a present-progress phrase.
Describe intent, not tool mechanics.
Start with an -ing verb.
Return exactly one plain-text phrase under 60 characters.
Do not use tool names, JSON, markdown, bullets, quotes, or a trailing period.`;

const ANSI_RE =
	/\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1B\\))/g;
const CONTROL_RE = /[\x00-\x1F\x7F-\x9F]/g;

const defaultScheduler: TimerScheduler = {
	setTimeout(callback, delayMs) {
		return setTimeout(callback, delayMs);
	},
	clearTimeout(handle) {
		clearTimeout(handle as ReturnType<typeof setTimeout>);
	},
};

function parseModelRef(
	value: string,
): { provider: string; id: string } | undefined {
	const separator = value.indexOf("/");
	if (separator <= 0 || separator === value.length - 1) return undefined;
	return {
		provider: value.slice(0, separator),
		id: value.slice(separator + 1),
	};
}

function cap(text: string, maxChars: number): string {
	const chars = Array.from(text);
	return chars.length <= maxChars
		? text
		: `${chars.slice(0, maxChars - 1).join("")}…`;
}

function stringify(value: unknown): string {
	if (typeof value === "string") return value;
	try {
		return JSON.stringify(value) ?? String(value);
	} catch {
		return String(value);
	}
}

function sanitize(text: string): string {
	return cap(
		text
			.replace(ANSI_RE, "")
			.replace(CONTROL_RE, " ")
			.replace(/\s+/g, " ")
			.trim()
			.replace(/^["'`]+|["'`.]+$/g, ""),
		MAX_SUMMARY_CHARS,
	);
}

function prompt(input: ActionSummaryInput): UserMessage {
	return {
		role: "user",
		content: [
			{
				type: "text",
				text: [
					`Task: ${cap(input.taskSummary, MAX_INPUT_CHARS)}`,
					`Tool: ${input.toolName}`,
					`Arguments: ${cap(stringify(input.args), MAX_INPUT_CHARS)}`,
				].join("\n"),
			},
		],
		timestamp: Date.now(),
	};
}

class ActionSummarizerEngine {
	private readonly pending = new Map<string, Job>();
	private readonly queue: Job[] = [];
	private readonly timers = new Map<string, unknown>();
	private inFlight: Job | undefined;
	private disposed = false;

	constructor(
		private readonly model: Model<any>,
		private readonly generate: GenerateActionSummary,
		private readonly scheduler: TimerScheduler,
	) {}

	enqueue(input: ActionSummaryInput): Promise<string | undefined> {
		if (this.disposed) return Promise.resolve(undefined);
		return new Promise((resolve) => {
			this.pending.get(input.agentId)?.resolve(undefined);
			this.pending.set(input.agentId, { input, resolve });
			this.clearTimer(input.agentId);
			this.timers.set(
				input.agentId,
				this.scheduler.setTimeout(() => {
					this.timers.delete(input.agentId);
					this.flush(input.agentId);
				}, COALESCE_MS),
			);
		});
	}

	finish(agentId: string): void {
		this.clearTimer(agentId);
		this.flush(agentId);
	}

	cancel(agentId: string): void {
		this.clearTimer(agentId);
		const pending = this.pending.get(agentId);
		if (pending) {
			this.pending.delete(agentId);
			pending.resolve(undefined);
		}
		for (let i = this.queue.length - 1; i >= 0; i--) {
			if (this.queue[i].input.agentId !== agentId) continue;
			this.queue.splice(i, 1)[0].resolve(undefined);
		}
		if (this.inFlight?.input.agentId === agentId) {
			this.inFlight.cancelled = true;
			this.inFlight.controller?.abort();
		}
	}

	dispose(): void {
		this.disposed = true;
		const agentIds = new Set([
			...this.pending.keys(),
			...this.queue.map((job) => job.input.agentId),
			...(this.inFlight ? [this.inFlight.input.agentId] : []),
		]);
		for (const agentId of agentIds) this.cancel(agentId);
	}

	private clearTimer(agentId: string): void {
		const timer = this.timers.get(agentId);
		if (timer === undefined) return;
		this.scheduler.clearTimeout(timer);
		this.timers.delete(agentId);
	}

	private flush(agentId: string): void {
		const job = this.pending.get(agentId);
		if (!job) return;
		this.pending.delete(agentId);
		const queued = this.queue.findIndex(
			(entry) => entry.input.agentId === agentId,
		);
		if (queued < 0) this.queue.push(job);
		else {
			this.queue[queued].resolve(undefined);
			this.queue[queued] = job;
		}
		this.pump();
	}

	private pump(): void {
		if (this.inFlight) return;
		const job = this.queue.shift();
		if (!job) return;
		this.inFlight = job;
		void this.run(job)
			.then(job.resolve, () => job.resolve(undefined))
			.finally(() => {
				if (this.inFlight === job) this.inFlight = undefined;
				this.pump();
			});
	}

	private async run(job: Job): Promise<string | undefined> {
		if (job.cancelled) return undefined;
		const controller = new AbortController();
		job.controller = controller;
		if (job.input.signal?.aborted) controller.abort();
		else
			job.input.signal?.addEventListener("abort", () => controller.abort(), {
				once: true,
			});
		const response = await this.generate(
			this.model,
			{ systemPrompt: SYSTEM_PROMPT, messages: [prompt(job.input)] },
			{
				maxTokens: MAX_TOKENS,
				maxRetries: 0,
				cacheRetention: "none",
				timeoutMs: TIMEOUT_MS,
				signal: controller.signal,
			},
		);
		if (response.stopReason !== "stop") return undefined;
		const text = response.content
			.filter(
				(part): part is { type: "text"; text: string } =>
					part.type === "text" && typeof part.text === "string",
			)
			.map((part) => part.text)
			.join("");
		return sanitize(text) || undefined;
	}
}

export function createActionSummarizer(
	ctx: Pick<ExtensionContext, "modelRegistry">,
	modelRef: string | null,
	options: ActionSummarizerOptions = {},
): ActionSummarizer | undefined {
	if (!modelRef) return undefined;
	const parsed = parseModelRef(modelRef);
	if (!parsed) return undefined;
	const model = ctx.modelRegistry.find(parsed.provider, parsed.id);
	if (!model) return undefined;
	const generate =
		options.generate ??
		((model, context, options) =>
			ctx.modelRegistry.complete(model, context, options));
	const engine = new ActionSummarizerEngine(
		model,
		generate,
		options.scheduler ?? defaultScheduler,
	);
	const summarize = ((input: ActionSummaryInput) =>
		engine.enqueue(input)) as ActionSummarizer;
	summarize.finish = (agentId) => engine.finish(agentId);
	summarize.cancel = (agentId) => engine.cancel(agentId);
	summarize.dispose = () => engine.dispose();
	return summarize;
}

export const __test = {
	ActionSummarizerEngine,
	parseModelRef,
	sanitize,
	prompt,
	SYSTEM_PROMPT,
};