/**
 * /context — Visualize current context usage as a colored overlay.
 *
 * Shows a grid of colored squares representing token usage, broken down by:
 * - System prompt
 * - User messages
 * - Assistant text
 * - Assistant thinking
 * - Tool results (per tool: read, bash, edit, write, grep, find, ls, custom)
 * - Compaction summaries
 * - Custom/injected messages
 * - Images (estimated)
 * - Free space
 *
 * Interactive: arrow keys navigate categories, Enter shows details,
 * T toggles table/grid view, S saves a text report.
 * Also shows cache stats, cost projections, and optimization suggestions.
 */

import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ContextUsage,
	Theme,
} from "@mariozechner/pi-coding-agent";
import type { AssistantMessage, ToolResultMessage, UserMessage } from "@mariozechner/pi-ai";
import { matchesKey, truncateToWidth, visibleWidth } from "@mariozechner/pi-tui";

// ═══════════════════════════════════════════════════════════════════════
// Constants
// ═══════════════════════════════════════════════════════════════════════

const TOKENS_PER_CHAR = 4;
const TOKENS_PER_IMAGE = 1600;
const GRID_SQUARE_W = 2;
const GRID_MIN_ROWS = 6;
const GRID_MAX_ROWS = 15;

const CONTEXT_WARN = 0.8;
const CONTEXT_CRITICAL = 0.95;
const CATEGORY_WARN = 0.15;
const CATEGORY_DANGER = 0.25;
const TOOL_BIG_CONSUMER = 0.2;
const THINKING_HIGH = 0.4;
const SINGLE_RESOURCE_HIGH = 0.3;

const COMPACT_MODE_WIDTH = 60;

// ═══════════════════════════════════════════════════════════════════════
// Types
// ═══════════════════════════════════════════════════════════════════════

interface Category {
	key: string;
	label: string;
	tokens: number;
	colorCode: number;
	square: string;
	highlightedSquare: string;
	pct: number; // precomputed percentage of contextWindow
}

interface ToolStats {
	tokens: number;
	callCount: number;
	maxCallTokens: number;
}

interface ContextBreakdown {
	categories: Category[];
	totalTokens: number;
	contextWindow: number;
	percent: number | null;
	cacheRead: number;
	cacheWrite: number;
	totalCost: number;
	messageCount: number;
	turnCount: number;
	userMessageCount: number;
	imageCount: number;
	toolStats: Record<string, ToolStats>;
}

interface ContextSnapshot {
	turnCount: number;
	totalTokens: number;
	contextWindow: number;
}

type ViewMode = "grid" | "table" | "detail";

// ═══════════════════════════════════════════════════════════════════════
// ANSI helpers
// ═══════════════════════════════════════════════════════════════════════

function ansi256Bg(code: number, text: string): string {
	return `\x1b[48;5;${code}m${text}\x1b[0m`;
}

function ansi256Fg(code: number, text: string): string {
	return `\x1b[38;5;${code}m${text}\x1b[0m`;
}

function brightenColor(code: number): number {
	// Grays 232–255: move toward white
	if (code >= 232) return Math.min(255, code + 12);
	// Color cube 16–231: jump one row lighter
	return Math.min(231, code + 36);
}

// ═══════════════════════════════════════════════════════════════════════
// History (module-level, persists across /context invocations)
// ═══════════════════════════════════════════════════════════════════════

const contextHistory: ContextSnapshot[] = [];

function addSnapshot(b: ContextBreakdown): void {
	contextHistory.push({
		turnCount: b.turnCount,
		totalTokens: b.totalTokens,
		contextWindow: b.contextWindow,
	});
	// Keep last 20 snapshots
	if (contextHistory.length > 20) contextHistory.shift();
}

// ═══════════════════════════════════════════════════════════════════════
// Token estimation
// ═══════════════════════════════════════════════════════════════════════

function estimateStringTokens(text: string): number {
	return Math.ceil(text.length / TOKENS_PER_CHAR);
}

function estimateContentTokens(
	content: string | Array<{ type: string; [k: string]: any }>,
): { textTokens: number; imageTokens: number } {
	if (typeof content === "string") {
		return { textTokens: estimateStringTokens(content), imageTokens: 0 };
	}
	let text = 0;
	let img = 0;
	for (const block of content) {
		if (block.type === "text") {
			text += estimateStringTokens(block.text ?? "");
		} else if (block.type === "image") {
			img += TOKENS_PER_IMAGE;
		}
	}
	return { textTokens: text, imageTokens: img };
}

// ═══════════════════════════════════════════════════════════════════════
// Breakdown computation helpers
// ═══════════════════════════════════════════════════════════════════════

function estimateSystemPrompt(ctx: any): number {
	try {
		const sp = ctx.getSystemPrompt();
		return sp ? estimateStringTokens(sp) : 0;
	} catch {
		return 0;
	}
}

function processUserMessage(
	msg: UserMessage,
): { userTokens: number; imageTokens: number; imageCount: number } {
	const { textTokens, imageTokens } = estimateContentTokens(msg.content);
	return {
		userTokens: textTokens,
		imageTokens,
		imageCount: imageTokens > 0 ? 1 : 0,
	};
}

function processAssistantMessage(
	msg: AssistantMessage,
): {
	assistantTextTokens: number;
	thinkingTokens: number;
	cacheRead: number;
	cacheWrite: number;
	totalCost: number;
} {
	let text = 0;
	let thinking = 0;

	for (const block of msg.content) {
		if (block.type === "text") {
			text += estimateStringTokens(block.text);
		} else if (block.type === "thinking") {
			thinking += estimateStringTokens(block.thinking);
		}
		// Tool call blocks: small (function name + args JSON)
		if ((block as any).type === "tool_use") {
			text += estimateStringTokens(JSON.stringify((block as any).arguments ?? {}));
		}
	}

	return {
		assistantTextTokens: text,
		thinkingTokens: thinking,
		cacheRead: msg.usage.cacheRead,
		cacheWrite: msg.usage.cacheWrite,
		totalCost: msg.usage.cost.total,
	};
}

function processToolResult(
	msg: ToolResultMessage,
): { name: string; tokens: number } {
	const name = msg.toolName || "unknown";
	const { textTokens, imageTokens } = estimateContentTokens(msg.content);
	return { name, tokens: textTokens + imageTokens };
}

// ── Category builder ──────────────────────────────────────────────────

interface BuiltinToolInfo {
	label: string;
	colorCode: number;
}

const BUILTIN_TOOLS: Record<string, BuiltinToolInfo> = {
	read: { label: "Tool: read", colorCode: 73 },
	bash: { label: "Tool: bash", colorCode: 167 },
	edit: { label: "Tool: edit", colorCode: 179 },
	write: { label: "Tool: write", colorCode: 143 },
	grep: { label: "Tool: grep", colorCode: 109 },
	find: { label: "Tool: find", colorCode: 146 },
	ls: { label: "Tool: ls", colorCode: 108 },
	subagent: { label: "Tool: subagent", colorCode: 175 },
	web_search: { label: "Tool: web_search", colorCode: 74 },
	web_fetch: { label: "Tool: web_fetch", colorCode: 38 },
	ask_user_question: { label: "Tool: ask_user", colorCode: 183 },
	video_extract: { label: "Tool: video", colorCode: 204 },
	google_image_search: { label: "Tool: img_search", colorCode: 214 },
	youtube_search: { label: "Tool: yt_search", colorCode: 196 },
};

const CUSTOM_TOOL_COLORS = [132, 166, 130, 97, 136, 169, 103, 172];

function buildCategories(
	ctxWindow: number,
	systemPromptTokens: number,
	userTokens: number,
	assistantTextTokens: number,
	thinkingTokens: number,
	compactionTokens: number,
	customMessageTokens: number,
	imageTokens: number,
	toolStats: Record<string, ToolStats>,
	totalTokens: number,
): Category[] {
	const categories: Category[] = [];

	const addCat = (key: string, label: string, tokens: number, colorCode: number) => {
		if (tokens <= 0) return;
		categories.push({
			key,
			label,
			tokens,
			colorCode,
			square: ansi256Bg(colorCode, "  "),
			highlightedSquare: ansi256Bg(brightenColor(colorCode), "▐▌"),
			pct: (tokens / ctxWindow) * 100,
		});
	};

	addCat("system", "System Prompt", systemPromptTokens, 141);
	addCat("user", "User Messages", userTokens, 75);
	addCat("assistant", "Assistant Text", assistantTextTokens, 114);
	addCat("thinking", "Thinking", thinkingTokens, 216);

	// Tools — sorted by tokens descending
	const sortedTools = Object.entries(toolStats).sort((a, b) => b[1].tokens - a[1].tokens);
	let customColorIdx = 0;
	for (const [name, stats] of sortedTools) {
		const builtin = BUILTIN_TOOLS[name];
		const colorCode = builtin?.colorCode ?? CUSTOM_TOOL_COLORS[customColorIdx++ % CUSTOM_TOOL_COLORS.length]!;
		const label = builtin?.label ?? `Tool: ${name}`;
		addCat(`tool:${name}`, label, stats.tokens, colorCode);
	}

	addCat("compaction", "Compaction", compactionTokens, 245);
	addCat("custom", "Custom Messages", customMessageTokens, 183);
	addCat("images", "Images", imageTokens, 219);

	// Free space
	const freeTokens = Math.max(0, ctxWindow - totalTokens);
	addCat("free", "Free", freeTokens, 236);

	return categories;
}

// ── Main breakdown ────────────────────────────────────────────────────

function computeBreakdown(ctx: any): ContextBreakdown | null {
	const usage: ContextUsage | undefined = ctx.getContextUsage();
	if (!usage) return null;

	const { contextWindow } = usage;
	const branch = ctx.sessionManager.getBranch();

	let systemPromptTokens = estimateSystemPrompt(ctx);
	let userTokens = 0;
	let assistantTextTokens = 0;
	let thinkingTokens = 0;
	let compactionTokens = 0;
	let customMessageTokens = 0;
	let imageTokens = 0;
	const toolStats: Record<string, ToolStats> = {};
	let cacheRead = 0;
	let cacheWrite = 0;
	let totalCost = 0;
	let turnCount = 0;
	let messageCount = 0;
	let userMessageCount = 0;
	let imageCount = 0;

	for (const entry of branch) {
		if (entry.type === "message") {
			const msg = entry.message;
			messageCount++;

			if (msg.role === "user") {
				userMessageCount++;
				const u = processUserMessage(msg as UserMessage);
				userTokens += u.userTokens;
				imageTokens += u.imageTokens;
				imageCount += u.imageCount;
			} else if (msg.role === "assistant") {
				turnCount++;
				const a = processAssistantMessage(msg as AssistantMessage);
				assistantTextTokens += a.assistantTextTokens;
				thinkingTokens += a.thinkingTokens;
				cacheRead += a.cacheRead;
				cacheWrite += a.cacheWrite;
				totalCost += a.totalCost;
			} else if (msg.role === "toolResult") {
				const tr = processToolResult(msg as ToolResultMessage);
				const existing = toolStats[tr.name] ?? { tokens: 0, callCount: 0, maxCallTokens: 0 };
				existing.tokens += tr.tokens;
				existing.callCount++;
				existing.maxCallTokens = Math.max(existing.maxCallTokens, tr.tokens);
				toolStats[tr.name] = existing;
			}
		} else if (entry.type === "compaction" || entry.type === "branch_summary") {
			compactionTokens += estimateStringTokens((entry as any).summary ?? "");
		} else if (entry.type === "custom_message") {
			const content = (entry as any).content ?? (entry as any).message?.content;
			if (typeof content === "string") {
				customMessageTokens += estimateStringTokens(content);
			} else if (Array.isArray(content)) {
				const { textTokens, imageTokens: imgs } = estimateContentTokens(content);
				customMessageTokens += textTokens;
				imageTokens += imgs;
				if (imgs > 0) imageCount++;
			}
		}
	}

	const toolTokens = Object.values(toolStats).reduce((s, t) => s + t.tokens, 0);
	const usedFromCategories =
		systemPromptTokens + userTokens + assistantTextTokens + thinkingTokens +
		toolTokens + compactionTokens + customMessageTokens + imageTokens;

	const totalTokens = usage.tokens ?? usedFromCategories;

	const categories = buildCategories(
		contextWindow,
		systemPromptTokens,
		userTokens,
		assistantTextTokens,
		thinkingTokens,
		compactionTokens,
		customMessageTokens,
		imageTokens,
		toolStats,
		totalTokens,
	);

	return {
		categories,
		totalTokens,
		contextWindow,
		percent: usage.percent,
		cacheRead,
		cacheWrite,
		totalCost,
		messageCount,
		turnCount,
		userMessageCount,
		imageCount,
		toolStats,
	};
}

// ═══════════════════════════════════════════════════════════════════════
// Suggestions
// ═══════════════════════════════════════════════════════════════════════

function generateSuggestions(breakdown: ContextBreakdown): string[] {
	const s: string[] = [];
	const pct = breakdown.percent;

	// Context-level warnings
	if (pct !== null && pct > CONTEXT_WARN) {
		s.push("⚠ Context usage above 80% — consider /compact");
	}
	if (pct !== null && pct > CONTEXT_CRITICAL) {
		s.push("🔴 Near context limit — compaction strongly recommended");
	}

	// Per-category warnings
	for (const cat of breakdown.categories) {
		const catPct = cat.tokens / breakdown.contextWindow;

		if (cat.key === "tools" || cat.key.startsWith("tool:")) {
			if (catPct > CATEGORY_DANGER) {
				s.push(
					`🔴 ${cat.label} uses ${(catPct * 100).toFixed(0)}% of context — consider /compact`,
				);
			} else if (catPct > CATEGORY_WARN) {
				s.push(
					`💡 ${cat.label} uses ${(catPct * 100).toFixed(0)}% — summarize large outputs`,
				);
			}
		}

		if (cat.key === "thinking" && catPct > THINKING_HIGH) {
			s.push("💭 Thinking is >40% of context — try simplifying the task");
		}

		if (cat.key === "compaction" && catPct > 0.3) {
			s.push("📦 Compaction summary is large — previous compaction may have been too verbose");
		}
	}

	// Total tool usage
	const totalToolPct =
		breakdown.categories
			.filter((c) => c.key.startsWith("tool:"))
			.reduce((sum, c) => sum + c.tokens / breakdown.contextWindow, 0);
	if (totalToolPct > 0.6) {
		s.push("🔧 Tools consume >60% — /compact can summarize tool outputs");
	}

	return s;
}

// ═══════════════════════════════════════════════════════════════════════
// Formatting
// ═══════════════════════════════════════════════════════════════════════

function formatTokens(n: number): string {
	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
	return `${n}`;
}

function formatPct(n: number, total: number): string {
	return `${((n / total) * 100).toFixed(1)}%`;
}

// ═══════════════════════════════════════════════════════════════════════
// Render: progress bar
// ═══════════════════════════════════════════════════════════════════════

function renderProgressBar(
	breakdown: ContextBreakdown,
	width: number,
): string[] {
	const barW = Math.min(width, 80);
	const lines: string[] = [];
	const nonFree = breakdown.categories.filter((c) => c.key !== "free");

	let bar = "";
	let remaining = barW;
	for (const cat of nonFree) {
		const segW = Math.max(1, Math.round((cat.tokens / breakdown.contextWindow) * barW));
		const actualW = Math.min(segW, remaining);
		if (actualW <= 0) continue;
		// Use a thin slice of the category's color
		bar += ansi256Bg(cat.colorCode, " ".repeat(actualW));
		remaining -= actualW;
	}
	// Free space
	if (remaining > 0) {
		const freeCat = breakdown.categories.find((c) => c.key === "free");
		bar += ansi256Bg(freeCat?.colorCode ?? 236, " ".repeat(remaining));
	}

	lines.push(bar);
	return lines;
}

// ═══════════════════════════════════════════════════════════════════════
// Render: grid
// ═══════════════════════════════════════════════════════════════════════

function renderGrid(
	breakdown: ContextBreakdown,
	width: number,
	highlightedKey: string | null,
): string[] {
	const squareW = GRID_SQUARE_W;
	const cols = Math.floor(width / squareW);
	if (cols <= 0) return [];

	const targetRows = Math.min(GRID_MAX_ROWS, Math.max(GRID_MIN_ROWS, Math.floor(width / 8)));
	const cellsTotal = cols * targetRows;
	const tokensPerCell = breakdown.contextWindow / cellsTotal;

	// Build cell array (skip free — it's fill below)
	const cells: string[] = [];
	let cellIdx = 0;
	for (const cat of breakdown.categories) {
		if (cat.key === "free") continue;
		const isHighlighted = cat.key === highlightedKey;
		const square = isHighlighted ? cat.highlightedSquare : cat.square;
		const numCells = Math.max(
			cat.tokens > 0 ? 1 : 0,
			Math.round(cat.tokens / tokensPerCell),
		);
		for (let i = 0; i < numCells && cellIdx < cellsTotal; i++) {
			cells.push(square);
			cellIdx++;
		}
	}

	// Fill remaining with free space
	const freeCat = breakdown.categories.find((c) => c.key === "free");
	const freeSquare = freeCat?.square ?? ansi256Bg(236, "  ");
	while (cellIdx < cellsTotal) {
		cells.push(freeSquare);
		cellIdx++;
	}

	// Center grid
	const gridW = cols * squareW;
	const leftPad = Math.floor((width - gridW) / 2);
	const leftPadStr = leftPad > 0 ? " ".repeat(leftPad) : "";

	const lines: string[] = [];
	for (let row = 0; row < targetRows; row++) {
		const start = row * cols;
		const end = Math.min(start + cols, cells.length);
		let line = leftPadStr;
		for (let i = start; i < end; i++) {
			line += cells[i];
		}
		lines.push(line);
	}

	return lines;
}

// ═══════════════════════════════════════════════════════════════════════
// Render: table view
// ═══════════════════════════════════════════════════════════════════════

function renderTableView(
	breakdown: ContextBreakdown,
	width: number,
	selectedIdx: number,
): string[] {
	const lines: string[] = [];
	const nonFree = breakdown.categories.filter((c) => c.key !== "free");
	const freeCat = breakdown.categories.find((c) => c.key === "free");
	const allCats = [...nonFree];
	if (freeCat) allCats.push(freeCat);

	const barW = Math.floor(width * 0.35);

	for (let i = 0; i < allCats.length; i++) {
		const cat = allCats[i]!;
		const isSelected = i === selectedIdx;
		const prefix = isSelected ? "▶ " : "  ";
		const marker = isSelected
			? ansi256Bg(cat.colorCode, " ") + ansi256Fg(cat.colorCode, "●") + " "
			: cat.square + " ";

		// Mini bar
		const segW = Math.max(1, Math.round((cat.tokens / breakdown.contextWindow) * barW));
		const bar = ansi256Bg(
			isSelected ? brightenColor(cat.colorCode) : cat.colorCode,
			" ".repeat(segW),
		) + " ".repeat(Math.max(0, barW - segW));

		// Danger indicator
		const catPct = cat.tokens / breakdown.contextWindow;
		let dangerIcon = "";
		if (cat.key !== "free" && catPct > CATEGORY_DANGER) dangerIcon = " 🔴";
		else if (cat.key !== "free" && catPct > CATEGORY_WARN) dangerIcon = " ⚠";

		const info = `${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)${dangerIcon}`;
		const line = prefix + marker + bar + " " + ansi256Fg(cat.colorCode, cat.label) + "  " + info;
		lines.push(truncateToWidth(line, width));
	}

	return lines;
}

// ═══════════════════════════════════════════════════════════════════════
// Render: detail view (selected category)
// ═══════════════════════════════════════════════════════════════════════

function renderDetailView(
	breakdown: ContextBreakdown,
	cat: Category,
	width: number,
): string[] {
	const lines: string[] = [];
	const cw = breakdown.contextWindow;
	const pct = cat.pct;

	lines.push(ansi256Fg(cat.colorCode, `▌ ${cat.label}`));
	lines.push("");
	lines.push(`${formatTokens(cat.tokens)} tokens (${pct.toFixed(1)}% of context)`);

	if (cat.key.startsWith("tool:")) {
		const toolName = cat.key.slice(5);
		const stats = breakdown.toolStats[toolName];
		if (stats) {
			const avg = stats.callCount > 0 ? Math.round(stats.tokens / stats.callCount) : 0;
			lines.push("");
			lines.push(`${stats.callCount} calls, avg ${formatTokens(avg)} tokens/call`);
			lines.push(`Largest call: ${formatTokens(stats.maxCallTokens)} tokens`);
		}
	} else if (cat.key === "user") {
		lines.push(`${breakdown.userMessageCount} messages`);
	} else if (cat.key === "assistant") {
		const thinkingCat = breakdown.categories.find((c) => c.key === "thinking");
		const thinkingPct = thinkingCat ? thinkingCat.pct : 0;
		lines.push(`Text: ${formatTokens(cat.tokens)} (${pct.toFixed(1)}%)`);
		lines.push(`Thinking: ${thinkingPct.toFixed(1)}% of context`);
	} else if (cat.key === "images") {
		lines.push(`${breakdown.imageCount} images`);
	} else if (cat.key === "free") {
		const est = breakdown.turnCount > 0
			? Math.floor((cat.tokens / breakdown.totalTokens) * breakdown.turnCount)
			: null;
		if (est !== null) lines.push(`Est. ~${est} turns remaining`);
	}

	// Category-specific suggestions
	if (cat.tokens / cw > CATEGORY_DANGER && cat.key !== "free") {
		lines.push("");
		lines.push(`⚠ This category alone uses >${(CATEGORY_DANGER * 100).toFixed(0)}% of context.`);
	}

	return lines.map((l) => truncateToWidth(l, width));
}

// ═══════════════════════════════════════════════════════════════════════
// Render: history trend
// ═══════════════════════════════════════════════════════════════════════

function renderHistory(_breakdown: ContextBreakdown, width: number): string[] {
	if (contextHistory.length < 2) return [];

	const lines: string[] = [];
	lines.push("Context trend:");
	const barW = Math.min(width - 12, 40);

	for (const snap of contextHistory) {
		const pct = snap.totalTokens / snap.contextWindow;
		const filled = Math.round(pct * barW);
		const bar = "█".repeat(filled) + "░".repeat(Math.max(0, barW - filled));
		const label = `T${snap.turnCount}`.padEnd(4);
		lines.push(
			`${label} ${bar} ${formatTokens(snap.totalTokens)}/${formatTokens(snap.contextWindow)}`,
		);
	}

	return lines;
}

// ═══════════════════════════════════════════════════════════════════════
// Interactive overlay component
// ═══════════════════════════════════════════════════════════════════════

type CachedRender = { width: number; lines: string[] } | null;

class ContextOverlay {
	private mode: ViewMode = "grid";
	private selectedIdx: number = -1; // index into non-free categories
	private cache: CachedRender = null;
	private cachedSelectedIdx: number = -1;

	constructor(
		private breakdown: ContextBreakdown,
		private theme: Theme,
		private done: (value: void) => void,
	) {
		addSnapshot(breakdown);
	}

	handleInput(data: string): void {
		const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
		const freeCat = this.breakdown.categories.find((c) => c.key === "free");
		const allCats = [...nonFree];
		if (freeCat) allCats.push(freeCat);

		if (this.mode === "detail") {
			if (matchesKey(data, "escape") || matchesKey(data, "q")) {
				this.mode = "grid";
				this.cache = null;
			}
			return;
		}

		if (matchesKey(data, "escape") || matchesKey(data, "q")) {
			this.done(undefined);
			return;
		}

		if (matchesKey(data, "enter")) {
			if (this.selectedIdx >= 0 && this.selectedIdx < nonFree.length) {
				this.mode = "detail";
				this.cache = null;
			}
			return;
		}

		if (matchesKey(data, "t")) {
			this.mode = this.mode === "table" ? "grid" : "table";
			this.cache = null;
			return;
		}

		if (matchesKey(data, "s")) {
			this.saveReport();
			return;
		}

		if (matchesKey(data, "up") || matchesKey(data, "k")) {
			if (this.selectedIdx <= 0) {
				this.selectedIdx = allCats.length - 1;
			} else {
				this.selectedIdx--;
			}
			this.cache = null;
			return;
		}

		if (matchesKey(data, "down") || matchesKey(data, "j")) {
			if (this.selectedIdx >= allCats.length - 1) {
				this.selectedIdx = 0;
			} else {
				this.selectedIdx++;
			}
			this.cache = null;
			return;
		}
	}

	invalidate(): void {
		this.cache = null;
	}

	render(width: number): string[] {
		if (this.cache && this.cache.width === width && this.cachedSelectedIdx === this.selectedIdx) {
			return this.cache.lines;
		}

		const lines: string[] = [];
		const innerW = width - 2;

		const pad = (s: string, len: number) => {
			const vis = visibleWidth(s);
			return s + " ".repeat(Math.max(0, len - vis));
		};

		const row = (content: string) =>
			this.theme.fg("border", "│") + pad(` ${content}`, innerW) + this.theme.fg("border", "│");

		const emptyRow = () => row("");

		const hr = () =>
			this.theme.fg("border", "│") +
			this.theme.fg("dim", "─".repeat(innerW)) +
			this.theme.fg("border", "│");

		// ═══ Top border & title ═══
		lines.push(this.theme.fg("border", `╭${"─".repeat(innerW)}╮`));

		const modeTag = this.mode === "table" ? " [Table]" : this.mode === "detail" ? " [Detail]" : "";
		const pctTag = this.breakdown.percent !== null ? ` (${this.breakdown.percent.toFixed(1)}%)` : "";
		const title = `Context Window Usage${pctTag}${modeTag}`;
		lines.push(row(this.theme.bold(this.theme.fg("accent", title))));

		const sub = `${formatTokens(this.breakdown.totalTokens)} / ${formatTokens(this.breakdown.contextWindow)} tokens`;
		lines.push(row(this.theme.fg("muted", sub)));

		lines.push(emptyRow());

		// ═══ Detail mode ═══
		if (this.mode === "detail") {
			const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
			const cat = nonFree[this.selectedIdx];
			if (cat) {
				const detailLines = renderDetailView(this.breakdown, cat, innerW);
				for (const dl of detailLines) {
					lines.push(row(dl));
				}
			}
			lines.push(emptyRow());
			lines.push(hr());
			lines.push(emptyRow());
			lines.push(row(this.theme.fg("dim", "Esc: back  Arrows: navigate")));
			lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));

			this.cache = { width, lines };
			this.cachedSelectedIdx = this.selectedIdx;
			return lines;
		}

		// ═══ Compact mode (narrow terminals) ═══
		if (width <= COMPACT_MODE_WIDTH) {
			const tableLines = renderTableView(this.breakdown, innerW, this.selectedIdx);
			for (const tl of tableLines) {
				lines.push(row(tl));
			}
			lines.push(emptyRow());
			lines.push(hr());
			lines.push(emptyRow());
			this.renderStatsCompact(lines, row);
			lines.push(emptyRow());
			lines.push(row(this.theme.fg("dim", "Arrows: navigate  Enter: detail  Esc: close")));
			lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));

			this.cache = { width, lines };
			this.cachedSelectedIdx = this.selectedIdx;
			return lines;
		}

		// ═══ Progress bar ═══
		const barLines = renderProgressBar(this.breakdown, innerW);
		for (const bl of barLines) {
			lines.push(
				this.theme.fg("border", "│") +
				pad(bl, innerW) +
				this.theme.fg("border", "│"),
			);
		}
		lines.push(emptyRow());

		// ═══ Grid or Table ═══
		if (this.mode === "table") {
			const tableLines = renderTableView(this.breakdown, innerW, this.selectedIdx);
			for (const tl of tableLines) {
				lines.push(row(tl));
			}
		} else {
			const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
			const highlightedKey = this.selectedIdx >= 0 && this.selectedIdx < nonFree.length
				? nonFree[this.selectedIdx]!.key
				: null;

			const gridLines = renderGrid(this.breakdown, innerW, highlightedKey);
			for (const gl of gridLines) {
				lines.push(
					this.theme.fg("border", "│") +
					pad(gl, innerW) +
					this.theme.fg("border", "│"),
				);
			}
		}

		// ═══ Remaining turns ═══
		if (this.breakdown.turnCount > 0) {
			const freeCat = this.breakdown.categories.find((c) => c.key === "free");
			if (freeCat && freeCat.tokens > 0) {
				const avgPerTurn = this.breakdown.totalTokens / this.breakdown.turnCount;
				const remaining = Math.floor(freeCat.tokens / avgPerTurn);
				lines.push(emptyRow());
				lines.push(row(this.theme.fg("muted", `Est. ~${remaining} turns remaining (avg ${formatTokens(Math.round(avgPerTurn))} tokens/turn)`)));
			}
		}

		lines.push(emptyRow());
		lines.push(hr());
		lines.push(emptyRow());

		// ═══ Legend ═══
		const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
		const freeCat = this.breakdown.categories.find((c) => c.key === "free");
		const colW = Math.floor((innerW - 2) / 2);

		const formatEntry = (cat: Category, i: number, w: number): string => {
			const isSelected = i === this.selectedIdx;
			const prefix = isSelected ? "▶" : " ";
			const square = isSelected ? cat.highlightedSquare : cat.square;
			const label = `${prefix}${square} ${ansi256Fg(cat.colorCode, cat.label)}`;
			const value = this.theme.fg(
				isSelected ? "accent" : "dim",
				`${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)`,
			);
			return pad(`${label} ${value}`, w);
		};

		for (let i = 0; i < nonFree.length; i += 2) {
			const left = nonFree[i]!;
			const right = nonFree[i + 1];
			let content = " " + formatEntry(left, i, colW);
			if (right) {
				content += formatEntry(right, i + 1, colW);
			}
			lines.push(row(content));
		}

		if (freeCat && freeCat.tokens > 0) {
			const isSelected = this.selectedIdx === nonFree.length;
			const prefix = isSelected ? "▶" : " ";
			const square = isSelected ? freeCat.highlightedSquare : freeCat.square;
			const label = `${prefix}${square} ${ansi256Fg(freeCat.colorCode, freeCat.label)}`;
			const value = this.theme.fg(
				isSelected ? "accent" : "dim",
				`${formatTokens(freeCat.tokens)} (${freeCat.pct.toFixed(1)}%)`,
			);
			lines.push(row(`${label} ${value}`));
		}

		lines.push(emptyRow());
		lines.push(hr());
		lines.push(emptyRow());

		// ═══ Stats ═══
		lines.push(row(this.theme.fg("accent", this.theme.bold("Session Stats"))));
		lines.push(row(this.renderStatsLine(innerW)));

		// Cost projection
		if (this.breakdown.totalCost > 0 && this.breakdown.totalTokens > 0) {
			const costPerTok = this.breakdown.totalCost / this.breakdown.totalTokens;
			const projected = costPerTok * this.breakdown.contextWindow;
			lines.push(row(
				this.theme.fg("muted", `Cost: $${this.breakdown.totalCost.toFixed(4)}`) +
				this.theme.fg("dim", `  |  ~$${projected.toFixed(2)} projected for full context`),
			));
		}

		// ═══ Warnings ═══
		const suggestions = generateSuggestions(this.breakdown);
		if (suggestions.length > 0) {
			lines.push(emptyRow());
			lines.push(hr());
			lines.push(emptyRow());
			for (const sugg of suggestions) {
				lines.push(row(this.theme.fg("warning", sugg)));
			}
		}

		// ═══ History ═══
		const histLines = renderHistory(this.breakdown, innerW);
		if (histLines.length > 0) {
			lines.push(emptyRow());
			lines.push(hr());
			lines.push(emptyRow());
			for (const hl of histLines) {
				lines.push(row(this.theme.fg("dim", hl)));
			}
		}

		lines.push(emptyRow());

		// ═══ Help ═══
		lines.push(row(this.theme.fg("dim", "Arrows: navigate  Enter: details  T: grid/table  S: save report  Esc: close")));

		lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));

		this.cache = { width, lines };
		this.cachedSelectedIdx = this.selectedIdx;
		return lines;
	}

	private renderStatsLine(innerW: number): string {
		const b = this.breakdown;
		const parts = [
			`Turns: ${b.turnCount}`,
			`Messages: ${b.messageCount}`,
			`Cache read: ${formatTokens(b.cacheRead)}`,
			`Cache write: ${formatTokens(b.cacheWrite)}`,
			`Cost: $${b.totalCost.toFixed(4)}`,
		];
		return parts.map((p) => this.theme.fg("muted", p)).join(this.theme.fg("dim", "  │  "));
	}

	private renderStatsCompact(lines: string[], row: (s: string) => string): void {
		const b = this.breakdown;
		lines.push(row(this.theme.fg("accent", "Session Stats")));
		lines.push(row(this.theme.fg("muted", `Turns: ${b.turnCount}  Messages: ${b.messageCount}  Cost: $${b.totalCost.toFixed(4)}`)));
		if (b.cacheRead > 0 || b.cacheWrite > 0) {
			lines.push(row(this.theme.fg("muted", `Cache read: ${formatTokens(b.cacheRead)}  Cache write: ${formatTokens(b.cacheWrite)}`)));
		}
	}

	private saveReport(): void {
		const b = this.breakdown;
		const lines: string[] = [];
		lines.push("=== Context Usage Report ===");
		lines.push(`Total: ${formatTokens(b.totalTokens)} / ${formatTokens(b.contextWindow)} (${b.percent?.toFixed(1) ?? "?"}%)`);
		lines.push(`Turns: ${b.turnCount}  Messages: ${b.messageCount}`);
		lines.push(`Cost: $${b.totalCost.toFixed(4)}`);
		lines.push("");
		lines.push("--- Categories ---");
		for (const cat of b.categories) {
			lines.push(`${cat.label}: ${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)`);
		}
		lines.push("");
		lines.push("--- Tool Details ---");
		for (const [name, stats] of Object.entries(b.toolStats)) {
			lines.push(`${name}: ${stats.callCount} calls, ${formatTokens(stats.tokens)} tokens, max ${formatTokens(stats.maxCallTokens)}`);
		}
		lines.push("");
		const suggestions = generateSuggestions(b);
		if (suggestions.length > 0) {
			lines.push("--- Suggestions ---");
			for (const s of suggestions) lines.push(s);
		}
		lines.push("");

		const report = lines.join("\n");
		// Copy to clipboard via OSC 52 if supported, otherwise just dump
		process.stdout.write(`\x1b]52;c;${Buffer.from(report).toString("base64")}\x07`);
	}
}

// ═══════════════════════════════════════════════════════════════════════
// Extension entry point
// ═══════════════════════════════════════════════════════════════════════

export default function (pi: ExtensionAPI) {
	pi.registerCommand("context", {
		description: "Visualize current context usage as a colored grid",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			const breakdown = computeBreakdown(ctx);
			if (!breakdown) {
				ctx.ui.notify("No context usage data available yet. Send a message first.", "warning");
				return;
			}

			await ctx.ui.custom<void>(
				(tui, theme, _keybindings, done) => {
					const overlay = new ContextOverlay(breakdown, theme, done);

					return {
						handleInput(data: string) {
							overlay.handleInput(data);
							tui.requestRender();
						},

						render(width: number): string[] {
							return overlay.render(width);
						},

						invalidate() {
							overlay.invalidate();
						},
					};
				},
				{
					overlay: true,
					overlayOptions: {
						anchor: "center",
						width: "80%",
						maxWidth: 100,
						minWidth: 40,
						maxHeight: "90%",
					},
				},
			);
		},
	});
}