import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { textTokens, type MessageEstimator } from "../estimate";
import type { Line, Section } from "../types";

/**
 * Conversation breakdown by role, using the injected estimator — pi's
 * `estimateTokens` in production (see `breakdown/index.ts`), a fake in tests.
 */
export const messagesSection = (
	ctx: ExtensionCommandContext,
	estimate: MessageEstimator,
): Section => {
	const entries = ctx.sessionManager.buildContextEntries();
	const byRole = new Map<string, number>();
	let compactionTokens = 0;

	for (const entry of entries) {
		if (entry.type === "message" && entry.message) {
			const t = estimate(entry.message);
			const role = entry.message.role;
			byRole.set(role, (byRole.get(role) ?? 0) + t);
		} else if (entry.type === "compaction" && "summary" in entry && entry.summary) {
			// Compaction summaries are part of what the LLM sees.
			compactionTokens += textTokens(entry.summary);
		}
	}

	const lines: Line[] = [];
	const total = [...byRole.values()].reduce((a, b) => a + b, 0) + compactionTokens;

	if (compactionTokens > 0) {
		lines.push({ label: "compaction summary", tokens: compactionTokens });
	}
	for (const role of ["assistant", "toolResult", "user"] as const) {
		const t = byRole.get(role);
		if (t && t > 0) {
			lines.push({ label: role, tokens: t });
		}
	}

	return { title: "Messages", total, lines };
};