/**
 * foldable-tools.ts — make pi tool blocks foldable / hideable.
 *
 * Why this exists
 * ---------------
 * pi's built-in tool blocks (the green/red "Tool execution box") are always
 * rendered with a full colored shell. The only built-in control,
 * `ctrl+o` (app.tools.expand), toggles output *verbosity* — which is invisible
 * for short outputs and never collapses the block itself. There is no way to
 * fold or hide the block. This extension adds that.
 *
 * What it does
 * ------------
 * Re-registers the built-in tools (read, bash, edit, write, grep, find, ls)
 * with `renderShell: "self"` so the always-on colored box is gone, and adds a
 * 3-state view model for COMPLETED blocks (running blocks always stream live so
 * you can see activity):
 *
 *   folded   (default) — a compact 2-line card: the call header (command/path)
 *                        + a one-line status (✓ N lines / ✗ exit N / +A -R).
 *   expanded            — the full original output (syntax-highlighted reads,
 *                        diffs, truncation notes) for read/grep/find/ls, or the
 *                        full text output for bash/edit/write.
 *   hidden              — completed blocks are removed from the transcript
 *                        entirely (0 lines).
 *
 * Controls
 * --------
 *   ctrl+q              cycle  folded → expanded → hidden → folded
 *   /tools [mode]       set mode directly (`folded` | `expanded` | `hidden`);
 *                        with no argument, cycles like ctrl+q
 *   ctrl+o (built-in)   still expands output within the folded/expanded modes
 *
 * Execution behavior is unchanged: every tool delegates execute() to the
 * original built-in tool definition, so reads/edits/bash/grep behave identically.
 *
 * Configuration
 * -------------
 *   PI_TOOL_FOLD_MODE   default mode on startup: "folded" (default),
 *                        "expanded", or "hidden".
 *
 * Loading
 * -------
 * This file is auto-discovered from ~/.pi/agent/extensions/. After adding it,
 * restart pi (or run /reload). In tmux, ctrl+q is a single byte (0x11) that
 * tmux passes through (only ctrl+b is tmux's prefix); if your terminal eats
 * ctrl+q, use the /tools command instead.
 *
 * Notes
 * ------
 * - `read` of an image file still renders the image in hidden mode (image
 *   rendering is owned by the tool component, not the result renderer); the
 *   text part is hidden as expected. Rare edge case.
 * - Rendering for read/grep/find/ls is delegated to the originals (they use a
 *   Text-based renderer, safe to delegate). bash/edit/write use a Container +
 *   shared-state renderer, so they are rendered here (Text-based) to avoid
 *   lastComponent type-mismatch crashes.
 */

import type {
	ExtensionAPI,
	ToolDefinition,
	ToolRenderContext,
	ToolRenderResultOptions,
	AgentToolResult,
} from "@earendil-works/pi-coding-agent";
import {
	createReadToolDefinition,
	createBashToolDefinition,
	createEditToolDefinition,
	createWriteToolDefinition,
	createGrepToolDefinition,
	createFindToolDefinition,
	createLsToolDefinition,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";

// ----------------------------------------------------------------------------
// View state
// ----------------------------------------------------------------------------
type Mode = "folded" | "expanded" | "hidden";
const MODES: Mode[] = ["folded", "expanded", "hidden"];

function readDefaultMode(): Mode {
	const v = (process.env.PI_TOOL_FOLD_MODE || "").trim().toLowerCase();
	return v === "expanded" || v === "hidden" ? v : "folded";
}

let mode: Mode = readDefaultMode();

// toolCallId -> invalidate() for every rendered tool block, so a mode toggle
// can force every block (past + present) to re-render with the new state.
const invalidators = new Map<string, () => void>();

function track(context: ToolRenderContext): void {
	if (context.toolCallId) invalidators.set(context.toolCallId, context.invalidate);
}

function rerenderAll(): void {
	for (const inv of invalidators.values()) {
		try {
			inv();
		} catch {
			// Component may be gone (compaction/session switch) — ignore.
		}
	}
}

// Tools whose built-in renderResult uses the Text pattern (safe to delegate to,
// preserving their rich rendering). bash/edit/write use Container + shared state
// and are rendered locally instead.
const DELEGATABLE = new Set(["read", "grep", "find", "ls"]);

// ----------------------------------------------------------------------------
// Small text helpers (all renderers return Text so lastComponent reuse is safe)
// ----------------------------------------------------------------------------
function asText(context: ToolRenderContext): Text {
	return (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
}

function setText(context: ToolRenderContext, content: string): Text {
	const t = asText(context);
	t.setText(content);
	return t;
}

function emptyText(context: ToolRenderContext): Text {
	return setText(context, "");
}

function firstLine(s: string): string {
	return s ? (s.split("\n")[0] ?? "") : "";
}

function lineCount(s: string): number {
	return s ? s.split("\n").length : 0;
}

function resultText(result: AgentToolResult): string {
	const content = (result?.content ?? []) as Array<{ type: string; text?: string }>;
	return content
		.filter((c) => c.type === "text")
		.map((c) => c.text ?? "")
		.join("\n");
}

function hasImage(result: AgentToolResult): boolean {
	const content = (result?.content ?? []) as Array<{ type: string }>;
	return content.some((c) => c.type === "image");
}

// ----------------------------------------------------------------------------
// Call header (renderCall) — one compact line per tool
// ----------------------------------------------------------------------------
function formatCall(name: string, args: any, theme: any): string {
	const title = (s: string) => theme.fg("toolTitle", theme.bold(s));
	const accent = (s: string) => theme.fg("accent", s);
	const dim = (s: string) => theme.fg("dim", s);
	switch (name) {
		case "read": {
			let s = title("read ") + accent(String(args?.path ?? ""));
			const parts: string[] = [];
			if (args?.offset) parts.push(`offset=${args.offset}`);
			if (args?.limit) parts.push(`limit=${args.limit}`);
			if (parts.length) s += dim(` (${parts.join(", ")})`);
			return s;
		}
		case "bash": {
			const cmd = String(args?.command ?? "");
			const trunc = cmd.length > 120 ? cmd.slice(0, 117) + "..." : cmd;
			let s = title("$ ") + accent(trunc);
			if (args?.timeout) s += dim(` (timeout ${args.timeout}s)`);
			return s;
		}
		case "edit":
			return title("edit ") + accent(String(args?.path ?? ""));
		case "write": {
			const n = typeof args?.content === "string" ? args.content.split("\n").length : 0;
			return title("write ") + accent(String(args?.path ?? "")) + dim(` (${n} lines)`);
		}
		case "grep":
			return title("grep ") + accent(String(args?.pattern ?? "")) + dim(` ${args?.path ?? ""}`);
		case "find":
			return (
				title("find ") +
				accent(String(args?.path ?? "")) +
				(args?.pattern ? dim(` ${args.pattern}`) : "")
			);
		case "ls":
			return title("ls ") + accent(String(args?.path ?? ""));
		default:
			return title(name);
	}
}

// ----------------------------------------------------------------------------
// Folded (compact) result summary — one line
// ----------------------------------------------------------------------------
function foldedResult(name: string, result: AgentToolResult, context: ToolRenderContext, theme: any): string {
	const isError = context.isError;
	const text = resultText(result);
	const lines = lineCount(text);
	const icon = isError ? theme.fg("error", "✗") : theme.fg("success", "✓");

	if (isError) {
		if (name === "bash") {
			const m = text.match(/Command exited with code (\d+)/);
			if (m) return `${icon} ${theme.fg("error", `exit ${m[1]}`)}${theme.fg("dim", ` · ${lines} lines`)}`;
		}
		const msg = firstLine(text) || "failed";
		const trimmed = msg.length > 80 ? msg.slice(0, 77) + "..." : msg;
		return `${icon} ${theme.fg("error", trimmed)}${theme.fg("dim", ` · ${lines} lines`)}`;
	}

	if (name === "edit") {
		const diff = (result as any)?.details?.diff;
		if (typeof diff === "string") {
			let add = 0;
			let rem = 0;
			for (const l of diff.split("\n")) {
				if (l.startsWith("+") && !l.startsWith("+++")) add++;
				else if (l.startsWith("-") && !l.startsWith("---")) rem++;
			}
			return `${icon} ${theme.fg("success", `+${add}`)} ${theme.fg("dim", "/")} ${theme.fg("error", `-${rem}`)}`;
		}
		return `${icon} ${theme.fg("success", "applied")}`;
	}

	if (hasImage(result) && !text) return `${icon} ${theme.fg("dim", "image")}`;
	return `${icon} ${theme.fg("dim", `${lines} lines`)}`;
}

// ----------------------------------------------------------------------------
// Expanded result for bash/edit/write (rendered locally; Container-based tools)
// ----------------------------------------------------------------------------
function expandedOwn(name: string, result: AgentToolResult, context: ToolRenderContext, theme: any): string {
	const text = resultText(result);
	const details = (result as any)?.details ?? {};

	if (name === "edit") {
		const diff: string | undefined = details.diff;
		if (typeof diff === "string" && diff) {
			return diff
				.split("\n")
				.map((l) => {
					if (l.startsWith("+") && !l.startsWith("+++")) return theme.fg("success", l);
					if (l.startsWith("-") && !l.startsWith("---")) return theme.fg("error", l);
					return theme.fg("dim", l);
				})
				.join("\n");
		}
		return context.isError ? theme.fg("error", firstLine(text) || "failed") : theme.fg("success", "Applied");
	}

	// bash, write
	const out = text
		? text
				.split("\n")
				.map((l) => theme.fg("muted", l))
				.join("\n")
		: theme.fg("success", name === "write" ? "Written" : "done");
	let s = out;
	const trunc = details.truncation;
	if (trunc?.truncated) {
		const by = trunc.truncatedBy === "lines"
			? `${trunc.outputLines} of ${trunc.totalLines} lines`
			: `${trunc.outputLines} lines`;
		s += `\n${theme.fg("warning", `[truncated: ${by}]`)}`;
	}
	if (details.fullOutputPath) {
		s += `\n${theme.fg("dim", `Full output: ${details.fullOutputPath}`)}`;
	}
	return s;
}

// ----------------------------------------------------------------------------
// Per-tool registration
// ----------------------------------------------------------------------------
function registerFoldable(pi: ExtensionAPI, cwd: string): void {
	const originals: Record<string, ToolDefinition> = {
		read: createReadToolDefinition(cwd),
		bash: createBashToolDefinition(cwd),
		edit: createEditToolDefinition(cwd),
		write: createWriteToolDefinition(cwd),
		grep: createGrepToolDefinition(cwd),
		find: createFindToolDefinition(cwd),
		ls: createLsToolDefinition(cwd),
	};

	for (const name of Object.keys(originals)) {
		const orig = originals[name];
		pi.registerTool({
			name: orig.name,
			label: orig.label,
			description: orig.description,
			parameters: orig.parameters,
			prepareArguments: orig.prepareArguments,
			executionMode: orig.executionMode,
			renderShell: "self",

			// Behavior is identical to the built-in tool — we only override rendering.
			execute: (toolCallId, params, signal, onUpdate, ctx) =>
				orig.execute(toolCallId, params, signal, onUpdate, ctx),

			renderCall: (args, theme, context) => {
				track(context);
				// Hide the whole block once completed in hidden mode; keep the
				// command visible while running so activity is visible.
				if (mode === "hidden" && !context.isPartial) return emptyText(context);
				return setText(context, formatCall(name, args, theme));
			},

			renderResult: (result, options, theme, context) => {
				track(context);

				// While running: show live status. Delegate to the original for the
				// Text-based tools (nice "Reading…"/streaming indicators); render a
				// simple live line for the Container-based tools.
				if (options.isPartial) {
					if (DELEGATABLE.has(name)) {
						return orig.renderResult!(result, options, theme, context);
					}
					if (name === "bash") {
						const out = resultText(result);
						const tail = out ? out.split("\n").slice(-10) : [];
						let s = theme.fg("warning", "Running…");
						if (tail.length) s += "\n" + tail.map((l: string) => theme.fg("muted", l)).join("\n");
						return setText(context, s);
					}
					return setText(context, theme.fg("warning", name === "edit" ? "Editing…" : "Writing…"));
				}

				// Completed
				if (mode === "hidden") return emptyText(context);

				if (mode === "expanded" || options.expanded) {
					if (DELEGATABLE.has(name)) {
						// Force full output regardless of the global ctrl+o state.
						return orig.renderResult!(result, { ...options, expanded: true }, theme, context);
					}
					return setText(context, expandedOwn(name, result, context, theme));
				}

				// folded
				return setText(context, foldedResult(name, result, context, theme));
			},
		});
	}
}

// ----------------------------------------------------------------------------
// Mode toggle + status
// ----------------------------------------------------------------------------
function applyMode(next: Mode, ctx: { ui: any }): void {
	mode = next;
	rerenderAll();
	const label = `tools: ${mode}`;
	try {
		ctx.ui.setStatus("tool-fold", label);
	} catch {
		// setStatus not available in this context (e.g. print mode) — ignore.
	}
	try {
		ctx.ui.notify(`Tool blocks: ${mode}`, "info");
	} catch {
		// ignore
	}
}

function cycleMode(ctx: { ui: any }): void {
	const idx = MODES.indexOf(mode);
	applyMode(MODES[(idx + 1) % MODES.length], ctx);
}

function setModeFromArg(arg: string, ctx: { ui: any }): void {
	const a = (arg || "").trim().toLowerCase();
	if (a === "folded" || a === "expanded" || a === "hidden") {
		applyMode(a as Mode, ctx);
		return;
	}
	cycleMode(ctx);
}

// ----------------------------------------------------------------------------
// Extension entry point
// ----------------------------------------------------------------------------
export default function foldableTools(pi: ExtensionAPI): void {
	// Register immediately (covers the common single-session case) and re-register
	// on every session_start so the tools use the live session cwd for path
	// resolution (the built-in execute closes over the construction cwd).
	registerFoldable(pi, process.cwd());

	pi.on("session_start", async (_event, ctx) => {
		invalidators.clear();
		registerFoldable(pi, ctx.cwd ?? process.cwd());
		try {
			ctx.ui.setStatus("tool-fold", `tools: ${mode}`);
		} catch {
			// ignore
		}
	});

	// Keyboard: ctrl+q cycles folded → expanded → hidden.
	pi.registerShortcut("ctrl+q", {
		description: "Cycle tool-block view: folded → expanded → hidden",
		handler: async (ctx) => {
			cycleMode(ctx);
		},
	});

	// Slash command: /tools [folded|expanded|hidden]  (cycles when no arg).
	pi.registerCommand("tools", {
		description: "Set or cycle tool-block view (folded | expanded | hidden)",
		handler: async (args, ctx) => {
			setModeFromArg(args, ctx);
		},
	});
}