/**
 * The `ipython` tool — one persistent Python REPL for the agent, with the
 * rlm() recursion builtin wired through the kernel.
 *
 * Output discipline follows pi's built-in tools: 50KB / 2000 lines for
 * stdout, smaller caps for stderr and trailing-expression reprs, with
 * explicit truncation notices. rlm() results never pass through here — they
 * stay inside the REPL as Python values; only their token usage is reported
 * back as nested-call accounting on the tool result.
 */

import {
	DEFAULT_MAX_BYTES,
	DEFAULT_MAX_LINES,
	type ExtensionAPI,
	type ExtensionContext,
	truncateHead,
	truncateTail,
} from "@earendil-works/pi-coding-agent";
import type { Usage } from "@earendil-works/pi-ai";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import type { KernelProcess } from "./kernel.ts";

export interface NestedUsage {
	usage: Usage;
	rlmCalls: number;
}

export interface IpythonToolDeps {
	getKernel: () => KernelProcess;
	getMaxDepth: () => number;
	/** reset the per-exec nested-usage accumulator */
	beginExec: () => void;
	/** read the per-exec nested-usage accumulator */
	endExec: () => NestedUsage;
	/** called on every execution with the live context (cwd/model capture) */
	onContext: (ctx: ExtensionContext) => void;
	/** best-effort flush of the pending user_prompt into a fresh kernel */
	flushPendingUserPrompt: (kernel: KernelProcess) => void;
}

const STDERR_MAX_BYTES = 16 * 1024;
const REPR_MAX_BYTES = 16 * 1024;
const REPR_MAX_LINES = 200;
const STREAM_UPDATE_MS = 200;

export function registerIpythonTool(pi: ExtensionAPI, deps: IpythonToolDeps): void {
	pi.registerTool({
		name: "ipython",
		label: "IPython",
		description: [
			"Execute Python 3 code in a persistent kernel. Variables, imports and definitions survive between calls — use the namespace as long-term memory.",
			"The kernel provides rlm(prompt, context=None) -> str, which spawns a recursive sub-agent with a fresh context window and returns its answer as a Python string. The sub-agent sees only what you pass it (f-strings or context=). Delegate chunk-level subtasks (summarize/extract/classify/transform) to rlm() and aggregate results in Python; rlm() results never enter this conversation unless you print them.",
			"Also available: user_prompt (the current user request as a string).",
			"Output: stdout, stderr, and the repr of a trailing expression. stdout is truncated (~50KB/2000 lines); store large data in variables and print small slices.",
		].join(" "),
		promptSnippet: "Execute Python in a persistent REPL with rlm() sub-agent recursion",
		promptGuidelines: [
			"Use ipython for computation, data processing, and inspecting large content programmatically instead of pasting it into the conversation.",
			"Call rlm() inside ipython to delegate chunk-level subtasks over large data to a sub-agent, then aggregate results in Python.",
			"Print only small, decisive values from ipython cells; keep large data in kernel variables.",
		],
		parameters: Type.Object({
			code: Type.String({
				description: "Python 3 code to execute in the persistent kernel namespace.",
			}),
		}),

		async execute(_toolCallId, params, signal, onUpdate, ctx) {
			deps.onContext(ctx);
			const kernel = deps.getKernel();
			await kernel.ensureStarted(); // boot first…
			deps.flushPendingUserPrompt(kernel); // …so user_prompt lands before the first cell
			deps.beginExec();

			let lastUpdate = 0;
			let streamWrites = 0;
			const result = await kernel.exec(params.code, {
				signal,
				onStream: () => {
					streamWrites++;
					const now = Date.now();
					if (onUpdate && now - lastUpdate > STREAM_UPDATE_MS) {
						lastUpdate = now;
						onUpdate({
							content: [{ type: "text", text: `cell running… (${streamWrites} output writes so far)` }],
							details: {},
						});
					}
				},
			});
			const nested = deps.endExec();

			const parts: string[] = [];
			if (result.restoredVars.length > 0) {
				const names = result.restoredVars.slice(0, 10).join(", ");
				parts.push(
					`[kernel restored ${result.restoredVars.length} variable(s) from the session snapshot: ${names}${result.restoredVars.length > 10 ? ", …" : ""}]`,
				);
			}
			if (result.restarted) {
				parts.push("[kernel restarted — all previous variables were lost]");
			}

			const stdout = truncateHead(result.stdout, {
				maxLines: DEFAULT_MAX_LINES,
				maxBytes: DEFAULT_MAX_BYTES,
			});
			if (stdout.content.trim()) {
				parts.push(
					stdout.content +
						(stdout.truncated
							? `\n[stdout truncated: ${stdout.outputLines} of ${stdout.totalLines} lines; the data is still in the kernel — print smaller slices]`
							: ""),
				);
			}

			if (result.result) {
				const repr = truncateHead(result.result, {
					maxLines: REPR_MAX_LINES,
					maxBytes: REPR_MAX_BYTES,
				});
				parts.push(`Out: ${repr.content}${repr.truncated ? "\n[result repr truncated]" : ""}`);
			}

			if (result.stderr.trim()) {
				const stderr = truncateTail(result.stderr, {
					maxLines: REPR_MAX_LINES,
					maxBytes: STDERR_MAX_BYTES,
				});
				parts.push(
					`${result.ok ? "stderr" : "Error"}:\n${stderr.content}${stderr.truncated ? "\n[stderr truncated]" : ""}`,
				);
			}

			if (nested.rlmCalls > 0) {
				parts.push(
					`[rlm: ${nested.rlmCalls} sub-call(s), nested tokens ↑${nested.usage.input} ↓${nested.usage.output}]`,
				);
			}

			return {
				content: [{ type: "text", text: parts.join("\n\n") || "(no output)" }],
				details: {
					ok: result.ok,
					cancelled: result.cancelled,
					restarted: result.restarted,
					rlmCalls: nested.rlmCalls,
				},
				...(nested.rlmCalls > 0 ? { usage: nested.usage } : {}),
			};
		},

		renderCall(args, theme) {
			const code = String(args.code ?? "");
			const lines = code.split("\n");
			const first = lines[0] ?? "";
			let text = theme.fg("toolTitle", theme.bold("ipython "));
			text +=
				lines.length === 1
					? theme.fg("dim", first.length > 80 ? `${first.slice(0, 80)}…` : first)
					: theme.fg("dim", `(${lines.length} lines) ${first.slice(0, 60)}${first.length > 60 ? "…" : ""}`);
			return new Text(text, 0, 0);
		},
	});
}