/**
 * pi-rlm — Recursive Language Models for pi.
 *
 * One extension, no fork:
 *   - `ipython` tool backed by a persistent Python kernel (unix socket +
 *     JSONL, stdlib-only shim — no ipykernel, no pip install)
 *   - `rlm(prompt, context=None)` inside the kernel → recursive sub-agents
 *     spawned as child `pi` processes that load this same extension
 *   - RLM mode (`pi --rlm` or `/rlm`): active tools collapse to ipython
 *     only and the system prompt is replaced with the RLM prompt
 *   - namespace snapshots: picklable kernel data is saved per session and
 *     restored when the session resumes
 *   - continual harness (`/refine`, `refine()` in the kernel): persistent
 *     memories/policies, refined by a tool-less child LLM analyzing the
 *     trajectory, injected into the system prompt on every turn
 *   - nested token/cost accounting rolls up onto the ipython tool result
 *
 * Depth: the root agent is depth 0. Each rlm() child runs with
 * PI_RLM_DEPTH = parent depth + 1; requests beyond the depth budget are
 * refused inside the kernel (rlm() raises, the model adapts).
 */

import { mkdirSync, readdirSync, statSync, unlinkSync } from "node:fs";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import { type ExtensionAPI, type ExtensionContext, getAgentDir } from "@earendil-works/pi-coding-agent";
import { addUsage, emptyUsage, runChildRlm } from "./child.ts";
import { registerIpythonTool, type NestedUsage } from "./ipython-tool.ts";
import { KernelProcess, type RlmReply, type RlmRequest } from "./kernel.ts";
import { buildRlmAddendum, buildRlmSystemPrompt } from "./prompt.ts";
import {
	applyEdits,
	buildHarnessSection,
	HARNESS_ENTRY_TYPE,
	type HarnessEntry,
	type HarnessScope,
	loadGlobalHarness,
	loadLocalHarness,
	REFINEMENT_ENTRY_TYPE,
	runRefinement,
	serializeConversation,
} from "./refine.ts";

const ENTRY_PATH = fileURLToPath(import.meta.url);
const KERNEL_PATH = join(dirname(ENTRY_PATH), "..", "kernel", "rlm_kernel.py");
const DEFAULT_MAX_DEPTH = 2;
const MODE_ENTRY_TYPE = "pi-rlm:mode";
const FALLBACK_DEFAULT_TOOLS = ["read", "bash", "edit", "write", "grep", "find", "ls", "ipython"];
const SNAPSHOT_MAX_AGE_MS = 30 * 24 * 60 * 60 * 1000;

function cleanupOldSnapshots(dir: string): void {
	try {
		const now = Date.now();
		for (const name of readdirSync(dir)) {
			if (!name.endsWith(".pickle")) continue;
			const full = join(dir, name);
			try {
				if (now - statSync(full).mtimeMs > SNAPSHOT_MAX_AGE_MS) unlinkSync(full);
			} catch {
				/* ignore */
			}
		}
	} catch {
		/* ignore */
	}
}

export default function (pi: ExtensionAPI) {
	// Our own depth in the rlm() recursion (0 = root interactive agent).
	const depth = Number.parseInt(process.env.PI_RLM_DEPTH ?? "", 10) || 0;
	const envParsed = Number.parseInt(process.env.PI_RLM_MAX_DEPTH ?? "", 10);
	const envMaxDepth = Number.isFinite(envParsed) && envParsed >= 0 ? envParsed : null;

	pi.registerFlag("rlm", {
		description: "Start in RLM mode: ipython-only toolset + RLM system prompt",
		type: "boolean",
		default: false,
	});
	pi.registerFlag("rlm-max-depth", {
		description: "Maximum rlm() recursion depth (default: 2, env: PI_RLM_MAX_DEPTH, 0 disables)",
		type: "string",
		default: "",
	});

	let maxDepthOverride: number | null = null;
	const getMaxDepth = (): number => {
		if (maxDepthOverride !== null) return maxDepthOverride;
		const flag = Number.parseInt(String(pi.getFlag("rlm-max-depth") ?? ""), 10);
		if (Number.isFinite(flag) && flag >= 0) return flag;
		return envMaxDepth ?? DEFAULT_MAX_DEPTH;
	};

	const globalHarnessPath = join(getAgentDir(), "pi-rlm", "harness.jsonl");

	let rlmMode = Boolean(pi.getFlag("rlm"));
	let kernel: KernelProcess | null = null;
	let defaultTools: string[] | null = null;
	let pendingUserPrompt: string | null = null;
	let snapshotPath: string | null = null;
	let lastCwd = process.cwd();
	let lastModel: string | undefined;
	let pendingRefinement: { instructions?: string; scope: HarnessScope } | null = null;
	let refineInFlight = false;

	// Per-exec nested-usage accumulator. Kernel executions are serialized and
	// rlm() requests only arrive while a cell runs, so this is race-free.
	let nestedUsage = emptyUsage();
	let nestedRlmCalls = 0;

	const captureContext = (ctx: ExtensionContext) => {
		lastCwd = ctx.cwd;
		if (ctx.model) lastModel = `${ctx.model.provider}/${ctx.model.id}`;
	};

	async function handleRlmRequest(req: RlmRequest, signal?: AbortSignal): Promise<RlmReply> {
		const maxDepth = getMaxDepth();
		if (req.depth > maxDepth) {
			return {
				ok: false,
				error: `rlm() max depth ${maxDepth} reached (this would be depth ${req.depth}). Continue without recursion.`,
			};
		}
		const result = await runChildRlm({
			prompt: req.prompt,
			depth: req.depth,
			maxDepth,
			extensionPath: ENTRY_PATH,
			model: lastModel,
			cwd: lastCwd,
			signal,
		});
		addUsage(nestedUsage, result.usage);
		nestedRlmCalls++;
		if (!result.ok) {
			return { ok: false, error: result.error ?? "rlm() child failed without an error message" };
		}
		return { ok: true, result: result.output };
	}

	const getKernel = (): KernelProcess => {
		if (!kernel) {
			kernel = new KernelProcess({
				pythonPath: process.env.PI_RLM_PYTHON || "python3",
				kernelPath: KERNEL_PATH,
				depth,
				snapshotPath: snapshotPath ?? undefined,
				onRlmRequest: handleRlmRequest,
				onRefineRequest: async (req) => {
					if (refineInFlight) return { scheduled: false, reason: "a refinement is already running" };
					pendingRefinement = {
						instructions: req.instructions ?? undefined,
						scope: req.global ? "global" : "local",
					};
					return { scheduled: true };
				},
			});
		}
		return kernel;
	};

	const flushPendingUserPrompt = (k: KernelProcess) => {
		if (pendingUserPrompt !== null && k.isReady) {
			void k.setVar("user_prompt", pendingUserPrompt).catch(() => {});
			pendingUserPrompt = null;
		}
	};

	registerIpythonTool(pi, {
		getKernel,
		getMaxDepth,
		beginExec: () => {
			nestedUsage = emptyUsage();
			nestedRlmCalls = 0;
		},
		endExec: (): NestedUsage => ({ usage: nestedUsage, rlmCalls: nestedRlmCalls }),
		onContext: captureContext,
		flushPendingUserPrompt,
	});

	const applyRlmTools = () => {
		const active = pi.getActiveTools();
		if (!defaultTools) {
			defaultTools = active.includes("ipython") ? active : [...active, "ipython"];
		}
		pi.setActiveTools(["ipython"]);
	};

	const restoreTools = () => {
		pi.setActiveTools(defaultTools ?? FALLBACK_DEFAULT_TOOLS);
	};

	const ensureIpythonActive = () => {
		const active = pi.getActiveTools();
		if (!active.includes("ipython")) pi.setActiveTools([...active, "ipython"]);
	};

	// ------------------------------------------------------------- refinement

	async function runRefinementJob(ctx: ExtensionContext, job: { instructions?: string; scope: HarnessScope }) {
		// notify is a no-op without a UI; keep print/json mode informed via stderr
		const say = (msg: string, level: "info" | "warning" | "error") => {
			if (ctx.hasUI) ctx.ui.notify(msg, level);
			else console.error(`pi-rlm: ${msg}`);
		};
		if (refineInFlight) {
			say("a refinement is already running", "warning");
			return;
		}
		refineInFlight = true;
		try {
			if (ctx.hasUI) ctx.ui.setStatus("pi-rlm", "refining…");
			const globalEntries = loadGlobalHarness(globalHarnessPath);
			const localEntries = loadLocalHarness(ctx);
			const result = await runRefinement({
				instructions: job.instructions,
				scope: job.scope,
				entries: [...globalEntries, ...localEntries],
				conversation: serializeConversation(ctx),
				model: lastModel,
				cwd: lastCwd,
			});
			if (!result.ok || !result.proposal) {
				say(`refinement failed: ${result.error ?? "no proposal"}`, "error");
				return;
			}
			const { applied, errors } = applyEdits({
				edits: result.proposal.edits,
				defaultScope: job.scope,
				globalPath: globalHarnessPath,
				globalEntries,
				localEntries,
				appendLocal: (entry) => pi.appendEntry(HARNESS_ENTRY_TYPE, entry),
			});
			pi.appendEntry(REFINEMENT_ENTRY_TYPE, {
				summary: result.proposal.summary,
				instructions: job.instructions,
				scope: job.scope,
				applied: applied.length,
				errors,
			});
			const errNote = errors.length > 0 ? ` (${errors.length} failed: ${errors[0]})` : "";
			say(
				`refinement: ${result.proposal.summary} — ${applied.length} edit(s) applied${errNote}`,
				applied.length > 0 ? "info" : "warning",
			);
		} finally {
			refineInFlight = false;
			if (ctx.hasUI && rlmMode) ctx.ui.setStatus("pi-rlm", depth > 0 ? `rlm d${depth}` : "rlm");
		}
	}

	// ---------------------------------------------------------------- events

	pi.on("session_start", async (_event, ctx) => {
		// Namespace snapshot follows the session: /resume restores it,
		// /new starts fresh, ephemeral sessions skip it.
		const sessionFile = ctx.sessionManager.getSessionFile();
		if (sessionFile) {
			try {
				const dir = join(getAgentDir(), "pi-rlm", "snapshots");
				mkdirSync(dir, { recursive: true });
				snapshotPath = join(dir, `${ctx.sessionManager.getSessionId()}.pickle`);
				cleanupOldSnapshots(dir);
			} catch {
				snapshotPath = null;
			}
		} else {
			snapshotPath = null;
		}

		if (!rlmMode) {
			// restore a mode persisted by /rlm in a previous run of this session
			for (const entry of ctx.sessionManager.getEntries()) {
				if (entry.type === "custom" && entry.customType === MODE_ENTRY_TYPE) {
					const data = entry.data as { enabled?: boolean } | undefined;
					rlmMode = Boolean(data?.enabled);
				}
			}
		}
		if (rlmMode) {
			applyRlmTools();
			if (ctx.hasUI) ctx.ui.setStatus("pi-rlm", depth > 0 ? `rlm d${depth}` : "rlm");
		} else {
			ensureIpythonActive();
		}
	});

	pi.on("before_agent_start", async (event, ctx) => {
		captureContext(ctx);
		const k = kernel;
		if (k?.isReady) {
			void k.setVar("user_prompt", event.prompt).catch(() => {});
		} else {
			pendingUserPrompt = event.prompt;
		}
		const info = { cwd: ctx.cwd, depth, maxDepth: getMaxDepth() };
		const harness = buildHarnessSection(loadGlobalHarness(globalHarnessPath), loadLocalHarness(ctx));
		const base = rlmMode
			? buildRlmSystemPrompt(info)
			: `${event.systemPrompt}\n\n${buildRlmAddendum(info)}`;
		return { systemPrompt: harness ? `${base}\n\n${harness}` : base };
	});

	// A refinement scheduled via refine() in the kernel runs once the turn
	// has fully settled — never mid-cell.
	pi.on("agent_settled", async (_event, ctx) => {
		if (!pendingRefinement || refineInFlight) return;
		const job = pendingRefinement;
		pendingRefinement = null;
		await runRefinementJob(ctx, job);
	});

	// The kernel process survives compaction, but the model forgets which
	// variables exist — re-ground it with a namespace snapshot.
	pi.on("session_compact", async () => {
		const k = kernel;
		if (!k?.isReady) return;
		try {
			const vars = await k.listVars();
			if (vars.length === 0) return;
			const lines = vars
				.slice(0, 50)
				.map((v) => `- ${v.name}: ${v.type} = ${v.repr}`)
				.join("\n");
			pi.sendMessage({
				customType: "pi-rlm:kernel-state",
				content: `[pi-rlm] The Python kernel survived compaction. Variables currently defined:\n${lines}`,
				display: false,
			});
		} catch {
			/* best effort */
		}
	});

	pi.on("session_shutdown", async () => {
		const k = kernel;
		kernel = null;
		if (k) {
			if (snapshotPath && k.isReady) await k.snapshot().catch(() => {});
			await k.shutdown();
		}
	});

	// --------------------------------------------------------------- commands

	pi.registerCommand("rlm", {
		description: "Toggle RLM mode. Subcommands: status | restart | depth <n>",
		handler: async (args, ctx) => {
			const sub = args.trim();
			if (sub === "status") {
				const s = kernel?.stats;
				const lines = [
					`RLM mode: ${rlmMode ? "ON" : "OFF"}`,
					`Depth: ${depth} / max ${getMaxDepth()}`,
					`Model for rlm() children: ${lastModel ?? "(session default)"}`,
					s?.ready
						? `Kernel: running (python ${s.python}, pid ${s.pid}, ${s.execs} cells, ${s.rlmCalls} rlm calls, ${s.restarts} restarts)`
						: "Kernel: not started (boots on first ipython call)",
					snapshotPath
						? `Snapshot: ${snapshotPath}${s?.restored?.length ? ` (restored: ${s.restored.join(", ")})` : ""}`
						: "Snapshot: disabled (ephemeral session)",
					`Harness: ${loadGlobalHarness(globalHarnessPath).length} global, ${loadLocalHarness(ctx).length} local entries`,
				];
				ctx.ui.notify(lines.join("\n"), "info");
				return;
			}
			if (sub === "restart") {
				const k = kernel;
				kernel = null;
				if (k) await k.shutdown();
				if (snapshotPath) {
					try {
						unlinkSync(snapshotPath); // intentional wipe — don't restore on reboot
					} catch {
						/* ignore */
					}
				}
				ctx.ui.notify(
					"pi-rlm: kernel stopped and snapshot deleted. It will reboot (empty namespace) on the next ipython call.",
					"info",
				);
				return;
			}
			if (sub.startsWith("depth")) {
				const n = Number.parseInt(sub.slice(5).trim(), 10);
				if (Number.isFinite(n) && n >= 0) {
					maxDepthOverride = n;
					ctx.ui.notify(`pi-rlm: max rlm() depth set to ${n} (this session)`, "info");
				} else {
					ctx.ui.notify(`pi-rlm: current max depth is ${getMaxDepth()}`, "info");
				}
				return;
			}
			rlmMode = !rlmMode;
			pi.appendEntry(MODE_ENTRY_TYPE, { enabled: rlmMode });
			if (rlmMode) {
				applyRlmTools();
				if (ctx.hasUI) ctx.ui.setStatus("pi-rlm", depth > 0 ? `rlm d${depth}` : "rlm");
			} else {
				restoreTools();
				if (ctx.hasUI) ctx.ui.setStatus("pi-rlm", "");
			}
			ctx.ui.notify(
				rlmMode
					? "RLM mode ON — only the ipython tool is active, RLM system prompt engaged"
					: "RLM mode OFF — full toolset restored (ipython stays available)",
				"info",
			);
		},
	});

	pi.registerCommand("refine", {
		description: "Refine the persistent harness: /refine [global] [instructions] | /refine list",
		handler: async (args, ctx) => {
			const sub = args.trim();
			if (sub === "list") {
				const entries = [...loadGlobalHarness(globalHarnessPath), ...loadLocalHarness(ctx)];
				if (entries.length === 0) {
					if (ctx.hasUI) ctx.ui.notify("pi-rlm harness is empty", "info");
					else console.error("pi-rlm: harness is empty");
					return;
				}
				const fmt = (e: HarnessEntry) => `[${e.id}] (${e.kind}/${e.scope}) ${e.title} — ${e.content}`;
				const text = entries.map(fmt).join("\n");
				if (ctx.hasUI) ctx.ui.notify(text, "info");
				else console.error(`pi-rlm harness:\n${text}`);
				return;
			}
			let scope: HarnessScope = "local";
			let instructions = sub;
			if (sub.startsWith("global")) {
				scope = "global";
				instructions = sub.slice(6).trim();
			}
			await ctx.waitForIdle();
			await runRefinementJob(ctx, { instructions: instructions || undefined, scope });
		},
	});
}