pi-ext

import type { Theme } from "@earendil-works/pi-coding-agent";
import type { Component, TUI } from "@earendil-works/pi-tui";
import { isProtectedKind, type MemoryAtom, type RecordChunk } from "./core.js";

const FRAME_MS = 55;
const HEIGHT = 5;
const MAP_START = 13;
const SPIN = ["◉", "◍", "◎", "◌"] as const;
const SUMMARY_HEAD = [" ", "░", "▒", "▓", "▒", "░"] as const;

type CellStyle =
	| "title"
	| "text"
	| "dim"
	| "sourceText"
	| "sourceBlock"
	| "keepBlock"
	| "durableBlock"
	| "contextBlock"
	| "pendingBlock"
	| "completeBlock"
	| "selectedBlock";

interface Cell {
	char: string;
	style: CellStyle;
}

export interface DisplayRecord {
	chunk: number;
	of: number;
	atoms: RecordChunk["atoms"];
}

export interface DecayDisplayState {
	model: string;
	inputTokens: number;
	keptTokens: number | undefined;
	totalChunks: number;
	phase: "classify" | "merge";
	records: DisplayRecord[];
	startedAt: number;
	candidates?: MemoryAtom[];
	selectedKeys?: string[];
	summaryTokens?: number;
}

export class DecayComponent implements Component {
	private readonly timer: ReturnType<typeof setInterval>;
	private disposed = false;

	constructor(
		private readonly tui: TUI,
		private readonly theme: Theme,
		readonly state: DecayDisplayState,
		private readonly now: () => number = Date.now,
	) {
		this.timer = setInterval(() => this.tui.requestRender(), FRAME_MS);
		this.timer.unref?.();
	}

	record(record: RecordChunk): void {
		if (
			record.of !== this.state.totalChunks ||
			record.chunk < 1 ||
			record.chunk > this.state.totalChunks
		)
			return;
		const displayRecord: DisplayRecord = {
			chunk: record.chunk,
			of: record.of,
			atoms: record.atoms,
		};
		const index = this.state.records.findIndex(
			(existing) => existing.chunk === record.chunk,
		);
		if (index >= 0) this.state.records[index] = displayRecord;
		else this.state.records.push(displayRecord);
		this.state.records.sort((left, right) => left.chunk - right.chunk);
		this.tui.requestRender();
	}

	merge(
		candidates: readonly MemoryAtom[],
		selected: readonly MemoryAtom[],
	): void {
		this.state.phase = "merge";
		this.state.candidates = candidates.map((atom) => ({ ...atom }));
		this.state.selectedKeys = selected.map((atom) => normalizeKey(atom.key));
		this.tui.requestRender();
	}

	setSummary(tokens: number): void {
		this.state.summaryTokens = Math.max(0, tokens);
		this.tui.requestRender();
	}

	invalidate(): void {
		this.tui.requestRender(true);
	}

	render(width: number): string[] {
		return styleFrame(
			renderDecayFrame(width, this.tui.terminal.rows, this.state, this.now()),
			this.theme,
		);
	}

	dispose(): void {
		if (this.disposed) return;
		this.disposed = true;
		clearInterval(this.timer);
	}
}

export function renderDecayFrame(
	viewportWidth: number,
	_viewportHeight: number,
	state: DecayDisplayState,
	now: number,
): Cell[][] {
	const width = Math.max(1, viewportWidth);
	const canvas = blankCanvas(width, HEIGHT);
	const elapsedMs = Math.max(0, now - state.startedAt);
	const frame = Math.floor(elapsedMs / FRAME_MS);
	const completedChunks = completedChunkCount(state);
	const phase =
		state.phase === "classify"
			? `CLASSIFYING ${completedChunks}/${state.totalChunks}`
			: "WRITING SUMMARY";

	putFragments(canvas, 0, 0, [
		["▣", "sourceText"],
		[" DECAY ", "title"],
		["· ", "dim"],
		[SPIN[frame % SPIN.length], "sourceText"],
		[` ${phase}`, "title"],
		[
			` · ${Math.floor(elapsedMs / 1_000)}s · ${modelLabel(state.model)}`,
			"dim",
		],
	]);

	renderContext(canvas, state);
	renderFacts(canvas, state);
	renderNewContext(canvas, state, frame, completedChunks);
	put(
		canvas,
		0,
		4,
		state.phase === "merge"
			? "↵ queues a message · Esc cancels · ◆ recurring · █ newer"
			: "↵ queues a message · Esc cancels",
		"dim",
	);
	return canvas;
}

function renderContext(canvas: Cell[][], state: DecayDisplayState): void {
	const capacity = Math.max(0, canvas[0].length - MAP_START);
	renderMapPrefix(canvas, 1, "CTX", compactTokens(state.inputTokens));
	if (capacity === 0) return;
	const kept = state.keptTokens;
	if (kept === undefined || kept < 0 || kept > state.inputTokens) {
		fillBlock(canvas, MAP_START, 1, capacity, "sourceBlock", [
			` ${compactTokens(state.inputTokens)} CTX `,
			" CTX ",
		]);
		return;
	}
	const summary = state.summaryTokens ?? 0;
	const scale = Math.max(1, state.inputTokens, kept + summary);
	const keepWidth = scaledWidth(kept, scale, capacity);
	const summarized = Math.max(0, state.inputTokens - kept);
	const sourceWidth =
		state.inputTokens === scale
			? Math.max(0, capacity - keepWidth)
			: Math.min(
					capacity - keepWidth,
					scaledWidth(summarized, scale, capacity),
				);
	fillBlock(canvas, MAP_START, 1, sourceWidth, "sourceBlock", [
		` ~${compactTokens(summarized)} TO SUMMARY `,
		" TO SUMMARY ",
		" S ",
	]);
	fillBlock(canvas, MAP_START + sourceWidth, 1, keepWidth, "keepBlock", [
		` ~${compactTokens(kept)} KEEP `,
		" KEEP ",
		" K ",
	]);
}

function renderFacts(canvas: Cell[][], state: DecayDisplayState): void {
	const capacity = Math.max(0, canvas[0].length - MAP_START);
	const records = uniqueRecords(state.records);
	const observed = records.flatMap((record) =>
		record.atoms.map((atom) => ({ key: atom.key, kind: atom.kind })),
	);
	const candidates =
		state.phase === "merge" ? (state.candidates ?? []) : observed;
	if (state.phase === "classify") {
		renderMapPrefix(canvas, 2, "FACTS", compactCount(observed.length));
	} else {
		const selected = new Set(state.selectedKeys ?? []);
		const selectedCount = candidates.filter((atom) =>
			selected.has(normalizeKey(atom.key)),
		).length;
		renderMapPrefix(canvas, 2, "MEM", `${selectedCount}/${candidates.length}`);
	}
	if (capacity === 0) return;

	if (state.phase === "classify") {
		const byChunk = new Map(records.map((record) => [record.chunk, record]));
		for (let index = 0; index < state.totalChunks; index++) {
			const start = Math.floor((index * capacity) / state.totalChunks);
			const end = Math.floor(((index + 1) * capacity) / state.totalChunks);
			const width = end - start;
			const record = byChunk.get(index + 1);
			if (record) {
				renderClassifiedGroups(
					canvas,
					2,
					record.atoms.map((atom) => ({ key: atom.key, kind: atom.kind })),
					width,
					undefined,
					MAP_START + start,
				);
			} else {
				fillBlock(canvas, MAP_START + start, 2, width, "pendingBlock", [], "░");
			}
		}
		return;
	}

	const selected = new Set(state.selectedKeys ?? []);
	renderClassifiedGroups(canvas, 2, candidates, capacity, selected);
}

type AtomView = Pick<MemoryAtom, "key" | "kind"> &
	Partial<Pick<MemoryAtom, "lastSeen" | "seenCount">>;

function renderClassifiedGroups(
	canvas: Cell[][],
	y: number,
	atoms: readonly AtomView[],
	width: number,
	selectedKeys: ReadonlySet<string> | undefined,
	startX = MAP_START,
): void {
	if (width <= 0) return;
	const durable = atoms.filter((atom) => isProtectedKind(atom.kind));
	const context = atoms.filter((atom) => !isProtectedKind(atom.kind));
	if (atoms.length === 0) {
		fillBlock(canvas, startX, y, width, "completeBlock", []);
		return;
	}
	const durableWidth = partitionWidth(durable.length, atoms.length, width);
	const contextWidth = width - durableWidth;
	renderAtomGroup(
		canvas,
		startX,
		y,
		durableWidth,
		durable,
		"durableBlock",
		[` ◆ ${durable.length} DURABLE `, ` ${durable.length} DUR `, " ◆ "],
		selectedKeys,
	);
	renderAtomGroup(
		canvas,
		startX + durableWidth,
		y,
		contextWidth,
		context,
		"contextBlock",
		[` ◇ ${context.length} CONTEXT `, ` ${context.length} CTX `, " ◇ "],
		selectedKeys,
	);
}

function renderAtomGroup(
	canvas: Cell[][],
	x: number,
	y: number,
	width: number,
	atoms: readonly AtomView[],
	baseStyle: CellStyle,
	labels: readonly string[],
	selectedKeys: ReadonlySet<string> | undefined,
): void {
	if (width <= 0) return;
	for (let offset = 0; offset < width; offset++) {
		const atom =
			atoms[
				Math.min(
					atoms.length - 1,
					Math.floor((offset * atoms.length) / Math.max(1, width)),
				)
			];
		const selected = atom && selectedKeys?.has(normalizeKey(atom.key));
		const glyph = selectedKeys ? atomGlyph(atom, atoms) : " ";
		put(canvas, x + offset, y, glyph, selected ? "selectedBlock" : baseStyle);
	}
	const label = fittingLabel(width, labels);
	if (label) putPreservingStyle(canvas, x, y, label);
}

function renderNewContext(
	canvas: Cell[][],
	state: DecayDisplayState,
	frame: number,
	completedChunks: number,
): void {
	const capacity = Math.max(0, canvas[0].length - MAP_START);
	const summary = state.summaryTokens;
	const kept = state.keptTokens;
	const knownKept =
		kept !== undefined && kept >= 0 && kept <= state.inputTokens;
	renderMapPrefix(
		canvas,
		3,
		"NEW",
		knownKept && summary !== undefined
			? `~${compactTokens(summary + kept)}`
			: "?",
	);
	if (capacity === 0) return;
	if (!knownKept) {
		const waiting =
			state.phase === "classify"
				? `○ waiting for ${Math.max(0, state.totalChunks - completedChunks)} chunks`
				: "○ retained size unavailable";
		put(canvas, MAP_START, 3, waiting, "dim");
		return;
	}

	const scale = Math.max(1, state.inputTokens, kept + (summary ?? 0));
	const keepWidth = scaledWidth(kept, scale, capacity);
	const keepStart = MAP_START + capacity - keepWidth;
	const summaryCapacity = Math.max(0, keepStart - MAP_START);
	const summaryWidth =
		summary === undefined
			? 0
			: Math.min(summaryCapacity, scaledWidth(summary, scale, capacity));
	if (summaryWidth > 0) {
		const summaryStart = keepStart - summaryWidth;
		fillBlock(canvas, summaryStart, 3, summaryWidth, "selectedBlock", []);
		put(
			canvas,
			summaryStart,
			3,
			SUMMARY_HEAD[frame % SUMMARY_HEAD.length],
			"selectedBlock",
		);
	} else {
		const waiting =
			state.phase === "classify"
				? `○ waiting for ${Math.max(0, state.totalChunks - completedChunks)} chunks`
				: "○ waiting for summary";
		put(canvas, MAP_START, 3, waiting.slice(0, summaryCapacity), "dim");
	}
	fillBlock(canvas, keepStart, 3, keepWidth, "keepBlock", [
		` ~${compactTokens(kept)} KEEP `,
		" KEEP ",
		" K ",
	]);
}

function renderMapPrefix(
	canvas: Cell[][],
	y: number,
	label: string,
	value: string,
): void {
	put(canvas, 0, y, label.slice(0, 6).padEnd(6), "text");
	put(canvas, 6, y, value.slice(-5).padStart(5), "title");
	put(canvas, 11, y, " │", "dim");
}

function fillBlock(
	canvas: Cell[][],
	x: number,
	y: number,
	width: number,
	style: CellStyle,
	labels: readonly string[],
	fill = " ",
): void {
	for (let offset = 0; offset < width; offset++)
		put(canvas, x + offset, y, fill, style);
	const label = fittingLabel(width, labels);
	if (label) put(canvas, x, y, label, style);
}

function fittingLabel(width: number, labels: readonly string[]): string {
	return labels.find((label) => label.length <= width) ?? "";
}

function partitionWidth(part: number, total: number, width: number): number {
	if (part <= 0) return 0;
	if (part >= total) return width;
	if (width <= 1) return Math.round(part / total);
	return Math.min(width - 1, Math.max(1, Math.round((part / total) * width)));
}

function scaledWidth(value: number, scale: number, capacity: number): number {
	if (value <= 0 || capacity <= 0) return 0;
	return Math.min(
		capacity,
		Math.max(1, Math.round((value / Math.max(1, scale)) * capacity)),
	);
}

function completedChunkCount(state: DecayDisplayState): number {
	return uniqueRecords(state.records).length;
}

function uniqueRecords(records: readonly DisplayRecord[]): DisplayRecord[] {
	const unique = new Map<number, DisplayRecord>();
	for (const record of records) unique.set(record.chunk, record);
	return [...unique.values()].sort((left, right) => left.chunk - right.chunk);
}

function atomGlyph(
	atom: AtomView | undefined,
	atoms: readonly AtomView[],
): string {
	if (!atom) return " ";
	if ((atom.seenCount ?? 0) > 1) return "◆";
	const seen = atom.lastSeen ?? 0;
	const newest = Math.max(...atoms.map((item) => item.lastSeen ?? 0), seen);
	const oldest = Math.min(...atoms.map((item) => item.lastSeen ?? 0), seen);
	const ratio = (seen - oldest) / Math.max(1, newest - oldest);
	return ratio < 0.25 ? "░" : ratio < 0.5 ? "▒" : ratio < 0.75 ? "▓" : "█";
}

function normalizeKey(key: string): string {
	return key.trim().toLowerCase().replace(/\s+/g, " ");
}

function blankCanvas(width: number, height: number): Cell[][] {
	return Array.from({ length: height }, () =>
		Array.from({ length: width }, () => ({ char: " ", style: "dim" })),
	);
}

function putFragments(
	canvas: Cell[][],
	x: number,
	y: number,
	parts: readonly (readonly [string, CellStyle])[],
): void {
	let column = x;
	for (const [text, style] of parts) {
		put(canvas, column, y, text, style);
		column += [...text].length;
	}
}

function put(
	canvas: Cell[][],
	x: number,
	y: number,
	text: string,
	style: CellStyle,
): void {
	if (y < 0 || y >= canvas.length) return;
	let column = x;
	for (const char of text) {
		if (column >= 0 && column < canvas[y].length)
			canvas[y][column] = { char, style };
		column++;
	}
}

function putPreservingStyle(
	canvas: Cell[][],
	x: number,
	y: number,
	text: string,
): void {
	if (y < 0 || y >= canvas.length) return;
	let column = x;
	for (const char of text) {
		if (column >= 0 && column < canvas[y].length) canvas[y][column].char = char;
		column++;
	}
}

function modelLabel(reference: string): string {
	const separator = reference.indexOf("/");
	return separator >= 0 ? reference.slice(separator + 1) : reference;
}

function compactTokens(tokens: number): string {
	const value = Math.max(0, tokens);
	if (value >= 1_000_000) {
		const millions = value / 1_000_000;
		return `${millions >= 10 ? Math.round(millions) : millions.toFixed(1).replace(/\.0$/, "")}m`;
	}
	if (value >= 1_000) {
		const thousands = value / 1_000;
		return `${thousands >= 10 ? Math.round(thousands) : thousands.toFixed(1).replace(/\.0$/, "")}k`;
	}
	return Math.round(value).toLocaleString("en-US");
}

function compactCount(value: number): string {
	return value >= 1_000 ? compactTokens(value) : value.toLocaleString("en-US");
}

function styleFrame(frame: Cell[][], theme: Theme): string[] {
	const style = (name: CellStyle, text: string): string => {
		switch (name) {
			case "title":
				return theme.bold(theme.fg("text", text));
			case "text":
				return theme.fg("text", text);
			case "sourceText":
				return theme.fg("success", text);
			case "sourceBlock":
				return theme.inverse(theme.fg("success", text));
			case "selectedBlock":
				return theme.inverse(theme.fg("accent", text));
			case "keepBlock":
				return theme.inverse(theme.fg("mdLink", text));
			case "durableBlock":
				return theme.inverse(theme.fg("warning", text));
			case "contextBlock":
				return theme.inverse(theme.fg("borderAccent", text));
			case "pendingBlock":
				return theme.inverse(theme.fg("muted", text));
			case "completeBlock":
				return theme.inverse(theme.fg("border", text));
			case "dim":
				return theme.fg("dim", text);
		}
	};
	return frame.map((row) => {
		let result = "";
		let active = row[0]?.style ?? "dim";
		let text = "";
		for (const cell of row) {
			if (cell.style !== active) {
				result += style(active, text);
				active = cell.style;
				text = "";
			}
			text += cell.char;
		}
		return result + style(active, text);
	});
}

export function plainFrame(frame: Cell[][]): string[] {
	return frame.map((row) => row.map((cell) => cell.char).join(""));
}